From e24523fec045c484c79206aa5ecf539cc6b5057f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 20:15:09 +0000 Subject: [PATCH] fix(bedrock): map invoke streaming cacheRead/cacheWrite token counts into usage --- .../anthropic_claude3_transformation.py | 26 ++- .../anthropic_passthrough_logging_handler.py | 6 + .../test_anthropic_claude3_transformation.py | 163 ++++++++++++++++++ ...t_anthropic_passthrough_logging_handler.py | 39 +++++ 4 files changed, 227 insertions(+), 7 deletions(-) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 08c13448d8c..13f37811c3d 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -898,6 +898,14 @@ class AmazonAnthropicClaudeMessagesConfig( yield pending_delta +_INVOCATION_METRICS_TO_ANTHROPIC_USAGE: Dict[str, str] = { + "inputTokenCount": "input_tokens", + "outputTokenCount": "output_tokens", + "cacheReadInputTokenCount": "cache_read_input_tokens", + "cacheWriteInputTokenCount": "cache_creation_input_tokens", +} + + class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): def __init__( self, @@ -915,14 +923,18 @@ class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): Bedrock returns usage metrics using camelCase keys. Convert these to the Anthropic `/v1/messages` specification so callers receive a - consistent response shape when streaming. + consistent response shape when streaming. Cache counts only ever appear + in this camelCase block on the Invoke path, so dropping them would bill + cached tokens as fresh input. """ amazon_bedrock_invocation_metrics = chunk_data.pop("amazon-bedrock-invocationMetrics", {}) if amazon_bedrock_invocation_metrics: - anthropic_usage = {} - if "inputTokenCount" in amazon_bedrock_invocation_metrics: - anthropic_usage["input_tokens"] = amazon_bedrock_invocation_metrics["inputTokenCount"] - if "outputTokenCount" in amazon_bedrock_invocation_metrics: - anthropic_usage["output_tokens"] = amazon_bedrock_invocation_metrics["outputTokenCount"] - chunk_data["usage"] = anthropic_usage + chunk_data["usage"] = { + **(chunk_data.get("usage") or {}), + **{ + anthropic_key: amazon_bedrock_invocation_metrics[bedrock_key] + for bedrock_key, anthropic_key in _INVOCATION_METRICS_TO_ANTHROPIC_USAGE.items() + if bedrock_key in amazon_bedrock_invocation_metrics + }, + } return chunk_data diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 50e90699194..0de75b66895 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -709,6 +709,12 @@ class AnthropicPassthroughLoggingHandler: tool_search_requests = _stu.get("tool_search_requests") if usage.get("cache_read_input_tokens") is not None: cache_read = usage.get("cache_read_input_tokens") + if usage.get("cache_creation_input_tokens") is not None: + cache_creation = usage.get("cache_creation_input_tokens") + _delta_cc = usage.get("cache_creation") + if isinstance(_delta_cc, dict): + cache_creation_5m = _delta_cc.get("ephemeral_5m_input_tokens") + cache_creation_1h = _delta_cc.get("ephemeral_1h_input_tokens") if usage.get("inference_geo") is not None: inference_geo = usage.get("inference_geo") found_usage = True diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 3b8b4af78d9..9205404343a 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -264,6 +264,169 @@ def test_chunk_parser_usage_transformation(): assert parsed["usage"]["output_tokens"] == 5 +def test_chunk_parser_maps_cache_token_counts(): + """Cache counts only exist in the camelCase invocationMetrics block on the Invoke + path; dropping them bills cached traffic as fresh input (issue #34497).""" + + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/anthropic.claude-3-5-sonnet-20241022-v2:0" + ) + + parsed = decoder._chunk_parser( + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 1, + "outputTokenCount": 162, + "cacheReadInputTokenCount": 421714, + "cacheWriteInputTokenCount": 1139, + }, + } + ) + + assert parsed["usage"] == { + "input_tokens": 1, + "output_tokens": 162, + "cache_read_input_tokens": 421714, + "cache_creation_input_tokens": 1139, + } + + +def test_chunk_parser_preserves_existing_usage_fields(): + """Usage already present on the chunk (e.g. the 5m/1h cache split) must survive the + invocationMetrics merge.""" + + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/anthropic.claude-3-5-sonnet-20241022-v2:0" + ) + + parsed = decoder._chunk_parser( + { + "type": "message_stop", + "usage": { + "cache_creation": { + "ephemeral_5m_input_tokens": 1139, + "ephemeral_1h_input_tokens": 0, + } + }, + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 1, + "outputTokenCount": 162, + }, + } + ) + + assert parsed["usage"]["cache_creation"] == { + "ephemeral_5m_input_tokens": 1139, + "ephemeral_1h_input_tokens": 0, + } + assert parsed["usage"]["input_tokens"] == 1 + assert parsed["usage"]["output_tokens"] == 162 + + +@pytest.mark.asyncio +async def test_invoke_stream_cache_metrics_reach_usage_and_cost(): + """End-to-end Bedrock Invoke streaming: cache counts arrive only in + amazon-bedrock-invocationMetrics on the last chunk and must reach the + reconstructed usage so cached tokens are not priced as fresh input.""" + from litellm import completion_cost + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( + AnthropicPassthroughLoggingHandler, + ) + + cfg = AmazonAnthropicClaudeMessagesConfig() + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/us.anthropic.claude-sonnet-4-6" + ) + + raw_chunks = [ + { + "type": "message_start", + "message": { + "id": "msg_bdrk_1", + "type": "message", + "role": "assistant", + "content": [], + "model": "claude-sonnet-4-6", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "hi"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 162}, + }, + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 1, + "outputTokenCount": 162, + "cacheReadInputTokenCount": 421714, + "cacheWriteInputTokenCount": 1139, + }, + }, + ] + + async def _stream(): # type: ignore[return-type] + for raw_chunk in raw_chunks: + yield decoder._chunk_parser(raw_chunk) + + logging_obj = LiteLLMLoggingObj( + model="bedrock/us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + call_type="chat", + start_time=datetime.now(), + litellm_call_id="test_invoke_stream_cache_metrics", + function_id="test_invoke_stream_cache_metrics", + ) + + collected: list[bytes] = [] + async for sse in cfg.bedrock_sse_wrapper( + completion_stream=_stream(), + litellm_logging_obj=logging_obj, + request_body={"model": "us.anthropic.claude-sonnet-4-6"}, + ): + collected.append(sse) + + built = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=collected, + model="us.anthropic.claude-sonnet-4-6", + litellm_logging_obj=Mock(), + ) + assert built.usage is not None + assert built.usage.cache_read_input_tokens == 421714 + assert built.usage.cache_creation_input_tokens == 1139 + assert built.usage.prompt_tokens == 422854 + assert built.usage.completion_tokens == 162 + + cached_cost = completion_cost( + completion_response=built, + model="bedrock/us.anthropic.claude-sonnet-4-6", + custom_llm_provider="bedrock", + ) + built.usage.cache_read_input_tokens = 0 + built.usage.cache_creation_input_tokens = 0 + built.usage.prompt_tokens_details = None + fresh_input_cost = completion_cost( + completion_response=built, + model="bedrock/us.anthropic.claude-sonnet-4-6", + custom_llm_provider="bedrock", + ) + assert cached_cost < fresh_input_cost + + def test_remove_ttl_from_cache_control(): """Ensure ttl field is removed from cache_control in messages.""" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 947a7a64beb..e836e7ca756 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2005,6 +2005,45 @@ class TestAnthropicUsageOnlyFallback: assert usage.server_tool_use.web_search_requests == 1 assert usage.server_tool_use.tool_search_requests == 3 + def test_build_usage_only_recovers_cache_creation_from_message_delta(self): + """Bedrock Invoke lands the whole cache breakdown on the final message_delta + (promoted from amazon-bedrock-invocationMetrics) with nothing on message_start; + cache-write tokens must not be billed at zero (issue #34497)""" + chunks = [ + _sse_bytes( + { + "type": "message_start", + "message": { + "model": "claude-3-5-haiku-20241022", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + } + ), + _sse_bytes( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 1, + "output_tokens": 162, + "cache_read_input_tokens": 421714, + "cache_creation_input_tokens": 1139, + }, + } + ), + ] + response = ( + AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=chunks, model="claude-3-5-haiku-20241022" + ) + ) + assert response is not None + usage = response.usage + assert usage.prompt_tokens == 422854 + assert usage.completion_tokens == 162 + assert usage._cache_read_input_tokens == 421714 + assert usage._cache_creation_input_tokens == 1139 + @pytest.mark.parametrize( "event_str,expected", [