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:
parent
421d38c94a
commit
c5eb22381d
@ -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 = (
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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 = [
|
||||
{
|
||||
|
||||
Loading…
Reference in New Issue
Block a user