fix(bedrock): include cacheWriteInputTokens in prompt_tokens calculation

Fixes #15263

This PR fixes the cost calculation for Bedrock Anthropic models with prompt caching.

**Root Cause:**
PR #9838 incorrectly removed adding `cacheWriteInputTokens` to `prompt_tokens`
for Bedrock, based on the assumption that it would cause double counting (similar
to an Anthropic API issue). However, Bedrock's token structure is different:

- **Bedrock API**: `inputTokens`, `cacheReadInputTokens`, and `cacheWriteInputTokens`
  are ALL separate values that should be summed for total input tokens
- **Anthropic API**: Same structure - all three token types are separate

The fix in #9838 was later reverted for Anthropic (correctly re-adding
`cache_creation_input_tokens` to `prompt_tokens`), but Bedrock was never fixed.

**Changes:**
1. Re-add `cacheWriteInputTokens` to `input_tokens` in Bedrock transformation
2. Update test assertions to reflect correct behavior
3. Add regression test for prompt caching cost calculation
4. Fix typo in Anthropic transformation where `cache_creation_tokens` was
   incorrectly set to `cache_read_input_tokens`

**Testing:**
- All existing Bedrock transformation tests pass
- New test validates correct cost calculation with prompt caching
- Verified costs are non-negative and accurate
This commit is contained in:
Tim Elfrink 2025-10-07 20:28:46 +02:00
parent 421d38c94a
commit c5eb22381d
4 changed files with 40 additions and 9 deletions

View File

@ -878,7 +878,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cache_read_input_tokens,
cache_creation_tokens=cache_read_input_tokens,
cache_creation_tokens=cache_creation_input_tokens,
cache_creation_token_details=cache_creation_token_details,
)
completion_token_details = (

View File

@ -1092,10 +1092,8 @@ class AmazonConverseConfig(BaseConfig):
cache_read_input_tokens = usage["cacheReadInputTokens"]
input_tokens += cache_read_input_tokens
if "cacheWriteInputTokens" in usage:
"""
Do not increment prompt_tokens with cacheWriteInputTokens
"""
cache_creation_input_tokens = usage["cacheWriteInputTokens"]
input_tokens += cache_creation_input_tokens
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cache_read_input_tokens

View File

@ -626,3 +626,28 @@ def test_service_tier_fallback_pricing():
assert abs(std_cost[0] - expected_standard_prompt) < 1e-10, f"Standard prompt cost mismatch: {std_cost[0]} vs {expected_standard_prompt}"
assert abs(std_cost[1] - expected_standard_completion) < 1e-10, f"Standard completion cost mismatch: {std_cost[1]} vs {expected_standard_completion}"
def test_bedrock_anthropic_prompt_caching():
"""Test Bedrock Anthropic models with prompt caching return correct costs."""
model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
usage = Usage(
prompt_tokens=52123,
completion_tokens=497,
total_tokens=52620,
cache_creation_input_tokens=7183,
cache_read_input_tokens=22465,
)
custom_llm_provider = "bedrock"
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider,
)
assert prompt_cost >= 0
assert completion_cost >= 0
assert round(prompt_cost, 3) == 0.845
assert round(completion_cost, 5) == 0.00820

View File

@ -33,7 +33,7 @@ def test_transform_usage():
openai_usage = config._transform_usage(usage)
assert (
openai_usage.prompt_tokens
== usage["inputTokens"] + usage["cacheReadInputTokens"]
== usage["inputTokens"] + usage["cacheReadInputTokens"] + usage["cacheWriteInputTokens"]
)
assert openai_usage.completion_tokens == usage["outputTokens"]
assert openai_usage.total_tokens == usage["totalTokens"]
@ -1620,7 +1620,9 @@ async def test_no_cache_control_no_cache_point():
def test_guarded_text_wraps_in_guardrail_converse_content():
"""Test that guarded_text content type gets wrapped in guardContent blocks."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import (
_bedrock_converse_messages_pt,
)
messages = [
{
@ -1711,7 +1713,9 @@ def test_guarded_text_with_system_messages():
def test_guarded_text_with_mixed_content_types():
"""Test guarded_text with mixed content types including images."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import (
_bedrock_converse_messages_pt,
)
messages = [
{
@ -1752,7 +1756,9 @@ def test_guarded_text_with_mixed_content_types():
@pytest.mark.asyncio
async def test_async_guarded_text():
"""Test async version of guarded_text processing."""
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
from litellm.litellm_core_utils.prompt_templates.factory import (
BedrockConverseMessagesProcessor,
)
messages = [
{
@ -1789,7 +1795,9 @@ async def test_async_guarded_text():
def test_guarded_text_with_tool_calls():
"""Test guarded_text with tool calls in the conversation."""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import (
_bedrock_converse_messages_pt,
)
messages = [
{