feat(bedrock): Set token attributes for non-streaming responses

pyproject.toml

@@ -82,6 +82,7 @@ typing = [
     "openai-agents>=0.22.0",
     "mistralai>=2.10.1",
     "typesafe-sdk>=0.7.5",
+    "mypy-boto3-bedrock-runtime>=1.43.62",
 ]
 test = [
     "dataclasses ; python_full_version < '3.7'",

sentry_sdk/integrations/boto3/_services/bedrock.py

@@ -1,14 +1,16 @@
 import json
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
 from sentry_sdk.utils import has_data_collection_enabled
 
 if TYPE_CHECKING:
     from typing import Any, Optional
 
+    from mypy_boto3_bedrock_runtime.type_defs import ConverseResponseTypeDef
+
     from sentry_sdk._types import Attributes, ToolDefinition
     from sentry_sdk.integrations.boto3._context import AwsCallContext
 
@@ -122,4 +124,29 @@
         ctx: "AwsCallContext",
         response: "Any",
     ) -> "Attributes":
-        return {}
+        if ctx.operation_name != "Converse" or "Error" in response:
+            return {}
+
+        usage = cast("ConverseResponseTypeDef", response)["usage"]
+
+        input_tokens = usage["inputTokens"]
+        attributes: "Attributes" = {
+            SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: usage["outputTokens"],
+            SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS: usage["totalTokens"],
+        }
+
+        cache_read_tokens = usage.get("cacheReadInputTokens")
+        if cache_read_tokens is not None:
+            input_tokens += cache_read_tokens
+            attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] = cache_read_tokens
+
+        cache_write_tokens = usage.get("cacheWriteInputTokens")
+        if cache_write_tokens is not None:
+            input_tokens += cache_write_tokens
+            attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] = (
+                cache_write_tokens
+            )
+
+        attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] = input_tokens
+
+        return attributes

tests/integrations/boto3/test_bedrock.py

@@ -22,10 +22,15 @@
 
 @pytest.fixture
 def converse_response():
+    usage = {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}
+    if BOTOCORE_VERSION >= Version("1.37.24"):
+        usage["cacheReadInputTokens"] = 80
+        usage["cacheWriteInputTokens"] = 20
+
     return {
         "output": {"message": {"role": "assistant", "content": [{"text": "Hello"}]}},
         "stopReason": "end_turn",
-        "usage": {"inputTokens": 50, "outputTokens": 30, "totalTokens": 80},
+        "usage": usage,
         "metrics": {"latencyMs": 500},
     }
 
@@ -87,6 +92,15 @@
         assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 1024
         assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
         assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9
+
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30
+        if BOTOCORE_VERSION < Version("1.37.24"):
+            assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
+            return
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 110
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20
     else:
         events = capture_events()
 
@@ -116,6 +130,15 @@
         assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
         assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9
 
+        assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20
+        assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30
+        if BOTOCORE_VERSION < Version("1.37.24"):
+            assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
+            return
+        assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 110
+        assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80
+        assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20
+
 
 @pytest.mark.parametrize("span_streaming", [True, False])
 @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])