From c99daef6899aec2ad818f66e3b007bf43ccad820 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 5 Jun 2025 21:18:19 -0700 Subject: [PATCH] [Fix]: /v1/messages - return streaming usage statistics when using litellm with bedrock models (#11469) * fix: using litellm with claude code bedrock * fix: usage for bedrock with /messages * fix: bedrock_sse_wrapper * tests: test for test_chunk_parser_usage_transformation * test fix --- .../anthropic_claude3_transformation.py | 49 ++++++++++++-- .../test_anthropic_claude3_transformation.py | 67 +++++++++++++++++++ 2 files changed, 109 insertions(+), 7 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py 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 03623bf86d..52e751d24a 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -1,3 +1,4 @@ +import json from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Tuple, Union import httpx @@ -13,6 +14,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import GenericStreamingChunk from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.types.utils import ModelResponseStream @@ -113,9 +115,9 @@ class AmazonAnthropicClaude3MessagesConfig( # 1. anthropic_version is required for all claude models if "anthropic_version" not in anthropic_messages_request: - anthropic_messages_request[ - "anthropic_version" - ] = self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION + anthropic_messages_request["anthropic_version"] = ( + self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION + ) # 2. `stream` is not allowed in request body for bedrock invoke if "stream" in anthropic_messages_request: @@ -139,7 +141,26 @@ class AmazonAnthropicClaude3MessagesConfig( completion_stream = aws_decoder.aiter_bytes( httpx_response.aiter_bytes(chunk_size=aws_decoder.DEFAULT_CHUNK_SIZE) ) - return completion_stream + # Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients. + return self.bedrock_sse_wrapper(completion_stream) + + async def bedrock_sse_wrapper( + self, + completion_stream: AsyncIterator[ + Union[bytes, GenericStreamingChunk, ModelResponseStream, dict] + ], + ): + """ + Bedrock invoke does not return SSE formatted data. This function is a wrapper to ensure litellm chunks are SSE formatted. + """ + async for chunk in completion_stream: + if isinstance(chunk, dict): + event_type: str = str(chunk.get("type", "message")) + payload = f"event: {event_type}\n" f"data: {json.dumps(chunk)}\n\n" + yield payload.encode() + else: + # For non-dict chunks, forward the original value unchanged so callers can leverage the richer Python objects if they wish. + yield chunk class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): @@ -159,8 +180,22 @@ class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): """ Parse the chunk data into anthropic /messages format - No transformation is needed for anthropic /messages format - - since bedrock invoke returns the response in the correct format + Bedrock returns usage metrics using camelCase keys. Convert these to + the Anthropic `/v1/messages` specification so callers receive a + consistent response shape when streaming. """ + 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 return chunk_data 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 new file mode 100644 index 0000000000..c972c07ce1 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -0,0 +1,67 @@ +import asyncio +import json +import os +import sys + +import pytest + +# Ensure the project root is on the import path so `litellm` can be imported when +# tests are executed from any working directory. +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaude3MessagesConfig, + AmazonAnthropicClaudeMessagesStreamDecoder, +) + + +@pytest.mark.asyncio +async def test_bedrock_sse_wrapper_encodes_dict_chunks(): + """Verify that `bedrock_sse_wrapper` converts dictionary chunks to properly formatted Server-Sent Events and forwards non-dict chunks unchanged.""" + + cfg = AmazonAnthropicClaude3MessagesConfig() + + async def _dummy_stream(): # type: ignore[return-type] + yield {"type": "message_delta", "text": "hello"} + yield b"raw-bytes" + + # Collect all chunks returned by the wrapper + collected: list[bytes] = [] + async for chunk in cfg.bedrock_sse_wrapper(_dummy_stream()): + collected.append(chunk) + + assert collected, "No chunks returned from wrapper" + + # First chunk should be SSE encoded + first_chunk = collected[0] + assert first_chunk.startswith(b"event: message_delta\n"), first_chunk + assert first_chunk.endswith(b"\n\n"), first_chunk + # Ensure the JSON payload is present in the SSE data line + assert b'"hello"' in first_chunk # payload contains the text + + # Second chunk should be forwarded unchanged + assert collected[1] == b"raw-bytes" + + +def test_chunk_parser_usage_transformation(): + """Ensure Bedrock invocation metrics are transformed to Anthropic usage keys.""" + + decoder = AmazonAnthropicClaudeMessagesStreamDecoder( + model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0" + ) + + chunk = { + "type": "message_delta", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 10, + "outputTokenCount": 5, + }, + } + + parsed = decoder._chunk_parser(chunk.copy()) # use copy to avoid side-effects + + # The invocation metrics key should be removed and replaced by `usage` + assert "amazon-bedrock-invocationMetrics" not in parsed + assert "usage" in parsed + assert parsed["usage"]["input_tokens"] == 10 + assert parsed["usage"]["output_tokens"] == 5