mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(bedrock): map invoke streaming cacheRead/cacheWrite token counts into usage
This commit is contained in:
parent
35dc982692
commit
e24523fec0
4 changed files with 227 additions and 7 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue