diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py index 83f6aa1b09..202282908d 100644 --- a/sentry_sdk/_types.py +++ b/sentry_sdk/_types.py @@ -471,20 +471,44 @@ class DataCollection(TypedDict): class TextPart(TypedDict): type: Literal["text"] - content: str + content: NotRequired[str] class ReasoningPart(TypedDict): type: Literal["reasoning"] - content: str + content: NotRequired[str] class ToolCallPart(TypedDict): type: Literal["tool_call"] name: NotRequired[str] arguments: NotRequired[Any] + class ToolCallResponsePart(TypedDict): + type: Literal["tool_call_response"] + response: NotRequired[Any] + + Modality = Literal["image", "video", "audio", "document"] + + class UriPart(TypedDict): + type: Literal["uri"] + modality: NotRequired[Modality] + uri: NotRequired[str] + + class BlobPart(TypedDict): + type: Literal["blob"] + modality: NotRequired[Modality] + class InputMessage(TypedDict): role: Literal["user", "assistant", "tool"] - parts: list[Union[TextPart, ReasoningPart, ToolCallPart]] + parts: list[ + Union[ + TextPart, + ReasoningPart, + ToolCallPart, + ToolCallResponsePart, + UriPart, + BlobPart, + ] + ] class OutputMessage(TypedDict): role: Literal["assistant"] diff --git a/sentry_sdk/integrations/boto3/_services/bedrock.py b/sentry_sdk/integrations/boto3/_services/bedrock.py index 49b77ee0b9..5c9b16325e 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock.py @@ -30,7 +30,17 @@ ConverseStreamResponseTypeDef, ) - from sentry_sdk._types import Attributes, TextPart, ToolDefinition + from sentry_sdk._types import ( + Attributes, + BlobPart, + InputMessage, + ReasoningPart, + TextPart, + ToolCallPart, + ToolCallResponsePart, + ToolDefinition, + UriPart, + ) from sentry_sdk.integrations.boto3._context import AwsCallContext from sentry_sdk.tracing import Span @@ -160,6 +170,193 @@ def _transform_system_instructions( return instruction_text_parts +def _transform_messages( + input_messages: "Union[list[dict[str, Any]], tuple[dict[str, Any], ...]]", +) -> "list[InputMessage]": + """ + Unlike other providers, AWS declares that input and output messages use the same message type. + For this reason, the same transform is used for user and assistant messages. + """ + messages: "list[InputMessage]" = [] + for message in input_messages: + if not isinstance(message, dict): + continue + + content = message.get("content") + if not isinstance(content, (list, tuple)) or "role" not in message: + continue + + parts: "list[Union[TextPart, ReasoningPart, ToolCallPart, ToolCallResponsePart, UriPart, BlobPart]]" = [] + for block in content: + if not isinstance(block, dict): + continue + + if "text" in block: + parts.append({"type": "text", "content": block["text"]}) + continue + + tool_use = block.get("toolUse") + if tool_use is not None: + tool_call_part: "ToolCallPart" = {"type": "tool_call"} + + if not isinstance(tool_use, dict): + parts.append(tool_call_part) + continue + + if "name" in tool_use: + tool_call_part["name"] = tool_use["name"] + if "input" in tool_use: + tool_call_part["arguments"] = tool_use["input"] + parts.append(tool_call_part) + continue + + tool_result = block.get("toolResult") + if tool_result is not None: + tool_call_response_part: "ToolCallResponsePart" = { + "type": "tool_call_response", + } + + if ( + not isinstance(tool_result, dict) + or "content" not in tool_result + or not isinstance(tool_result["content"], (list, tuple)) + ): + parts.append(tool_call_response_part) + continue + + tool_call_response_part["response"] = [ + item + for item in tool_result["content"] + if isinstance(item, dict) and ("json" in item or "text" in item) + ] + parts.append(tool_call_response_part) + continue + + reasoning_content = block.get("reasoningContent") + if reasoning_content is not None: + reasoning_part: "ReasoningPart" = {"type": "reasoning"} + + if ( + not isinstance(reasoning_content, dict) + or "reasoningText" not in reasoning_content + or not isinstance(reasoning_content["reasoningText"], dict) + or "text" not in reasoning_content["reasoningText"] + ): + parts.append(reasoning_part) + continue + + reasoning_part["content"] = reasoning_content["reasoningText"]["text"] + parts.append(reasoning_part) + continue + + image = block.get("image") + if ( + isinstance(image, dict) + and isinstance(image.get("source"), dict) + and "s3Location" in image["source"] + ): + uri_part: "UriPart" = { + "type": "uri", + "modality": "image", + } + s3_location = image["source"]["s3Location"] + if not isinstance(s3_location, dict) or "uri" not in s3_location: + parts.append(uri_part) + continue + + uri_part["uri"] = s3_location["uri"] + parts.append(uri_part) + continue + elif ( + isinstance(image, dict) + and isinstance(image.get("source"), dict) + and "bytes" in image["source"] + ): + parts.append({"type": "blob", "modality": "image"}) + continue + + video = block.get("video") + if ( + isinstance(video, dict) + and isinstance(video.get("source"), dict) + and "s3Location" in video["source"] + ): + uri_part = { + "type": "uri", + "modality": "video", + } + s3_location = video["source"]["s3Location"] + if not isinstance(s3_location, dict) or "uri" not in s3_location: + parts.append(uri_part) + continue + + uri_part["uri"] = s3_location["uri"] + parts.append(uri_part) + continue + elif ( + isinstance(video, dict) + and isinstance(video.get("source"), dict) + and "bytes" in video["source"] + ): + parts.append({"type": "blob", "modality": "video"}) + continue + + audio = block.get("audio") + if ( + isinstance(audio, dict) + and isinstance(audio.get("source"), dict) + and "s3Location" in audio["source"] + ): + uri_part = { + "type": "uri", + "modality": "audio", + } + s3_location = audio["source"]["s3Location"] + if not isinstance(s3_location, dict) or "uri" not in s3_location: + parts.append(uri_part) + continue + + uri_part["uri"] = s3_location["uri"] + parts.append(uri_part) + continue + elif ( + isinstance(audio, dict) + and isinstance(audio.get("source"), dict) + and "bytes" in audio["source"] + ): + parts.append({"type": "blob", "modality": "audio"}) + continue + + document = block.get("document") + if ( + not isinstance(document, dict) + or "source" not in document + or not isinstance(document["source"], dict) + ): + continue + + if "s3Location" in document["source"]: + uri_part = { + "type": "uri", + "modality": "document", + } + s3_location = document["source"]["s3Location"] + if not isinstance(s3_location, dict) or "uri" not in s3_location: + parts.append(uri_part) + continue + + uri_part["uri"] = s3_location["uri"] + parts.append(uri_part) + continue + + if "bytes" in document["source"]: + parts.append({"type": "blob", "modality": "document"}) + continue + + messages.append({"role": message["role"], "parts": parts}) + return messages + + def _transform_tool_definitions( tools: "list[dict[str, Any]]", ) -> "list[ToolDefinition]": @@ -258,6 +455,14 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": _transform_system_instructions(system_instructions) ) + input_messages = ctx.params.get("messages") + if ( + has_data_collection_enabled(client_options) or should_send_default_pii() + ) and isinstance(input_messages, (list, tuple)): + attributes[SPANDATA.GEN_AI_INPUT_MESSAGES] = json.dumps( + _transform_messages(input_messages) + ) + tool_config = ctx.params.get("toolConfig") if not isinstance(tool_config, dict): return attributes diff --git a/tests/integrations/boto3/test_bedrock.py b/tests/integrations/boto3/test_bedrock.py index 6cf2b89673..9461bce69c 100644 --- a/tests/integrations/boto3/test_bedrock.py +++ b/tests/integrations/boto3/test_bedrock.py @@ -153,6 +153,272 @@ def converse_stream_error_response(): } +@pytest.fixture +def content_block_messages(): + """ + A list of messages exercising all types of content blocks from the Bedrock API + which the SDK transforms, and the corresponding transformed messages. + """ + messages = [ + {"role": "user", "content": [{"text": "What is the weather in Paris?"}]}, + { + "role": "user", + "content": [{"image": {"format": "png", "source": {"bytes": b"image"}}}], + }, + { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_1", + "name": "get_weather", + "input": {"location": "Paris"}, + } + } + ], + }, + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "tooluse_1", + "content": [ + {"json": {"temperature": 18, "unit": "celsius"}}, + {"text": "Partly cloudy"}, + ], + "status": "success", + } + } + ], + }, + ] + + expected_input_messages = [ + { + "role": "user", + "parts": [{"type": "text", "content": "What is the weather in Paris?"}], + }, + {"role": "user", "parts": [{"type": "blob", "modality": "image"}]}, + { + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "name": "get_weather", + "arguments": {"location": "Paris"}, + } + ], + }, + { + "role": "user", + "parts": [ + { + "type": "tool_call_response", + "response": [ + {"json": {"temperature": 18, "unit": "celsius"}}, + {"text": "Partly cloudy"}, + ], + } + ], + }, + ] + + if BOTOCORE_VERSION >= Version("1.34.131"): + messages.append( + { + "role": "user", + "content": [ + { + "document": { + "format": "pdf", + "name": "report", + "source": {"bytes": b"document"}, + } + } + ], + } + ) + expected_input_messages.append( + {"role": "user", "parts": [{"type": "blob", "modality": "document"}]} + ) + + if BOTOCORE_VERSION >= Version("1.35.74"): + messages.extend( + [ + { + "role": "user", + "content": [ + {"video": {"format": "mp4", "source": {"bytes": b"video"}}} + ], + }, + { + "role": "user", + "content": [ + { + "video": { + "format": "mp4", + "source": { + "s3Location": { + "uri": "s3://my-bucket/videos/clip.mp4" + } + }, + } + } + ], + }, + ] + ) + expected_input_messages.extend( + [ + {"role": "user", "parts": [{"type": "blob", "modality": "video"}]}, + { + "role": "user", + "parts": [ + { + "type": "uri", + "modality": "video", + "uri": "s3://my-bucket/videos/clip.mp4", + } + ], + }, + ] + ) + + if BOTOCORE_VERSION >= Version("1.37.0"): + messages.append( + { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": "The user wants the weather in Paris.", + "signature": "signature", + } + } + } + ], + } + ) + expected_input_messages.append( + { + "role": "assistant", + "parts": [ + { + "type": "reasoning", + "content": "The user wants the weather in Paris.", + } + ], + } + ) + + if BOTOCORE_VERSION >= Version("1.38.3"): + messages.extend( + [ + { + "role": "user", + "content": [ + { + "image": { + "format": "png", + "source": { + "s3Location": { + "uri": "s3://my-bucket/images/photo.png" + } + }, + } + } + ], + }, + { + "role": "user", + "content": [ + { + "document": { + "format": "pdf", + "name": "manual", + "source": { + "s3Location": { + "uri": "s3://my-bucket/documents/manual.pdf" + } + }, + } + } + ], + }, + ] + ) + expected_input_messages.extend( + [ + { + "role": "user", + "parts": [ + { + "type": "uri", + "modality": "image", + "uri": "s3://my-bucket/images/photo.png", + } + ], + }, + { + "role": "user", + "parts": [ + { + "type": "uri", + "modality": "document", + "uri": "s3://my-bucket/documents/manual.pdf", + } + ], + }, + ] + ) + + if BOTOCORE_VERSION >= Version("1.42.1"): + messages.extend( + [ + { + "role": "user", + "content": [ + {"audio": {"format": "mp3", "source": {"bytes": b"audio"}}} + ], + }, + { + "role": "user", + "content": [ + { + "audio": { + "format": "mp3", + "source": { + "s3Location": { + "uri": "s3://my-bucket/audio/question.mp3" + } + }, + } + } + ], + }, + ] + ) + expected_input_messages.extend( + [ + {"role": "user", "parts": [{"type": "blob", "modality": "audio"}]}, + { + "role": "user", + "parts": [ + { + "type": "uri", + "modality": "audio", + "uri": "s3://my-bucket/audio/question.mp3", + } + ], + }, + ] + ) + + return messages, expected_input_messages + + @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_converse( @@ -485,6 +751,7 @@ def test_converse_input_attributes( span_streaming, stream_gen_ai_spans, converse_response, + content_block_messages, ): if data_collection: sentry_init( @@ -503,6 +770,8 @@ def test_converse_input_attributes( trace_lifecycle="stream" if span_streaming else "static", ) + messages, expected_input_messages = content_block_messages + session = boto3.Session( aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" ) @@ -516,7 +785,7 @@ def test_converse_input_attributes( with stubber, sentry_sdk.start_transaction(name="parent"): client.converse( modelId="anthropic.claude-3-sonnet-20240229-v1:0", - messages=[{"role": "user", "content": [{"text": "Hi"}]}], + messages=messages, system=[ {"text": "You are a helpful assistant."}, {"text": "Be concise and clear."}, @@ -535,13 +804,17 @@ def test_converse_input_attributes( {"type": "text", "content": "You are a helpful assistant."}, {"type": "text", "content": "Be concise and clear."}, ] + assert ( + json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_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"}]}], + messages=messages, system=[ {"text": "You are a helpful assistant."}, {"text": "Be concise and clear."}, @@ -555,6 +828,10 @@ def test_converse_input_attributes( {"type": "text", "content": "You are a helpful assistant."}, {"type": "text", "content": "Be concise and clear."}, ] + assert ( + json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_messages + ) @pytest.mark.parametrize("span_streaming", [True, False]) @@ -861,6 +1138,7 @@ def test_converse_stream_input_attributes( span_streaming, stream_gen_ai_spans, converse_stream_response, + content_block_messages, ): if data_collection: sentry_init( @@ -879,6 +1157,8 @@ def test_converse_stream_input_attributes( trace_lifecycle="stream" if span_streaming else "static", ) + messages, expected_input_messages = content_block_messages + session = boto3.Session( aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" ) @@ -895,7 +1175,7 @@ def test_converse_stream_input_attributes( 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"}]}], + messages=messages, system=[ {"text": "You are a helpful assistant."}, {"text": "Be concise and clear."}, @@ -916,13 +1196,17 @@ def test_converse_stream_input_attributes( {"type": "text", "content": "You are a helpful assistant."}, {"type": "text", "content": "Be concise and clear."}, ] + assert ( + json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_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"}]}], + messages=messages, system=[ {"text": "You are a helpful assistant."}, {"text": "Be concise and clear."}, @@ -938,3 +1222,7 @@ def test_converse_stream_input_attributes( {"type": "text", "content": "You are a helpful assistant."}, {"type": "text", "content": "Be concise and clear."}, ] + assert ( + json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_messages + )