feat(boto3): Instrument Bedrock Runtime InvokeModel calls as gen_ai spans

sentry_sdk/consts.py

@@ -1321,6 +1321,7 @@ class OP:
     GEN_AI_EMBEDDINGS = "gen_ai.embeddings"
     GEN_AI_EVALUATE = "gen_ai.evaluate"
     GEN_AI_EXECUTE_TOOL = "gen_ai.execute_tool"
+    GEN_AI_GENERATE_CONTENT = "gen_ai.generate_content"
     GEN_AI_TEXT_COMPLETION = "gen_ai.text_completion"
     GEN_AI_HANDOFF = "gen_ai.handoff"
     GEN_AI_INVOKE_AGENT = "gen_ai.invoke_agent"

sentry_sdk/integrations/boto3/_client.py

@@ -144,9 +144,14 @@ def sentry_patched_make_api_call(
 
         streaming_body_instrumented = False
         with capture_internal_exceptions():
-            streaming_body_instrumented = _instrument_streaming_body(span, parsed)
-
-        # `StreamingBody`s finish their span when consumed or closed.
+            if service_ext is not None:
+                streaming_body_instrumented = service_ext.instrument_response(
+                    ctx, span, parsed
+                )
+            if not streaming_body_instrumented:
+                streaming_body_instrumented = _instrument_streaming_body(span, parsed)
+
+        # streaming responses finish their span when consumed or closed.
         if not streaming_body_instrumented:
             _finish_span(span)
         return parsed

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -25,7 +25,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Dict, Mapping, Optional, Union
+    from typing import Any, Callable, Dict, Mapping, Optional, Union
 
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
@@ -294,17 +294,8 @@ def _instrument_streaming_body(
             origin=ORIGIN,
         )
 
-    finished = False
-    read_in_progress = False
-
     def finish_span(error: "Optional[BaseException]" = None) -> None:
-        nonlocal finished
-        if finished:
-            return
-
-        finished = True
-        # finish the stream span before the client span, and only once across
-        # read and close.
+        # finish the stream span before the client span.
         if error is not None:
             with capture_internal_exceptions():
                 attributes = _get_error_attributes(error)
@@ -314,6 +305,32 @@ def finish_span(error: "Optional[BaseException]" = None) -> None:
         _finish_span(streaming_span, error)
         _finish_span(span, error)
 
+    _patch_streaming_body(body, finish_span)
+    return True
+
+
+def _patch_streaming_body(
+    body: "StreamingBody",
+    on_finish: "Callable[[Optional[BaseException]], None]",
+    on_read: "Optional[Callable[[bytes], None]]" = None,
+) -> None:
+    """
+    Call `on_finish` once, when `body` is consumed, closed, or fails to be read.
+    `on_read` receives the bytes returned by each `read()`.
+    """
+    finished = False
+    read_in_progress = False
+
+    def finish(error: "Optional[BaseException]" = None) -> None:
+        nonlocal finished
+        if finished:
+            return
+
+        # only finish once across read and close.
+        finished = True
+        with capture_internal_exceptions():
+            on_finish(error)
+
     def content_length_reached() -> bool:
         content_length = body._content_length  # type: ignore[attr-defined]
         return content_length is not None and body._amount_read >= int(content_length)  # type: ignore[attr-defined]
@@ -323,6 +340,9 @@ def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes:
         read_in_progress = True
         try:
             read_return_value = orig_read(*args, **kwargs)
+            if on_read is not None:
+                with capture_internal_exceptions():
+                    on_read(read_return_value)
             with capture_internal_exceptions():
                 amount_of_bytes_requested = args[0] if args else kwargs.get("amt")
                 # detect read-to-end, eof, or the known content length being consumed.
@@ -332,29 +352,29 @@ def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes:
                     or (amount_of_bytes_requested > 0 and not read_return_value)
                     or content_length_reached()
                 ):
-                    finish_span()
+                    finish()
             return read_return_value
         except BaseException as error:
-            finish_span(error)
+            finish(error)
             raise
         finally:
             read_in_progress = False
 
     def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None:
         try:
             orig_close(*args, **kwargs)
-            finish_span()
+            finish()
         except BaseException as error:
-            finish_span(error)
+            finish(error)
             raise
 
     def sentry_raw_stream_close(*args: "Any", **kwargs: "Any") -> None:
         try:
             orig_raw_close(*args, **kwargs)
             if not read_in_progress:
-                finish_span()
+                finish()
         except BaseException as error:
-            finish_span(error)
+            finish(error)
             raise
 
     try:
@@ -367,11 +387,9 @@ def sentry_raw_stream_close(*args: "Any", **kwargs: "Any") -> None:
         body.read = sentry_streaming_body_read  # type: ignore
         body.close = sentry_streaming_body_close  # type: ignore
     except Exception:
-        finish_span()
+        finish()
         raise
 
-    return True
-
 
 def _set_request_attributes(
     span: "Union[Span, StreamedSpan]",

sentry_sdk/integrations/boto3/_services/base.py

@@ -50,3 +50,16 @@ def get_response_attributes(
     ) -> "Attributes":
         """Return response attributes to add after the AWS request is made."""
         return {}
+
+    def instrument_response(
+        self,
+        ctx: "AwsCallContext",
+        span: "Union[Span, StreamedSpan]",
+        response: "Any",
+    ) -> bool:
+        """
+        Instrument a response that is read after the AWS request is made, e.g. a
+        streaming body. Return `True` if the extension finishes `span` once the
+        response is read, or `False` to finish it when the request is made.
+        """
+        return False