From db1a5a41f407a72d4b031f8700a32aef3c081a12 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 11 Nov 2026 18:09:47 +0100 Subject: [PATCH 1/4] feat(bedrock): Set gen_ai.output.messages for streamed responses --- .../integrations/boto3/_services/bedrock.py | 135 +++++++++++--- tests/integrations/boto3/test_bedrock.py | 165 ++++++++++++++++++ 2 files changed, 280 insertions(+), 20 deletions(-) diff --git a/sentry_sdk/integrations/boto3/_services/bedrock.py b/sentry_sdk/integrations/boto3/_services/bedrock.py index e7b7ddcf2e..6f663a6974 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock.py @@ -93,37 +93,132 @@ def sentry_patched_iter( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) + role: "Optional[str]" = None + content_blocks: "dict[int, dict[str, Any]]" = {} + try: for event in original_iter(self): - if "metadata" not in event: - yield event - continue + if "metadata" in event: + usage = event["metadata"]["usage"] - 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"] + ) - 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, + ) - 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 + 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 + yield event + continue + + client = sentry_sdk.get_client() + if ( + not has_data_collection_enabled(client.options) + or not client.options["data_collection"]["gen_ai"]["outputs"] + ) and ( + has_data_collection_enabled(client.options) + or not should_send_default_pii() + ): + yield event + continue + + if "messageStart" in event: + role = event["messageStart"]["role"] + + yield event + continue + + if "contentBlockStart" in event: + block_start = event["contentBlockStart"] + tool_use = block_start["start"].get("toolUse") + if tool_use is not None: + tool_use_block: "dict[str, Any]" = {"input": ""} + if "name" in tool_use: + tool_use_block["name"] = tool_use["name"] + content_blocks[block_start["contentBlockIndex"]] = { + "toolUse": tool_use_block + } + + yield event + continue + + if "contentBlockDelta" in event: + index = event["contentBlockDelta"]["contentBlockIndex"] + delta = event["contentBlockDelta"]["delta"] + + if "text" in delta: + block = content_blocks.setdefault(index, {"text": ""}) + block["text"] += delta["text"] + elif "toolUse" in delta: + block = content_blocks.setdefault( + index, {"toolUse": {"input": ""}} + ) + block["toolUse"]["input"] += delta["toolUse"]["input"] + elif "reasoningContent" in delta: + block = content_blocks.setdefault( + index, {"reasoningContent": {}} + ) + reasoning_text = delta["reasoningContent"].get("text") + if reasoning_text is not None: + block["reasoningContent"].setdefault( + "reasoningText", {"text": ""} + )["text"] += reasoning_text + + yield event + continue + + if "messageStop" in event and role is not None: + for block in content_blocks.values(): + if "toolUse" not in block: + continue + # Tool input is streamed as a JSON string; non-streaming responses contain the parsed object. + try: + block["toolUse"]["input"] = json.loads( + block["toolUse"]["input"] + ) + except ValueError: + pass + set_on_span( - SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, - cache_write_tokens, + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps( + _transform_messages( + [ + { + "role": role, + "content": [ + content_blocks[index] + for index in sorted(content_blocks) + ], + } + ] + ) + ), ) - set_on_span( - SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, - input_tokens, - ) + yield event + continue yield event except GeneratorExit: diff --git a/tests/integrations/boto3/test_bedrock.py b/tests/integrations/boto3/test_bedrock.py index 3ecd3f357e..cb7a57ff66 100644 --- a/tests/integrations/boto3/test_bedrock.py +++ b/tests/integrations/boto3/test_bedrock.py @@ -1320,3 +1320,168 @@ def test_converse_stream_input_attributes( json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == expected_input_messages ) + + +@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_stream_output_attributes( + sentry_init, + capture_items, + capture_events, + data_collection, + span_streaming, + stream_gen_ai_spans, + converse_stream_response, +): + 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", + ) + + events = [ + ("messageStart", {"role": "assistant"}), + ( + "contentBlockDelta", + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"text": "The user wants "}}, + }, + ), + ( + "contentBlockDelta", + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"text": "the weather."}}, + }, + ), + ( + "contentBlockDelta", + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"signature": "sig"}}, + }, + ), + ("contentBlockStop", {"contentBlockIndex": 0}), + ( + "contentBlockDelta", + {"contentBlockIndex": 1, "delta": {"text": "Let me "}}, + ), + ( + "contentBlockDelta", + {"contentBlockIndex": 1, "delta": {"text": "check."}}, + ), + ("contentBlockStop", {"contentBlockIndex": 1}), + ( + "contentBlockStart", + { + "contentBlockIndex": 2, + "start": {"toolUse": {"toolUseId": "tool-1", "name": "get_weather"}}, + }, + ), + ( + "contentBlockDelta", + {"contentBlockIndex": 2, "delta": {"toolUse": {"input": '{"city": '}}}, + ), + ( + "contentBlockDelta", + {"contentBlockIndex": 2, "delta": {"toolUse": {"input": '"Paris"}'}}}, + ), + ("contentBlockStop", {"contentBlockIndex": 2}), + ("messageStop", {"stopReason": "tool_use"}), + ( + "metadata", + { + "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}, + "metrics": {"latencyMs": 500}, + }, + ), + ] + converse_stream_response["stream"]._event_generator = iter( + EventStreamMessage( + prelude=None, + headers={":message-type": "event", ":event-type": event_type}, + payload=json.dumps(payload).encode(), + crc=None, + ) + for event_type, payload in events + ) + + expected_output_messages = [ + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "The user wants the weather."}, + {"type": "text", "content": "Let me check."}, + { + "type": "tool_call", + "name": "get_weather", + "arguments": {"city": "Paris"}, + }, + ], + } + ] + + session = boto3.Session( + aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" + ) + client = session.client("bedrock-runtime") + mock_http = Mock(status_code=200) + client.meta.events.register( + "before-call.bedrock-runtime.ConverseStream", + lambda **kwargs: (mock_http, converse_stream_response), + ) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + with sentry_sdk.start_transaction(name="parent"): + response = client.converse_stream( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + for _ in response["stream"]: + pass + + 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 sentry_sdk.start_transaction(name="parent"): + response = client.converse_stream( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + for _ in response["stream"]: + pass + + (transaction,) = events + (span,) = transaction["spans"] + + assert ( + json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) + == expected_output_messages + ) From 5253e2954426cce13cf4ed538bc5b711ef12a9ee Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 11 Nov 2026 18:14:02 +0100 Subject: [PATCH 2/4] feat(bedrock): Record gen_ai.response.finish_reasons for streamed responses --- sentry_sdk/integrations/boto3/_services/bedrock.py | 4 ++++ tests/integrations/boto3/test_bedrock.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/sentry_sdk/integrations/boto3/_services/bedrock.py b/sentry_sdk/integrations/boto3/_services/bedrock.py index 6f663a6974..4b24251e80 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock.py @@ -188,6 +188,10 @@ def sentry_patched_iter( yield event continue + if "messageStop" in event: + finish_reasons: "list[str]" = [event["messageStop"]["stopReason"]] + set_on_span(SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, finish_reasons) + if "messageStop" in event and role is not None: for block in content_blocks.values(): if "toolUse" not in block: diff --git a/tests/integrations/boto3/test_bedrock.py b/tests/integrations/boto3/test_bedrock.py index cb7a57ff66..2cf9a3cd71 100644 --- a/tests/integrations/boto3/test_bedrock.py +++ b/tests/integrations/boto3/test_bedrock.py @@ -990,6 +990,9 @@ def test_converse_stream( assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 1024 assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == [ + "end_turn" + ] assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 assert span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 @@ -1029,6 +1032,7 @@ def test_converse_stream( assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 1024 assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + assert span["data"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == ["end_turn"] assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 From 7d0b930cb1ce53230e9888f4de3e910f19a96c56 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 11 Nov 2026 21:38:42 +0100 Subject: [PATCH 3/4] empty commit From 80267a21c95cd353bdb39710ec431137ddd5939f Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 11 Nov 2026 22:45:37 +0100 Subject: [PATCH 4/4] empty commit