feat(bedrock): Set token attributes for nonstreaming responses

Author: alexander-alderman-webb

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,11 +1,13 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
 
 if TYPE_CHECKING:
     from typing import Any, Optional
 
+    from mypy_boto3_bedrock_runtime.type_defs import ConverseResponseTypeDef
+
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
 
@@ -39,4 +41,33 @@ def get_response_attributes(
         ctx: "AwsCallContext",
         response: "Any",
     ) -> "Attributes":
-        return {}
+        if ctx.operation_name != "Converse":
+            return {}
+
+        attributes: "Attributes" = {}
+
+        usage = cast("ConverseResponseTypeDef", response)["usage"]
+
+        input_tokens = usage["inputTokens"]
+        if input_tokens is not None:
+            attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] = input_tokens
+
+        cache_read_tokens = usage.get("cacheReadInputTokens")
+        if cache_read_tokens is not None:
+            attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] = cache_read_tokens
+
+        cache_write_tokens = usage.get("cacheWriteInputTokens")
+        if cache_write_tokens is not None:
+            attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] = (
+                cache_write_tokens
+            )
+
+        output_tokens = usage["outputTokens"]
+        if output_tokens is not None:
+            attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] = output_tokens
+
+        total_tokens = usage["totalTokens"]
+        if total_tokens is not None:
+            attributes[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] = total_tokens
+
+        return attributes

tests/integrations/boto3/test_bedrock.py

@@ -1,13 +1,17 @@
 import boto3
+import botocore
 import pytest
 from botocore.exceptions import ClientError
 from botocore.loaders import Loader
 from botocore.stub import Stubber
+from packaging.version import Version
 
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.boto3 import Boto3Integration
 
+BOTOCORE_VERSION = Version(botocore.__version__)
+
 pytestmark = pytest.mark.skipif(
     "bedrock-runtime" not in Loader().list_available_services("service-2"),
     reason="botocore is too old for bedrock-runtime.",
@@ -16,10 +20,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},
     }
 
@@ -68,6 +77,13 @@ def test_converse(
         assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock"
         assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
         assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
+        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"):
+            return
+        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()
 
@@ -83,6 +99,13 @@ def test_converse(
         assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock"
         assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
         assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
+        assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10
+        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"):
+            return
+        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])