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 98085c08a5e..14c476ece95 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -622,9 +622,9 @@ class AmazonAnthropicClaudeMessagesConfig( merges usage from message_start and message_delta but ignores message_stop. This method buffers message_delta and, when message_stop arrives with cache usage, merges those fields into the - message_delta usage and also updates the input_tokens on - message_delta to include the full count (uncached + cache_creation + - cache_read). + message_delta usage. input_tokens is kept as the uncached-only + count; downstream calculate_usage adds cache tokens to + prompt_tokens. """ _CACHE_FIELDS = ("cache_creation_input_tokens", "cache_read_input_tokens") pending_delta = None @@ -653,12 +653,7 @@ class AmazonAnthropicClaudeMessagesConfig( raw_input = stop_usage.get("input_tokens") if raw_input is not None: - uncached = raw_input if isinstance(raw_input, int) else 0 - raw_cc = delta_usage.get("cache_creation_input_tokens", 0) - cache_creation = raw_cc if isinstance(raw_cc, int) else 0 - raw_cr = delta_usage.get("cache_read_input_tokens", 0) - cache_read = raw_cr if isinstance(raw_cr, int) else 0 - delta_usage["input_tokens"] = uncached + cache_creation + cache_read + delta_usage["input_tokens"] = raw_input if isinstance(raw_input, int) else 0 if delta_usage: pending_delta["usage"] = delta_usage # type: ignore[arg-type] 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 fb5c2c37bfb..0cd09c1ec48 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 @@ -3,6 +3,7 @@ import json import os import sys from datetime import datetime +from unittest.mock import Mock import pytest @@ -130,7 +131,8 @@ async def test_bedrock_sse_wrapper_keeps_usage_in_message_start_and_message_delt assert "usage" in delta_json assert delta_json["usage"]["cache_creation_input_tokens"] == 1562 assert delta_json["usage"]["cache_read_input_tokens"] == 32392 - assert delta_json["usage"]["input_tokens"] == 3 + 1562 + 32392 + assert delta_json["usage"]["input_tokens"] == 3 + assert delta_json["usage"]["output_tokens"] == 8 def test_chunk_parser_usage_transformation(): @@ -599,3 +601,111 @@ def test_bedrock_messages_strips_output_config_with_output_format(): assert "output_config" not in result assert "output_format" not in result + + +@pytest.mark.asyncio +async def test_promote_message_stop_usage_preserves_message_delta_output_tokens(): + """ + Bedrock unified /messages streaming can send full usage on message_delta and a + conflicting smaller usage on message_stop (e.g. output_tokens 9 vs 12). + _promote_message_stop_usage must not replace message_delta output_tokens. + """ + cfg = AmazonAnthropicClaudeMessagesConfig() + + async def _stream(): # type: ignore[return-type] + yield { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": { + "input_tokens": 3, + "cache_creation_input_tokens": 10553, + "cache_read_input_tokens": 25490, + "output_tokens": 12, + }, + } + yield { + "type": "message_stop", + "usage": {"input_tokens": 3, "output_tokens": 9}, + } + + merged: list[dict] = [] + async for chunk in cfg._promote_message_stop_usage(_stream()): + if isinstance(chunk, dict): + merged.append(chunk) + + assert len(merged) >= 1 + delta_out = merged[0] + assert delta_out["type"] == "message_delta" + assert delta_out["usage"]["output_tokens"] == 12 + assert delta_out["usage"]["cache_creation_input_tokens"] == 10553 + assert delta_out["usage"]["cache_read_input_tokens"] == 25490 + assert delta_out["usage"]["input_tokens"] == 3 + + +@pytest.mark.asyncio +async def test_unified_bedrock_messages_sse_usage_and_cost_claude_sonnet_46(): + """ + End-to-end for Bedrock Invoke Anthropic Messages (unified) streaming path: + dict chunks -> _promote_message_stop_usage -> bedrock_sse_wrapper SSE bytes -> + same logging reconstruction as Anthropic /messages. Ensures token counts and + completion_cost match model_prices for us.anthropic.claude-sonnet-4-6. + """ + from litellm import completion_cost + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( + AnthropicPassthroughLoggingHandler, + ) + + cfg = AmazonAnthropicClaudeMessagesConfig() + + async def _stream(): # type: ignore[return-type] + yield { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": { + "input_tokens": 3, + "cache_creation_input_tokens": 10553, + "cache_read_input_tokens": 25490, + "output_tokens": 12, + }, + } + yield { + "type": "message_stop", + "usage": {"input_tokens": 3, "output_tokens": 9}, + } + + 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_unified_bedrock_messages_sse_cost", + function_id="test_unified_bedrock_messages_sse_cost", + ) + + 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.completion_tokens == 12 + assert built.usage.prompt_tokens == 36046 + assert built.usage.total_tokens == 36058 + assert built.usage.cache_creation_input_tokens == 10553 + assert built.usage.cache_read_input_tokens == 25490 + + cost = completion_cost( + completion_response=built, + model="bedrock/us.anthropic.claude-sonnet-4-6", + custom_llm_provider="bedrock", + ) + assert cost == pytest.approx(0.052150725, rel=0, abs=1e-9)