feat(bedrock): Set token attributes for streamed responses

sentry_sdk/integrations/boto3/__init__.py

@@ -1,5 +1,6 @@
 from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
 from sentry_sdk.integrations.boto3._client import _patch_botocore_client
+from sentry_sdk.integrations.boto3._services.bedrock import _patch_event_stream
 from sentry_sdk.integrations.boto3.consts import IDENTIFIER, ORIGIN
 from sentry_sdk.utils import parse_version
 
@@ -19,3 +20,4 @@ def setup_once() -> None:
         _check_minimum_version(Boto3Integration, version, "botocore")
 
         _patch_botocore_client()
+        _patch_event_stream()

sentry_sdk/integrations/boto3/_client.py

@@ -27,7 +27,6 @@
 
 try:
     from botocore.client import BaseClient
-    from botocore.exceptions import ClientError
 except ImportError:
     raise DidNotEnable("botocore not installed")
 
@@ -121,21 +120,11 @@
                 try:
                     parsed = orig_make_api_call(self, operation_name, api_params)
                 except BaseException as error:
-                    if service_ext is not None and isinstance(error, ClientError):
-                        with capture_internal_exceptions():
-                            attributes.update(
-                                service_ext.get_response_attributes(ctx, error.response)
-                            )
                     with capture_internal_exceptions():
                         attributes.update(_get_error_attributes(error))
                     with capture_internal_exceptions():
                         _set_span_attributes(span, attributes)
                     raise
-                if service_ext is not None:
-                    with capture_internal_exceptions():
-                        attributes.update(
-                            service_ext.get_response_attributes(ctx, parsed)
-                        )
                 with capture_internal_exceptions():
                     attributes.update(_get_response_attributes(parsed))
                 with capture_internal_exceptions():
@@ -144,6 +133,10 @@
             _finish_span(span, error)
             raise
 
+        if service_ext is not None:
+            service_ext.on_success_response(ctx, parsed, span)
+            return parsed
+
         streaming_body_instrumented = False
         with capture_internal_exceptions():
             streaming_body_instrumented = _instrument_streaming_body(span, parsed)

sentry_sdk/integrations/boto3/_services/base.py

@@ -1,10 +1,12 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Optional, Union
 
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.traces import StreamedSpan
+    from sentry_sdk.tracing import Span
 
 
 class _ServiceExtension:
@@ -30,8 +32,11 @@
         """Return request attributes to add before the AWS request is made."""
         return {}
 
-    def get_response_attributes(
-        self, ctx: "AwsCallContext", response: "Any"
-    ) -> "Attributes":
-        """Return response attributes to add after the AWS request is made."""
-        return {}
+    def on_success_response(
+        self,
+        ctx: "AwsCallContext",
+        response: "Any",
+        span: "Union[Span, StreamedSpan]",
+    ) -> None:
+        """Hook that fires after a successful response. Responsible for closing the span."""
+        return None

sentry_sdk/integrations/boto3/_services/bedrock.py

@@ -1,18 +1,149 @@
 import json
+import sys
 from typing import TYPE_CHECKING, cast
 
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
-from sentry_sdk.utils import has_data_collection_enabled
+from sentry_sdk.traces import SpanStatus, StreamedSpan
+from sentry_sdk.tracing import SPANSTATUS
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    has_data_collection_enabled,
+    reraise,
+    should_be_treated_as_error,
+)
+
+try:
+    from botocore.eventstream import EventStream
+except ImportError:
+    raise DidNotEnable("botocore is not installed")
+
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Iterator, Optional, Union
 
-    from mypy_boto3_bedrock_runtime.type_defs import ConverseResponseTypeDef
+    from mypy_boto3_bedrock_runtime.type_defs import (
+        ConverseResponseTypeDef,
+        ConverseStreamResponseTypeDef,
+    )
 
     from sentry_sdk._types import Attributes, ToolDefinition
     from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.tracing import Span
+
+
+def _finish_span(
+    stream: "EventStream[Any]",
+    exc_type: "Optional[type[BaseException]]",
+    exc_val: "Optional[BaseException]",
+    exc_tb: "Optional[Any]",
+) -> None:
+    """
+    Finishes the span attached to the stream.
+    Is a no-op if the stream has no span set, i.e., when the span has already been finished.
+    """
+    if not hasattr(stream, "_span"):
+        return
+
+    span = stream._span
+    del stream._span
+
+    if exc_type is not None and should_be_treated_as_error(exc_type, exc_val):
+        if isinstance(span, StreamedSpan):
+            span.status = SpanStatus.ERROR
+        else:
+            span.status = SPANSTATUS.INTERNAL_ERROR
+
+    if isinstance(span, StreamedSpan):
+        span.end()
+    else:
+        span.finish()
+
+
+def _patch_event_stream() -> None:
+    original_iter = EventStream.__iter__
+    original_close = EventStream.close
+
+    def sentry_patched_iter(
+        self: "EventStream[Any]",
+    ) -> "Iterator[Any]":
+        """
+        Sets information received while iterating the response stream on the AI Client Span.
+        Responsible for closing the AI Client Span unless the span has already been closed in the close() patch.
+        """
+        span = getattr(self, "_span", None)
+        if span is None:
+            yield from original_iter(self)
+            return
+
+        set_on_span = (
+            span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
+        )
+
+        try:
+            for event in original_iter(self):
+                if "metadata" not in event:
+                    yield event
+                    continue
+
+                usage = event["metadata"]["usage"]
+
+                input_tokens = usage["inputTokens"]
+                set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"])
+                set_on_span(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
+                    set_on_span(
+                        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
+                    set_on_span(
+                        SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE,
+                        cache_write_tokens,
+                    )
+
+                set_on_span(
+                    SPANDATA.GEN_AI_USAGE_INPUT_TOKENS,
+                    input_tokens,
+                )
+
+                yield event
+        except GeneratorExit:
+            raise  # Leave span open since underlying stream is not necessarily closed.
+        except BaseException:
+            exc_info = sys.exc_info()
+            with capture_internal_exceptions():
+                _finish_span(self, *exc_info)
+            reraise(*exc_info)
+
+        _finish_span(self, None, None, None)
+
+    def sentry_patched_close(self: "EventStream[Any]") -> None:
+        """
+        Closes the AI Client Span unless `sentry_patched_iter()` closes the span first.
+        """
+        if not hasattr(self, "_span"):
+            return original_close(self)
+
+        try:
+            original_close(self)
+        except BaseException:
+            exc_info = sys.exc_info()
+            with capture_internal_exceptions():
+                _finish_span(self, *exc_info)
+            reraise(*exc_info)
+
+        _finish_span(self, None, None, None)
+
+    EventStream.__iter__ = sentry_patched_iter  # type: ignore[method-assign]
+    EventStream.close = sentry_patched_close  # type: ignore[method-assign]
 
 
 def _transform_tool_definitions(
@@ -119,34 +250,54 @@
 
         return attributes
 
-    def get_response_attributes(
+    def on_success_response(
         self,
         ctx: "AwsCallContext",
         response: "Any",
-    ) -> "Attributes":
-        if ctx.operation_name != "Converse" or "Error" in response:
-            return {}
+        span: "Union[Span, StreamedSpan]",
+    ) -> None:
+        set_on_span = (
+            span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
+        )
 
-        usage = cast("ConverseResponseTypeDef", response)["usage"]
+        if ctx.operation_name == "Converse":
+            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"],
-        }
+            set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"])
+            set_on_span(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
+            input_tokens = usage["inputTokens"]
 
-        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
-            )
+            cache_read_tokens = usage.get("cacheReadInputTokens")
+            if cache_read_tokens is not None:
+                input_tokens += cache_read_tokens
+                set_on_span(
+                    SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, cache_read_tokens
+                )
 
-        attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] = input_tokens
+            cache_write_tokens = usage.get("cacheWriteInputTokens")
+            if cache_write_tokens is not None:
+                input_tokens += cache_write_tokens
+                set_on_span(
+                    SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, cache_write_tokens
+                )
 
-        return attributes
+            set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
+
+            if isinstance(span, StreamedSpan):
+                span.end()
+                return
+
+            span.finish()
+            return
+
+        if ctx.operation_name == "ConverseStream":
+            stream = cast("ConverseStreamResponseTypeDef", response)["stream"]
+            stream._span = span  # type: ignore[attr-defined]
+            return
+
+        if isinstance(span, StreamedSpan):
+            span.end()
+            return
+
+        span.finish()