[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

Author: ishaan-jaff

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 @@
     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 @@ def transform_anthropic_messages_request(
 
         # 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 @@ def get_async_streaming_response_iterator(
         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 @@ def _chunk_parser(
         """
         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

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