feat(bedrock): Set gen_ai.output.messages for non-streaming responses

sentry_sdk/integrations/boto3/_services/bedrock.py

@@ -23,7 +23,7 @@
 
 
 if TYPE_CHECKING:
-    from typing import Any, Iterator, Optional, Union
+    from typing import Any, Iterator, Mapping, Optional, Union
 
     from mypy_boto3_bedrock_runtime.type_defs import (
         ConverseResponseTypeDef,
@@ -171,7 +171,7 @@ def _transform_system_instructions(
 
 
 def _transform_messages(
-    input_messages: "Union[list[dict[str, Any]], tuple[dict[str, Any], ...]]",
+    input_messages: "Union[list[Mapping[str, Any]], tuple[Mapping[str, Any], ...]]",
 ) -> "list[InputMessage]":
     """
     Unlike other providers, AWS declares that input and output messages use the same message type.
@@ -488,7 +488,8 @@ def on_success_response(
         )
 
         if ctx.operation_name == "Converse":
-            usage = cast("ConverseResponseTypeDef", response)["usage"]
+            response = cast("ConverseResponseTypeDef", response)
+            usage = response["usage"]
 
             set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"])
             set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage["totalTokens"])
@@ -511,6 +512,22 @@ def on_success_response(
 
             set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
 
+            client = sentry_sdk.get_client()
+            if "message" in response["output"] and (
+                (
+                    has_data_collection_enabled(client.options)
+                    and client.options["data_collection"]["gen_ai"]["outputs"]
+                )
+                or (
+                    not has_data_collection_enabled(client.options)
+                    and should_send_default_pii()
+                )
+            ):
+                set_on_span(
+                    SPANDATA.GEN_AI_OUTPUT_MESSAGES,
+                    json.dumps(_transform_messages([response["output"]["message"]])),
+                )
+
             if isinstance(span, StreamedSpan):
                 span.end()
                 return

tests/integrations/boto3/test_bedrock.py

@@ -834,6 +834,96 @@ def test_converse_input_attributes(
         )
 
 
+@pytest.mark.parametrize("data_collection", [True, False])
+@pytest.mark.parametrize("span_streaming", [True, False])
+@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
+def test_converse_output_attributes(
+    sentry_init,
+    capture_items,
+    capture_events,
+    data_collection,
+    span_streaming,
+    stream_gen_ai_spans,
+    converse_response,
+    content_block_messages,
+):
+    if data_collection:
+        sentry_init(
+            traces_sample_rate=1.0,
+            integrations=[Boto3Integration()],
+            data_collection={"gen_ai": {"outputs": True}},
+            stream_gen_ai_spans=stream_gen_ai_spans,
+            trace_lifecycle="stream" if span_streaming else "static",
+        )
+    else:
+        sentry_init(
+            traces_sample_rate=1.0,
+            integrations=[Boto3Integration()],
+            send_default_pii=True,
+            stream_gen_ai_spans=stream_gen_ai_spans,
+            trace_lifecycle="stream" if span_streaming else "static",
+        )
+
+    messages, expected_messages = content_block_messages
+    converse_response["output"]["message"] = {
+        "role": "assistant",
+        "content": [block for message in messages for block in message["content"]],
+    }
+    expected_output_messages = [
+        {
+            "role": "assistant",
+            "parts": [
+                part for message in expected_messages for part in message["parts"]
+            ],
+        }
+    ]
+
+    session = boto3.Session(
+        aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1"
+    )
+    client = session.client("bedrock-runtime")
+    stubber = Stubber(client)
+    stubber.add_response("converse", converse_response)
+
+    if span_streaming or stream_gen_ai_spans:
+        items = capture_items("span")
+
+        with stubber, sentry_sdk.start_transaction(name="parent"):
+            client.converse(
+                modelId="anthropic.claude-3-sonnet-20240229-v1:0",
+                messages=[{"role": "user", "content": [{"text": "Hi"}]}],
+            )
+
+        sentry_sdk.flush()
+        spans = [item.payload for item in items]
+        (span,) = (
+            span
+            for span in spans
+            if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT
+        )
+
+        assert (
+            json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES])
+            == expected_output_messages
+        )
+    else:
+        events = capture_events()
+
+        with stubber, sentry_sdk.start_transaction(name="parent"):
+            client.converse(
+                modelId="anthropic.claude-3-sonnet-20240229-v1:0",
+                messages=[{"role": "user", "content": [{"text": "Hi"}]}],
+            )
+
+        (transaction,) = events
+        (span,) = transaction["spans"]
+
+        assert (
+            json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES])
+            == expected_output_messages
+        )
+
+
 @pytest.mark.parametrize("span_streaming", [True, False])
 @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
 def test_converse_stream(