diff --git a/sentry_sdk/integrations/boto3/_services/bedrock.py b/sentry_sdk/integrations/boto3/_services/bedrock.py index 5c9b16325e..d35fc26efb 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock.py +++ b/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 diff --git a/tests/integrations/boto3/test_bedrock.py b/tests/integrations/boto3/test_bedrock.py index 9461bce69c..4696293eda 100644 --- a/tests/integrations/boto3/test_bedrock.py +++ b/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(