fix: normalize token usage conversion in AmazonBedrockGenerator (#3247)

Author: mathislucka

integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/adapters.py

@@ -5,6 +5,86 @@
 from botocore.eventstream import EventStream
 from haystack.dataclasses import StreamingChunk, SyncStreamingCallbackT
 
+_USAGE_HEADER_MAP = {
+    "input_tokens": "x-amzn-bedrock-input-token-count",
+    "output_tokens": "x-amzn-bedrock-output-token-count",
+    "cache_read_input_tokens": "x-amzn-bedrock-cache-read-input-token-count",
+    "cache_write_input_tokens": "x-amzn-bedrock-cache-write-input-token-count",
+}
+
+_USAGE_FIELD_MAP = {
+    "input_tokens": "input_tokens",
+    "output_tokens": "output_tokens",
+    "cache_read_input_tokens": "cache_read_input_tokens",
+    "cache_write_input_tokens": "cache_creation_input_tokens",
+}
+
+
+def _set_usage_value(usage: dict[str, int], key: str, value: Any) -> None:
+    """
+    Sets a usage value coerced to int, ignoring values that are None or not int-convertible.
+
+    :param usage: The usage dictionary to update in place.
+    :param key: The destination key.
+    :param value: The raw value to coerce and store.
+    """
+    if value is None:
+        return
+    try:
+        usage[key] = int(value)
+    except (TypeError, ValueError):
+        return
+
+
+def _apply_usage(usage: dict[str, int], source: dict[str, Any], field_map: dict[str, str]) -> None:
+    """
+    Copies usage values from a source dictionary into the usage dictionary using the given field map.
+
+    :param usage: The usage dictionary to update in place.
+    :param source: The source dictionary holding raw usage values.
+    :param field_map: A mapping from destination key to source key.
+    """
+    for dst, src in field_map.items():
+        _set_usage_value(usage, dst, source.get(src))
+
+
+def _usage_from_response_metadata(metadata: dict[str, Any]) -> dict[str, int]:
+    """
+    Extracts normalized token usage from Bedrock InvokeModel ResponseMetadata HTTP headers.
+
+    :param metadata: The Bedrock response metadata dictionary.
+    :returns: A normalized usage dictionary, or an empty dictionary when no usage headers are present.
+    """
+    headers = metadata.get("HTTPHeaders") or metadata.get("http_headers") or {}
+    if not isinstance(headers, dict):
+        return {}
+
+    normalized_headers = {str(key).lower(): value for key, value in headers.items()}
+    usage: dict[str, int] = {}
+    _apply_usage(usage, normalized_headers, _USAGE_HEADER_MAP)
+    return usage
+
+
+def _merge_usage(metadata: dict[str, Any], usage: dict[str, int]) -> None:
+    """
+    Merges a usage dictionary into the metadata under the ``usage`` key.
+
+    Recomputes ``total_tokens`` after merging when both ``input_tokens`` and ``output_tokens``
+    are present, so partial usage from multiple sources is summed correctly.
+
+    :param metadata: The metadata dictionary to update in place.
+    :param usage: The normalized usage dictionary to merge in.
+    """
+    if not usage:
+        return
+
+    existing_usage = metadata.get("usage")
+    base = existing_usage if isinstance(existing_usage, dict) else {}
+    merged_usage = {**base, **usage}
+    if "input_tokens" in merged_usage and "output_tokens" in merged_usage:
+        merged_usage["total_tokens"] = merged_usage["input_tokens"] + merged_usage["output_tokens"]
+    metadata["usage"] = merged_usage
+
 
 class BedrockModelAdapter(ABC):
     """
@@ -54,6 +134,20 @@ def get_stream_responses(self, stream: EventStream, streaming_callback: SyncStre
         :param streaming_callback: The handler for the streaming response.
         :returns: A list of string responses.
         """
+        responses, _ = self.get_stream_responses_and_metadata(stream, streaming_callback)
+        return responses
+
+    def get_stream_responses_and_metadata(
+        self, stream: EventStream, streaming_callback: SyncStreamingCallbackT
+    ) -> tuple[list[str], dict[str, Any]]:
+        """
+        Extracts both the responses and normalized metadata from the Amazon Bedrock streaming response.
+
+        :param stream: The streaming response from the Amazon Bedrock request.
+        :param streaming_callback: The handler for the streaming response.
+        :returns: A tuple of ``(responses, metadata)`` where ``responses`` is a list of string
+            responses and ``metadata`` is a dictionary that may contain a normalized ``usage`` block.
+        """
         streaming_chunks: list[StreamingChunk] = []
         for event in stream:
             chunk = event.get("chunk")
@@ -64,7 +158,37 @@ def get_stream_responses(self, stream: EventStream, streaming_callback: SyncStre
                 streaming_callback(streaming_chunk)
 
         responses = ["".join(streaming_chunk.content for streaming_chunk in streaming_chunks).lstrip()]
-        return responses
+        metadata = self._extract_streaming_metadata(streaming_chunks)
+        return responses, metadata
+
+    def _extract_streaming_metadata(self, streaming_chunks: list[StreamingChunk]) -> dict[str, Any]:
+        """
+        Extracts normalized metadata from Bedrock streaming chunks.
+
+        The default implementation handles Anthropic Claude Messages API stream events, which
+        expose input usage in ``message_start.message.usage`` and output usage in
+        ``message_delta.usage``.
+
+        :param streaming_chunks: The streaming chunks emitted during the response.
+        :returns: A metadata dictionary with a ``usage`` block, or an empty dictionary when no
+            usage information is present.
+        """
+        usage: dict[str, int] = {}
+
+        for streaming_chunk in streaming_chunks:
+            meta = streaming_chunk.meta
+            if not isinstance(meta, dict):
+                continue
+            message = meta.get("message")
+            chunk_usage = meta.get("usage")
+            if message is None and chunk_usage is None:
+                continue
+            if isinstance(message, dict) and isinstance(message.get("usage"), dict):
+                _apply_usage(usage, message["usage"], _USAGE_FIELD_MAP)
+            if isinstance(chunk_usage, dict):
+                _apply_usage(usage, chunk_usage, _USAGE_FIELD_MAP)
+
+        return {"usage": usage} if usage else {}
 
     def _get_params(self, inference_kwargs: dict[str, Any], default_params: dict[str, Any]) -> dict[str, Any]:
         """

integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py

@@ -25,6 +25,8 @@
     CohereCommandRAdapter,
     MetaLlamaAdapter,
     MistralAdapter,
+    _merge_usage,
+    _usage_from_response_metadata,
 )
 
 logger = logging.getLogger(__name__)
@@ -215,6 +217,7 @@ def run(
         generation_kwargs["stream"] = streaming_callback is not None
 
         body = self.model_adapter.prepare_body(prompt=prompt, **generation_kwargs)
+        stream_metadata: dict[str, Any] = {}
         try:
             if streaming_callback:
                 response = self.client.invoke_model_with_response_stream(
@@ -224,7 +227,7 @@ def run(
                     contentType="application/json",
                 )
                 response_stream = response["body"]
-                replies = self.model_adapter.get_stream_responses(
+                replies, stream_metadata = self.model_adapter.get_stream_responses_and_metadata(
                     stream=response_stream, streaming_callback=streaming_callback
                 )
             else:
@@ -238,6 +241,8 @@ def run(
                 replies = self.model_adapter.get_responses(response_body=response_body)
 
             metadata = response.get("ResponseMetadata", {})
+            _merge_usage(metadata, _usage_from_response_metadata(metadata))
+            _merge_usage(metadata, stream_metadata.get("usage", {}))
 
         except ClientError as exception:
             msg = f"Could not perform inference for Amazon Bedrock model {self.model} due to:\n{exception}"

integrations/amazon_bedrock/tests/test_generator.py

@@ -1,3 +1,5 @@
+import json
+from io import BytesIO
 from typing import Any
 from unittest.mock import MagicMock, call
 
@@ -336,6 +338,137 @@ def test_run_client_error(mock_boto3_session):
         generator.run("Hello")
 
 
+def test_run_non_streaming_normalizes_usage_from_headers(mock_boto3_session):
+    generator = AmazonBedrockGenerator(model="anthropic.claude-v2")
+    mock_client = mock_boto3_session.return_value.client.return_value
+    mock_client.invoke_model.return_value = {
+        "body": BytesIO(json.dumps({"content": [{"type": "text", "text": "ok"}]}).encode()),
+        "ResponseMetadata": {
+            "HTTPHeaders": {
+                "x-amzn-bedrock-input-token-count": "20",
+                "x-amzn-bedrock-output-token-count": "10",
+                "x-amzn-bedrock-cache-read-input-token-count": "0",
+                "x-amzn-bedrock-cache-write-input-token-count": "0",
+            }
+        },
+    }
+
+    result = generator.run("hi")
+
+    assert result["replies"] == ["ok"]
+    assert result["meta"]["usage"] == {
+        "input_tokens": 20,
+        "output_tokens": 10,
+        "total_tokens": 30,
+        "cache_read_input_tokens": 0,
+        "cache_write_input_tokens": 0,
+    }
+
+
+def test_run_non_streaming_without_usage_headers_omits_usage(mock_boto3_session):
+    generator = AmazonBedrockGenerator(model="anthropic.claude-v2")
+    mock_client = mock_boto3_session.return_value.client.return_value
+    mock_client.invoke_model.return_value = {
+        "body": BytesIO(json.dumps({"content": [{"type": "text", "text": "ok"}]}).encode()),
+        "ResponseMetadata": {"HTTPHeaders": {}},
+    }
+
+    result = generator.run("hi")
+
+    assert "usage" not in result["meta"]
+
+
+def test_run_streaming_normalizes_anthropic_usage(mock_boto3_session):
+    generator = AmazonBedrockGenerator(model="anthropic.claude-v2")
+    mock_client = mock_boto3_session.return_value.client.return_value
+
+    stream_body = MagicMock()
+    stream_body.__iter__.return_value = [
+        {
+            "chunk": {
+                "bytes": json.dumps(
+                    {
+                        "type": "message_start",
+                        "message": {"usage": {"input_tokens": 20, "output_tokens": 1}},
+                    }
+                ).encode()
+            }
+        },
+        {"chunk": {"bytes": json.dumps({"type": "content_block_delta", "delta": {"text": "ok"}}).encode()}},
+        {"chunk": {"bytes": json.dumps({"type": "message_delta", "usage": {"output_tokens": 10}}).encode()}},
+    ]
+    mock_client.invoke_model_with_response_stream.return_value = {
+        "body": stream_body,
+        "ResponseMetadata": {"RequestId": "req-1"},
+    }
+
+    result = generator.run("hi", streaming_callback=lambda chunk: None)
+
+    assert result["replies"] == ["ok"]
+    assert result["meta"]["usage"]["input_tokens"] == 20
+    assert result["meta"]["usage"]["output_tokens"] == 10
+    assert result["meta"]["usage"]["total_tokens"] == 30
+
+
+def test_run_streaming_with_cache_usage(mock_boto3_session):
+    generator = AmazonBedrockGenerator(model="anthropic.claude-v2")
+    mock_client = mock_boto3_session.return_value.client.return_value
+
+    stream_body = MagicMock()
+    stream_body.__iter__.return_value = [
+        {
+            "chunk": {
+                "bytes": json.dumps(
+                    {
+                        "type": "message_start",
+                        "message": {
+                            "usage": {
+                                "input_tokens": 5,
+                                "output_tokens": 1,
+                                "cache_read_input_tokens": 100,
+                                "cache_creation_input_tokens": 50,
+                            }
+                        },
+                    }
+                ).encode()
+            }
+        },
+        {"chunk": {"bytes": json.dumps({"type": "message_delta", "usage": {"output_tokens": 7}}).encode()}},
+    ]
+    mock_client.invoke_model_with_response_stream.return_value = {
+        "body": stream_body,
+        "ResponseMetadata": {},
+    }
+
+    result = generator.run("hi", streaming_callback=lambda chunk: None)
+
+    assert result["meta"]["usage"] == {
+        "input_tokens": 5,
+        "output_tokens": 7,
+        "total_tokens": 12,
+        "cache_read_input_tokens": 100,
+        "cache_write_input_tokens": 50,
+    }
+
+
+def test_run_streaming_without_usage_omits_usage(mock_boto3_session):
+    generator = AmazonBedrockGenerator(model="anthropic.claude-v2")
+    mock_client = mock_boto3_session.return_value.client.return_value
+
+    stream_body = MagicMock()
+    stream_body.__iter__.return_value = [
+        {"chunk": {"bytes": b'{"delta": {"text": "ok"}}'}},
+    ]
+    mock_client.invoke_model_with_response_stream.return_value = {
+        "body": stream_body,
+        "ResponseMetadata": {"RequestId": "req-1"},
+    }
+
+    result = generator.run("hi", streaming_callback=lambda chunk: None)
+
+    assert "usage" not in result["meta"]
+
+
 def test_from_dict_with_streaming_callback(mock_boto3_session):
     data = {
         "type": "haystack_integrations.components.generators.amazon_bedrock.generator.AmazonBedrockGenerator",

integrations/amazon_bedrock/tests/test_generator_inference.py

@@ -0,0 +1,60 @@
+import os
+
+import pytest
+from haystack.utils import Secret
+
+from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockGenerator
+
+MODELS_TO_TEST = [
+    "global.anthropic.claude-haiku-4-5-20251001-v1:0",
+]
+
+
+def _generator(model: str) -> AmazonBedrockGenerator:
+    return AmazonBedrockGenerator(
+        model=model,
+        max_length=64,
+        aws_region_name=Secret.from_token(os.environ["AWS_REGION"]),
+    )
+
+
+def _assert_usage(usage: dict) -> None:
+    assert isinstance(usage["input_tokens"], int) and usage["input_tokens"] > 0
+    assert isinstance(usage["output_tokens"], int) and usage["output_tokens"] > 0
+    assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
+
+
+@pytest.mark.integration
+@pytest.mark.skipif(
+    not os.getenv("AWS_BEARER_TOKEN_BEDROCK") or not os.getenv("AWS_REGION"),
+    reason="AWS_BEARER_TOKEN_BEDROCK and AWS_REGION must be set",
+)
+class TestAmazonBedrockGeneratorInference:
+    @pytest.mark.parametrize("model", MODELS_TO_TEST)
+    def test_run_non_streaming_normalizes_usage(self, model: str) -> None:
+        generator = _generator(model)
+        result = generator.run("What is the capital of France? Reply in one word.")
+
+        assert result["replies"], "No replies received"
+        assert isinstance(result["replies"][0], str) and result["replies"][0]
+
+        meta = result["meta"]
+        assert "usage" in meta, f"meta does not contain a normalized 'usage' block: {meta}"
+        _assert_usage(meta["usage"])
+
+    @pytest.mark.parametrize("model", MODELS_TO_TEST)
+    def test_run_streaming_normalizes_usage(self, model: str) -> None:
+        generator = _generator(model)
+        chunks: list = []
+        result = generator.run(
+            "What is the capital of France? Reply in one word.",
+            streaming_callback=chunks.append,
+        )
+
+        assert chunks, "Streaming callback was not invoked"
+        assert result["replies"], "No replies received"
+        assert isinstance(result["replies"][0], str) and result["replies"][0]
+
+        meta = result["meta"]
+        assert "usage" in meta, f"meta does not contain a normalized 'usage' block: {meta}"
+        _assert_usage(meta["usage"])