fix(bedrock): map invoke streaming cacheRead/cacheWrite token counts into usage

This commit is contained in:
Devin AI 2026-07-24 20:15:09 +00:00
parent 35dc982692
commit e24523fec0
4 changed files with 227 additions and 7 deletions

View file

@ -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

View file

@ -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

View file

@ -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."""

View file

@ -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",
[