diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py index 83f6aa1b09..28b9a6562a 100644 --- a/sentry_sdk/_types.py +++ b/sentry_sdk/_types.py @@ -479,16 +479,22 @@ class ReasoningPart(TypedDict): class ToolCallPart(TypedDict): type: Literal["tool_call"] + id: NotRequired[str] name: NotRequired[str] arguments: NotRequired[Any] + class ToolCallResponsePart(TypedDict): + type: Literal["tool_call_response"] + id: NotRequired[str] + result: NotRequired[Any] + class InputMessage(TypedDict): role: Literal["user", "assistant", "tool"] - parts: list[Union[TextPart, ReasoningPart, ToolCallPart]] + parts: list[Union[TextPart, ReasoningPart, ToolCallPart, ToolCallResponsePart]] class OutputMessage(TypedDict): role: Literal["assistant"] - parts: list[Union[TextPart, ReasoningPart]] + parts: list[Union[TextPart, ReasoningPart, ToolCallPart]] class ToolDefinition(TypedDict): type: str diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index be2660433e..db212a83c3 100644 --- a/sentry_sdk/consts.py +++ b/sentry_sdk/consts.py @@ -733,6 +733,12 @@ class SPANDATA: Example: "1234567890" """ + GEN_AI_REQUEST_STOP_SEQUENCES = "gen_ai.request.stop_sequences" + """ + List of sequences that the model will use to stop generating further tokens. + Example: ["forest", "lived"] + """ + GEN_AI_REQUEST_TEMPERATURE = "gen_ai.request.temperature" """ The temperature parameter used to control randomness in the output. diff --git a/sentry_sdk/integrations/boto3/__init__.py b/sentry_sdk/integrations/boto3/__init__.py index d36b44d271..3688211be9 100644 --- a/sentry_sdk/integrations/boto3/__init__.py +++ b/sentry_sdk/integrations/boto3/__init__.py @@ -13,6 +13,11 @@ class Boto3Integration(Integration): identifier = IDENTIFIER origin = ORIGIN + def __init__(self: "Boto3Integration", include_prompts: bool = True) -> None: + # whether to record gen_ai prompts and responses, e.g. of Bedrock Runtime + # calls, when `send_default_pii` is on and `data_collection` isn't set. + self.include_prompts = include_prompts + @staticmethod def setup_once() -> None: version = parse_version(BOTOCORE_VERSION) diff --git a/sentry_sdk/integrations/boto3/_instrumentation.py b/sentry_sdk/integrations/boto3/_instrumentation.py index 3c180a025c..ac70b01464 100644 --- a/sentry_sdk/integrations/boto3/_instrumentation.py +++ b/sentry_sdk/integrations/boto3/_instrumentation.py @@ -176,12 +176,18 @@ def _start_client_span( with capture_internal_exceptions(): attributes.update(_get_client_attributes(ctx)) - # `sentry.span_op` and `sentry.span_origin` are set to generic defaults; - # a service extension can override them with `get_span_op()` and `get_span_origin()`. + # the span name, `sentry.span_op`, and `sentry.span_origin` are set to generic defaults; + # a service extension can override them with `get_span_name()`, `get_span_op()`, and + # `get_span_origin()`. span_op = OP.HTTP_CLIENT span_origin = ORIGIN if service_ext is not None: + with capture_internal_exceptions(): + service_span_name = service_ext.get_span_name(ctx) + if service_span_name is not None: + span_name = service_span_name + with capture_internal_exceptions(): service_op = service_ext.get_span_op(ctx) if service_op is not None: @@ -207,13 +213,19 @@ def _start_client_span( SPANDATA.SENTRY_ORIGIN: span_origin, } ) - return sentry_sdk.traces.start_span( + streamed_span = sentry_sdk.traces.start_span( name=span_name, attributes=attributes, # `StreamingBody` responses outlive `_make_api_call()`. `_activate_client_span()` # activates this span only while the call itself runs. active=False, ) + if service_ext is not None: + with capture_internal_exceptions(): + streamed_span.set_attributes( + service_ext.get_span_start_attributes(ctx, streamed_span) + ) + return streamed_span span = sentry_sdk.start_span( name=span_name, @@ -222,6 +234,9 @@ def _start_client_span( ) with capture_internal_exceptions(): _set_span_attributes(span, attributes) + if service_ext is not None: + with capture_internal_exceptions(): + _set_span_attributes(span, service_ext.get_span_start_attributes(ctx, span)) with capture_internal_exceptions(): if ctx.service_id_hyphenized is not None: span.set_tag("aws.service_id", ctx.service_id_hyphenized) diff --git a/sentry_sdk/integrations/boto3/_services/base.py b/sentry_sdk/integrations/boto3/_services/base.py index 0d9cb1c24d..be3f8add7b 100644 --- a/sentry_sdk/integrations/boto3/_services/base.py +++ b/sentry_sdk/integrations/boto3/_services/base.py @@ -1,10 +1,12 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Optional + from typing import Any, Optional, Union from sentry_sdk._types import Attributes from sentry_sdk.integrations.boto3._context import AwsCallContext + from sentry_sdk.traces import StreamedSpan + from sentry_sdk.tracing import Span class _ServiceExtension: @@ -17,6 +19,10 @@ class _ServiceExtension: __slots__ = () + def get_span_name(self, ctx: "AwsCallContext") -> "Optional[str]": + """Return an optional span name override, or `None` to keep the default.""" + return None + def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]": """Return an optional `sentry.op` override, or `None` to keep the default.""" return None @@ -29,6 +35,16 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": """Return request attributes to add before the AWS request is made.""" return {} + def get_span_start_attributes( + self, ctx: "AwsCallContext", span: "Union[Span, StreamedSpan]" + ) -> "Attributes": + """ + Return request attributes that need the started span, e.g. truncated + `gen_ai` messages. Unlike `get_request_attributes()`, they're added after + the span starts, so `ignore_spans` doesn't see them. + """ + return {} + def get_response_attributes( self, ctx: "AwsCallContext", response: "Any" ) -> "Attributes": diff --git a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py new file mode 100644 index 0000000000..516971a670 --- /dev/null +++ b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py @@ -0,0 +1,431 @@ +import json +from collections.abc import Mapping +from typing import TYPE_CHECKING, cast + +import sentry_sdk +from sentry_sdk.ai.utils import ( + MAX_SINGLE_MESSAGE_CONTENT_CHARS, + truncate_and_annotate_messages, +) +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.boto3._services.base import _ServiceExtension +from sentry_sdk.integrations.boto3.consts import AI_ORIGIN, IDENTIFIER +from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.traces import NoOpStreamedSpan +from sentry_sdk.tracing_utils import should_truncate_gen_ai_input +from sentry_sdk.utils import capture_internal_exceptions, has_data_collection_enabled + +if TYPE_CHECKING: + from typing import Any, Literal, Optional, Union + + from sentry_sdk._types import ( + Attributes, + InputMessage, + OutputMessage, + ReasoningPart, + TextPart, + ToolCallPart, + ToolCallResponsePart, + ToolDefinition, + ) + from sentry_sdk.integrations.boto3 import Boto3Integration + from sentry_sdk.integrations.boto3._context import AwsCallContext + from sentry_sdk.traces import StreamedSpan + from sentry_sdk.tracing import Span + + ContentPart = Union[TextPart, ReasoningPart, ToolCallPart, ToolCallResponsePart] + + +# `gen_ai.provider.name` of Amazon Bedrock. +# https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#gen-ai-provider-name +PROVIDER_NAME = "aws.bedrock" + +# `gen_ai.operation.name` and span op of each instrumented operation. +GEN_AI_OPERATIONS = { + "Converse": ("chat", OP.GEN_AI_CHAT), +} + + +class _BedrockRuntimeExtension(_ServiceExtension): + """ + Turns Bedrock Runtime model calls into `gen_ai` client spans. + https://develop.sentry.dev/sdk/telemetry/traces/modules/ai-agents/#ai-client-span + """ + + __slots__ = () + + def get_span_name(self, ctx: "AwsCallContext") -> "Optional[str]": + if ctx.operation_name not in GEN_AI_OPERATIONS: + return None + + operation_name, _ = GEN_AI_OPERATIONS[ctx.operation_name] + model_id = ctx.params.get("modelId") + return f"{operation_name} {model_id}" if model_id else operation_name + + def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]": + if ctx.operation_name not in GEN_AI_OPERATIONS: + return None + + _, span_op = GEN_AI_OPERATIONS[ctx.operation_name] + return span_op + + def get_span_origin(self, ctx: "AwsCallContext") -> "Optional[str]": + if ctx.operation_name not in GEN_AI_OPERATIONS: + return None + + return AI_ORIGIN + + def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": + if ctx.operation_name not in GEN_AI_OPERATIONS: + return {} + + operation_name, _ = GEN_AI_OPERATIONS[ctx.operation_name] + attributes: "Attributes" = { + SPANDATA.GEN_AI_PROVIDER_NAME: PROVIDER_NAME, + SPANDATA.GEN_AI_OPERATION_NAME: operation_name, + SPANDATA.GEN_AI_RESPONSE_STREAMING: False, + } + + model_id = ctx.params.get("modelId") + if model_id: + attributes[SPANDATA.GEN_AI_REQUEST_MODEL] = model_id + + with capture_internal_exceptions(): + attributes.update(_get_converse_request_attributes(ctx.params)) + + return attributes + + def get_span_start_attributes( + self, ctx: "AwsCallContext", span: "Union[Span, StreamedSpan]" + ) -> "Attributes": + if ( + ctx.operation_name not in GEN_AI_OPERATIONS + # e.g. a span dropped by `ignore_spans`. + or isinstance(span, NoOpStreamedSpan) + or not _should_record_gen_ai("inputs") + ): + return {} + + return _get_input_messages_attributes( + span, _transform_input_messages(ctx.params.get("messages")) + ) + + def get_response_attributes( + self, ctx: "AwsCallContext", response: "Any" + ) -> "Attributes": + if ctx.operation_name not in GEN_AI_OPERATIONS: + return {} + + return _get_converse_response_attributes(response) + + +def _should_record_gen_ai(kind: 'Literal["inputs", "outputs"]') -> bool: + """ + Record prompts (`inputs`) or responses (`outputs`) as `data_collection` allows, + or, if it isn't set, with `send_default_pii` and `include_prompts`, like the + other gen_ai integrations. + """ + client = sentry_sdk.get_client() + if has_data_collection_enabled(client.options): + return bool(client.options["data_collection"]["gen_ai"][kind]) + + integration = cast("Optional[Boto3Integration]", client.get_integration(IDENTIFIER)) + return ( + should_send_default_pii() + and integration is not None + and integration.include_prompts + ) + + +def _get_converse_request_attributes(params: "Mapping[str, Any]") -> "Attributes": + """ + The input messages are added by `get_span_start_attributes()`. + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html#API_runtime_Converse_RequestSyntax + """ + attributes: "Attributes" = {} + + inference_config = params.get("inferenceConfig") + if isinstance(inference_config, Mapping): + if inference_config.get("maxTokens") is not None: + attributes[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] = inference_config[ + "maxTokens" + ] + if inference_config.get("temperature") is not None: + attributes[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] = inference_config[ + "temperature" + ] + if inference_config.get("topP") is not None: + attributes[SPANDATA.GEN_AI_REQUEST_TOP_P] = inference_config["topP"] + if isinstance(inference_config.get("stopSequences"), (list, tuple)): + attributes[SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES] = list( + inference_config["stopSequences"] + ) + + if not _should_record_gen_ai("inputs"): + return attributes + + with capture_internal_exceptions(): + system_instructions = _transform_system_instructions(params.get("system")) + if system_instructions: + attributes[SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS] = json.dumps( + system_instructions + ) + + with capture_internal_exceptions(): + tool_config = params.get("toolConfig") + if isinstance(tool_config, Mapping): + tool_definitions = _transform_tool_definitions(tool_config.get("tools")) + if tool_definitions: + attributes[SPANDATA.GEN_AI_TOOL_DEFINITIONS] = json.dumps( + tool_definitions + ) + + return attributes + + +def _get_input_messages_attributes( + span: "Union[Span, StreamedSpan]", input_messages: "list[InputMessage]" +) -> "Attributes": + """ + Like the other `gen_ai` integrations, only the last message is kept when + `gen_ai` spans are sent in transactions, and its text, tool call arguments, + and tool results share one truncation limit, so that long histories, prompts, + and tool results don't exceed the transaction size limit. JSON arguments and + results count with their serialized length, and are replaced by their + truncated serialization when they don't fit. + """ + messages: "Optional[list[dict[str, Any]]]" = cast( + "list[dict[str, Any]]", input_messages + ) + if should_truncate_gen_ai_input(sentry_sdk.get_client().options): + messages = truncate_and_annotate_messages( + messages, span, sentry_sdk.get_current_scope() + ) + if messages: + # `truncate_and_annotate_messages()` only truncates a message's + # `content`, but this format keeps the text, tool call arguments, and + # tool results in each part. They share the message's limit. + remaining = MAX_SINGLE_MESSAGE_CONTENT_CHARS + + def truncate(value: "Any") -> "Any": + nonlocal remaining + text = value if isinstance(value, str) else json.dumps(value) + if len(text) > remaining: + value = text[:remaining] + "..." + remaining = max(remaining - len(text), 0) + return value + + # the last message is a copy, so its parts can be replaced. + for part in messages[-1].get("parts") or []: + for key in ("content", "arguments"): + if key in part: + part[key] = truncate(part[key]) + result = part.get("result") + if isinstance(result, list): + part["result"] = [truncate(item) for item in result] + elif result is not None: + part["result"] = truncate(result) + + if not messages: + return {} + + return {SPANDATA.GEN_AI_INPUT_MESSAGES: json.dumps(messages)} + + +def _get_converse_response_attributes(response: "Any") -> "Attributes": + """ + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html#API_runtime_Converse_ResponseSyntax + """ + attributes: "Attributes" = {} + if not isinstance(response, Mapping): + return attributes + + stop_reason = response.get("stopReason") + if isinstance(stop_reason, str): + attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] = [stop_reason] + + usage = response.get("usage") + if isinstance(usage, Mapping): + attributes.update(_get_usage_attributes(usage)) + + if not _should_record_gen_ai("outputs"): + return attributes + + with capture_internal_exceptions(): + output = response.get("output") + message = output.get("message") if isinstance(output, Mapping) else None + if isinstance(message, Mapping): + output_message: "OutputMessage" = {"role": "assistant", "parts": []} + for part in _transform_content_blocks(message.get("content")): + # models don't respond with tool results. + if part["type"] != "tool_call_response": + output_message["parts"].append(part) + attributes[SPANDATA.GEN_AI_OUTPUT_MESSAGES] = json.dumps([output_message]) + + return attributes + + +def _get_usage_attributes(usage: "Mapping[str, Any]") -> "Attributes": + """ + Bedrock's `inputTokens` excludes the tokens read from or written to the prompt + cache, but `gen_ai.usage.input_tokens` includes them. + https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html + """ + attributes: "Attributes" = {} + + input_tokens = usage.get("inputTokens") + output_tokens = usage.get("outputTokens") + total_tokens = usage.get("totalTokens") + cache_read_input_tokens = usage.get("cacheReadInputTokens") + cache_write_input_tokens = usage.get("cacheWriteInputTokens") + + if cache_read_input_tokens is not None: + attributes[SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] = ( + cache_read_input_tokens + ) + if cache_write_input_tokens is not None: + attributes[SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS] = ( + cache_write_input_tokens + ) + if input_tokens is not None: + attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] = ( + input_tokens + + (cache_read_input_tokens or 0) + + (cache_write_input_tokens or 0) + ) + if output_tokens is not None: + attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] = output_tokens + if total_tokens is not None: + attributes[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] = total_tokens + + return attributes + + +def _transform_system_instructions(system: "Any") -> "list[TextPart]": + if not isinstance(system, (list, tuple)): + return [] + + return [ + {"type": "text", "content": block["text"]} + for block in system + if isinstance(block, Mapping) and isinstance(block.get("text"), str) + ] + + +def _transform_input_messages(messages: "Any") -> "list[InputMessage]": + """ + Bedrock sends tool results in `user` messages; they're recorded as `tool` + messages, which is how `gen_ai.input.messages` represents them. + """ + if not isinstance(messages, (list, tuple)): + return [] + + input_messages: "list[InputMessage]" = [] + for message in messages: + if not isinstance(message, Mapping): + continue + + role = message.get("role") + if role != "user" and role != "assistant": + continue + + parts: "list[ContentPart]" = [] + tool_parts: "list[ContentPart]" = [] + for part in _transform_content_blocks(message.get("content")): + if part["type"] == "tool_call_response": + tool_parts.append(part) + else: + parts.append(part) + + if tool_parts: + input_messages.append({"role": "tool", "parts": tool_parts}) + if parts: + input_messages.append({"role": role, "parts": parts}) + + return input_messages + + +def _transform_content_blocks(content: "Any") -> "list[ContentPart]": + """ + Transform text, reasoning, tool use, and tool result content blocks. Media + content blocks, e.g. images and documents, are not recorded. + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlock.html + """ + if not isinstance(content, (list, tuple)): + return [] + + parts: "list[ContentPart]" = [] + for block in content: + if not isinstance(block, Mapping): + continue + + if isinstance(block.get("text"), str): + parts.append({"type": "text", "content": block["text"]}) + continue + + reasoning = block.get("reasoningContent") + if isinstance(reasoning, Mapping): + reasoning_text = reasoning.get("reasoningText") + if isinstance(reasoning_text, Mapping) and isinstance( + reasoning_text.get("text"), str + ): + parts.append({"type": "reasoning", "content": reasoning_text["text"]}) + continue + + tool_use = block.get("toolUse") + if isinstance(tool_use, Mapping): + tool_call: "ToolCallPart" = {"type": "tool_call"} + if "toolUseId" in tool_use: + tool_call["id"] = tool_use["toolUseId"] + if "name" in tool_use: + tool_call["name"] = tool_use["name"] + if "input" in tool_use: + tool_call["arguments"] = tool_use["input"] + parts.append(tool_call) + continue + + tool_result = block.get("toolResult") + if isinstance(tool_result, Mapping): + tool_response: "ToolCallResponsePart" = {"type": "tool_call_response"} + if "toolUseId" in tool_result: + tool_response["id"] = tool_result["toolUseId"] + + result = [ + item["text"] if "text" in item else item["json"] + for item in tool_result.get("content") or () + if isinstance(item, Mapping) and ("text" in item or "json" in item) + ] + if result: + tool_response["result"] = result[0] if len(result) == 1 else result + parts.append(tool_response) + + return parts + + +def _transform_tool_definitions(tools: "Any") -> "list[ToolDefinition]": + """ + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolSpecification.html + """ + if not isinstance(tools, (list, tuple)): + return [] + + tool_definitions: "list[ToolDefinition]" = [] + for tool in tools: + tool_spec = tool.get("toolSpec") if isinstance(tool, Mapping) else None + if not isinstance(tool_spec, Mapping): + continue + + tool_definition: "ToolDefinition" = {"type": "function"} + if "name" in tool_spec: + tool_definition["name"] = tool_spec["name"] + if "description" in tool_spec: + tool_definition["description"] = tool_spec["description"] + + input_schema = tool_spec.get("inputSchema") + if isinstance(input_schema, Mapping) and isinstance( + input_schema.get("json"), Mapping + ): + tool_definition["parameters"] = dict(input_schema["json"]) + tool_definitions.append(tool_definition) + + return tool_definitions diff --git a/sentry_sdk/integrations/boto3/_services/registry.py b/sentry_sdk/integrations/boto3/_services/registry.py index b387db0f26..379c32496f 100644 --- a/sentry_sdk/integrations/boto3/_services/registry.py +++ b/sentry_sdk/integrations/boto3/_services/registry.py @@ -1,5 +1,9 @@ from typing import TYPE_CHECKING +from sentry_sdk.integrations.boto3._services.bedrock_runtime import ( + _BedrockRuntimeExtension, +) + if TYPE_CHECKING: from typing import Dict, Optional @@ -10,7 +14,9 @@ # _SERVICE_EXTENSIONS = {"s3": _S3Extension()} # when py 3.15 drops, we might want to take a look at using # a lazy-loading approach using the new `lazy` keyword. -_SERVICE_EXTENSIONS: "Dict[str, _ServiceExtension]" = {} +_SERVICE_EXTENSIONS: "Dict[str, _ServiceExtension]" = { + "bedrock-runtime": _BedrockRuntimeExtension(), +} def _resolve_service( diff --git a/sentry_sdk/integrations/boto3/consts.py b/sentry_sdk/integrations/boto3/consts.py index 7d0d441257..2bc79b5d32 100644 --- a/sentry_sdk/integrations/boto3/consts.py +++ b/sentry_sdk/integrations/boto3/consts.py @@ -1,5 +1,7 @@ IDENTIFIER = "boto3" ORIGIN = f"auto.http.{IDENTIFIER}" +# origin of `gen_ai` spans, e.g. for Bedrock Runtime calls. +AI_ORIGIN = f"auto.ai.{IDENTIFIER}" # value is used by `rpc.system` (deprecated in OTel, but we still support it for now) and `rpc.system.name` # https://opentelemetry.io/docs/specs/semconv/cloud-providers/aws-sdk/#aws-sdk-spans diff --git a/tests/integrations/boto3/test_bedrock_runtime.py b/tests/integrations/boto3/test_bedrock_runtime.py new file mode 100644 index 0000000000..257d350f08 --- /dev/null +++ b/tests/integrations/boto3/test_bedrock_runtime.py @@ -0,0 +1,744 @@ +import copy +import json + +import boto3 +import pytest +from botocore.exceptions import ClientError +from botocore.stub import Stubber + +import sentry_sdk +from sentry_sdk.ai.utils import MAX_SINGLE_MESSAGE_CONTENT_CHARS +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.boto3 import Boto3Integration +from sentry_sdk.integrations.boto3.consts import AI_ORIGIN, ORIGIN + +session = boto3.Session( # type: ignore[attr-defined] + aws_access_key_id="-", + aws_secret_access_key="-", + region_name="us-east-1", +) + +if "bedrock-runtime" not in session.get_available_services() or not hasattr( + session.client("bedrock-runtime"), "converse" +): + pytest.skip( + "botocore doesn't support the Bedrock Runtime Converse API", + allow_module_level=True, + ) + + +def _has_shape_member(shape_name, member_name): + service_model = session.client("bedrock-runtime").meta.service_model + return member_name in service_model.shape_for(shape_name).members + + +MODEL_ID = "anthropic.claude-3-5-sonnet-20240620-v1:0" +# as long as the last input message's text can be without being truncated. +LONG_TEXT = "a" * MAX_SINGLE_MESSAGE_CONTENT_CHARS +HUGE_TEXT = "a" * (MAX_SINGLE_MESSAGE_CONTENT_CHARS * 30) +TRUNCATED_TEXT = LONG_TEXT + "..." + +CONVERSE_RESPONSE = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": "Hello! How can I help?"}], + } + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 12, "outputTokens": 8, "totalTokens": 20}, + "metrics": {"latencyMs": 100}, +} + +TRACE_MODES = [ + pytest.param({"trace_lifecycle": "stream"}, id="span_streaming"), + pytest.param({"trace_lifecycle": "static"}, id="static"), + pytest.param( + {"trace_lifecycle": "static", "stream_gen_ai_spans": False}, + id="static_without_gen_ai_span_streaming", + ), +] + +# `sentry_init()` options, `include_prompts`, and whether prompts (inputs) and +# responses (outputs) are recorded. +GEN_AI_DATA_COLLECTION_CASES = [ + pytest.param({}, True, False, False, id="default"), + pytest.param({"send_default_pii": True}, True, True, True, id="send_default_pii"), + pytest.param( + {"send_default_pii": True}, + False, + False, + False, + id="send_default_pii_without_include_prompts", + ), + pytest.param({"data_collection": {}}, True, True, True, id="data_collection"), + pytest.param( + {"data_collection": {"gen_ai": {"inputs": True, "outputs": False}}}, + True, + True, + False, + id="data_collection_inputs_only", + ), + pytest.param( + {"data_collection": {"gen_ai": {"inputs": False, "outputs": True}}}, + True, + False, + True, + id="data_collection_outputs_only", + ), + pytest.param( + { + "send_default_pii": True, + "data_collection": {"gen_ai": {"inputs": False, "outputs": False}}, + }, + True, + False, + False, + id="data_collection_wins_over_send_default_pii", + ), + pytest.param( + {"data_collection": {"gen_ai": {"inputs": True, "outputs": True}}}, + False, + True, + True, + id="data_collection_ignores_include_prompts", + ), +] + + +@pytest.fixture +def init_bedrock(sentry_init): + def inner(trace_mode, include_prompts=True, **kwargs): + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration(include_prompts=include_prompts)], + server_name="", + **trace_mode, + **kwargs, + ) + return session.client("bedrock-runtime") + + return inner + + +def _capture_spans(capture_items, trace_mode, call): + """ + Return the finished spans, normalized to the span streaming format, i.e. + with the span's data and `sentry.op` and `sentry.origin` in `attributes`. + """ + items = capture_items("transaction", "span") + + if trace_mode["trace_lifecycle"] == "stream": + with sentry_sdk.traces.start_span(name="parent"): # type: ignore[attr-defined] + call() + else: + with sentry_sdk.start_transaction(name="parent"): + call() + sentry_sdk.flush() + + spans = [] + for item in items: + if item.type == "span": + spans.append(item.payload) + continue + + for span in item.payload["spans"]: + assert span["timestamp"] is not None + spans.append( + { + "name": span["description"], + # like the conversion of `gen_ai` spans to the span streaming format. + "status": "ok" if span.get("status") in (None, "ok") else "error", + "attributes": { + SPANDATA.SENTRY_OP: span["op"], + SPANDATA.SENTRY_ORIGIN: span["origin"], + **span["data"], + }, + } + ) + + return spans + + +def _expected_input_messages(trace_mode, input_messages): + """ + Like the other `gen_ai` integrations, only the last input message is kept when + `gen_ai` spans are sent in transactions. + """ + if trace_mode.get("stream_gen_ai_spans") is False: + return input_messages[-1:] + return input_messages + + +def _get_bedrock_span(spans, origin=AI_ORIGIN): + (span,) = ( + span + for span in spans + if span["attributes"].get(SPANDATA.SENTRY_ORIGIN) == origin + ) + return span + + +# The last Bedrock message, how it's recorded, and how it's recorded when it's +# truncated. Its text, tool call arguments, and tool results share one limit, and +# JSON that doesn't fit is replaced by its truncated serialization. +LAST_MESSAGE_CASES = [ + pytest.param( + {"role": "user", "content": [{"text": "Hello!"}, {"text": LONG_TEXT}]}, + { + "role": "user", + "parts": [ + {"type": "text", "content": "Hello!"}, + {"type": "text", "content": LONG_TEXT}, + ], + }, + { + "role": "user", + "parts": [ + {"type": "text", "content": "Hello!"}, + { + "type": "text", + "content": LONG_TEXT[: MAX_SINGLE_MESSAGE_CONTENT_CHARS - 6] + + "...", + }, + ], + }, + id="text", + ), + pytest.param( + {"role": "user", "content": [{"text": LONG_TEXT}]}, + {"role": "user", "parts": [{"type": "text", "content": LONG_TEXT}]}, + {"role": "user", "parts": [{"type": "text", "content": LONG_TEXT}]}, + id="text_at_limit", + ), + pytest.param( + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "tool-1", + "content": [{"text": HUGE_TEXT}], + } + } + ], + }, + { + "role": "tool", + "parts": [ + {"type": "tool_call_response", "id": "tool-1", "result": HUGE_TEXT} + ], + }, + { + "role": "tool", + "parts": [ + {"type": "tool_call_response", "id": "tool-1", "result": TRUNCATED_TEXT} + ], + }, + id="tool_result", + ), + pytest.param( + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "tool-1", + "content": [ + {"json": {"weather": "rainy"}}, + {"text": HUGE_TEXT}, + {"text": "Done."}, + ], + } + } + ], + }, + { + "role": "tool", + "parts": [ + { + "type": "tool_call_response", + "id": "tool-1", + "result": [{"weather": "rainy"}, HUGE_TEXT, "Done."], + } + ], + }, + { + "role": "tool", + "parts": [ + { + "type": "tool_call_response", + "id": "tool-1", + # the JSON item counts with its serialized length. + "result": [ + {"weather": "rainy"}, + HUGE_TEXT[ + : MAX_SINGLE_MESSAGE_CONTENT_CHARS + - len(json.dumps({"weather": "rainy"})) + ] + + "...", + "...", + ], + } + ], + }, + id="tool_result_list", + ), + pytest.param( + { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tool-1", + "name": "get_weather", + "input": HUGE_TEXT, + } + } + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "id": "tool-1", + "name": "get_weather", + "arguments": HUGE_TEXT, + } + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "id": "tool-1", + "name": "get_weather", + "arguments": TRUNCATED_TEXT, + } + ], + }, + id="tool_call_arguments", + ), + pytest.param( + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "tool-1", + "content": [{"json": {"blob": HUGE_TEXT}}], + } + } + ], + }, + { + "role": "tool", + "parts": [ + { + "type": "tool_call_response", + "id": "tool-1", + "result": {"blob": HUGE_TEXT}, + } + ], + }, + { + "role": "tool", + "parts": [ + { + "type": "tool_call_response", + "id": "tool-1", + "result": json.dumps({"blob": HUGE_TEXT})[ + :MAX_SINGLE_MESSAGE_CONTENT_CHARS + ] + + "...", + } + ], + }, + id="tool_result_json", + ), + pytest.param( + { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tool-1", + "name": "get_weather", + "input": {"city": HUGE_TEXT}, + } + } + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "id": "tool-1", + "name": "get_weather", + "arguments": {"city": HUGE_TEXT}, + } + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "id": "tool-1", + "name": "get_weather", + "arguments": json.dumps({"city": HUGE_TEXT})[ + :MAX_SINGLE_MESSAGE_CONTENT_CHARS + ] + + "...", + } + ], + }, + id="tool_call_arguments_json", + ), +] + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +@pytest.mark.parametrize( + "init_kwargs, include_prompts, record_inputs, record_outputs", + GEN_AI_DATA_COLLECTION_CASES, +) +@pytest.mark.parametrize( + "last_message, expected_last_message, expected_truncated_last_message", + LAST_MESSAGE_CASES, +) +def test_converse( + init_bedrock, + capture_items, + trace_mode, + init_kwargs, + include_prompts, + record_inputs, + record_outputs, + last_message, + expected_last_message, + expected_truncated_last_message, +): + client = init_bedrock(trace_mode, include_prompts=include_prompts, **init_kwargs) + params = { + "modelId": MODEL_ID, + "system": [{"text": "You are a helpful assistant."}], + "messages": [ + {"role": "user", "content": [{"text": "Hi!"}]}, + {"role": "assistant", "content": [{"text": "Hi! What do you need?"}]}, + last_message, + ], + "inferenceConfig": { + "maxTokens": 100, + "temperature": 0.5, + "topP": 0.9, + "stopSequences": ["STOP"], + }, + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": "get_weather", + "description": "Get the weather for a city.", + "inputSchema": { + "json": { + "type": "object", + "properties": {"city": {"type": "string"}}, + } + }, + } + } + ] + }, + } + + original_params = copy.deepcopy(params) + + with Stubber(client) as stubber: + stubber.add_response("converse", CONVERSE_RESPONSE, original_params) + responses = [] + spans = _capture_spans( + capture_items, + trace_mode, + lambda: responses.append(client.converse(**params)), + ) + + assert responses[0]["output"] == CONVERSE_RESPONSE["output"] + # the request isn't changed by the truncation. + assert params == original_params + + span = _get_bedrock_span(spans) + attributes = span["attributes"] + assert span["name"] == f"chat {MODEL_ID}" + assert span["status"] == "ok" + assert attributes[SPANDATA.SENTRY_OP] == OP.GEN_AI_CHAT + assert attributes[SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert attributes[SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert attributes[SPANDATA.GEN_AI_REQUEST_MODEL] == MODEL_ID + assert attributes[SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + assert attributes[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert attributes[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.5 + assert attributes[SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + assert attributes[SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES] == ["STOP"] + assert attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == ["end_turn"] + assert attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 12 + assert attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 8 + assert attributes[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 20 + assert SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS not in attributes + # the generic AWS client attributes are kept. + assert attributes[SPANDATA.RPC_METHOD] == "Converse" + assert attributes[SPANDATA.CLOUD_REGION] == "us-east-1" + + if record_inputs: + assert json.loads(attributes[SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) == [ + {"type": "text", "content": "You are a helpful assistant."} + ] + if trace_mode.get("stream_gen_ai_spans") is False: + expected_input_messages = [expected_truncated_last_message] + else: + expected_input_messages = [ + {"role": "user", "parts": [{"type": "text", "content": "Hi!"}]}, + { + "role": "assistant", + "parts": [{"type": "text", "content": "Hi! What do you need?"}], + }, + expected_last_message, + ] + assert ( + json.loads(attributes[SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_messages + ) + assert json.loads(attributes[SPANDATA.GEN_AI_TOOL_DEFINITIONS]) == [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + } + ] + else: + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in attributes + assert SPANDATA.GEN_AI_INPUT_MESSAGES not in attributes + assert SPANDATA.GEN_AI_TOOL_DEFINITIONS not in attributes + + if record_outputs: + assert json.loads(attributes[SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [{"type": "text", "content": "Hello! How can I help?"}], + } + ] + else: + assert SPANDATA.GEN_AI_OUTPUT_MESSAGES not in attributes + + +@pytest.mark.skipif( + not _has_shape_member("ContentBlock", "reasoningContent"), + reason="botocore doesn't support reasoning content", +) +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_converse_tool_use_and_reasoning(init_bedrock, capture_items, trace_mode): + client = init_bedrock(trace_mode, send_default_pii=True) + params = { + "modelId": MODEL_ID, + "messages": [ + {"role": "user", "content": [{"text": "Weather in Paris?"}]}, + { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": {"text": "I need the weather tool."} + } + }, + { + "toolUse": { + "toolUseId": "tool-1", + "name": "get_weather", + "input": {"city": "Paris"}, + } + }, + ], + }, + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "tool-1", + "content": [{"json": {"weather": "rainy"}}], + } + }, + {"text": "And in London?"}, + {"image": {"format": "png", "source": {"bytes": b"\x89PNG"}}}, + ], + }, + ], + } + response = { + **CONVERSE_RESPONSE, + "output": { + "message": { + "role": "assistant", + "content": [ + {"text": "Let me check."}, + { + "toolUse": { + "toolUseId": "tool-2", + "name": "get_weather", + "input": {"city": "London"}, + } + }, + ], + } + }, + "stopReason": "tool_use", + } + + with Stubber(client) as stubber: + stubber.add_response("converse", response, params) + spans = _capture_spans( + capture_items, trace_mode, lambda: client.converse(**params) + ) + + attributes = _get_bedrock_span(spans)["attributes"] + expected_input_messages = [ + {"role": "user", "parts": [{"type": "text", "content": "Weather in Paris?"}]}, + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "I need the weather tool."}, + { + "type": "tool_call", + "id": "tool-1", + "name": "get_weather", + "arguments": {"city": "Paris"}, + }, + ], + }, + { + "role": "tool", + "parts": [ + { + "type": "tool_call_response", + "id": "tool-1", + "result": {"weather": "rainy"}, + } + ], + }, + # media content isn't recorded. + {"role": "user", "parts": [{"type": "text", "content": "And in London?"}]}, + ] + assert json.loads( + attributes[SPANDATA.GEN_AI_INPUT_MESSAGES] + ) == _expected_input_messages(trace_mode, expected_input_messages) + assert json.loads(attributes[SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Let me check."}, + { + "type": "tool_call", + "id": "tool-2", + "name": "get_weather", + "arguments": {"city": "London"}, + }, + ], + } + ] + assert attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == ["tool_use"] + + +@pytest.mark.skipif( + not _has_shape_member("TokenUsage", "cacheReadInputTokens"), + reason="botocore doesn't support prompt caching", +) +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_converse_prompt_cache_usage(init_bedrock, capture_items, trace_mode): + client = init_bedrock(trace_mode) + params = { + "modelId": MODEL_ID, + "messages": [{"role": "user", "content": [{"text": "Hello!"}]}], + } + response = { + **CONVERSE_RESPONSE, + "usage": { + "inputTokens": 4, + "outputTokens": 8, + "totalTokens": 1662, + "cacheReadInputTokens": 1000, + "cacheWriteInputTokens": 650, + }, + } + + with Stubber(client) as stubber: + stubber.add_response("converse", response, params) + spans = _capture_spans( + capture_items, trace_mode, lambda: client.converse(**params) + ) + + attributes = _get_bedrock_span(spans)["attributes"] + # Bedrock's `inputTokens` excludes cached tokens; `gen_ai.usage.input_tokens` doesn't. + assert attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 1654 + assert attributes[SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 1000 + assert attributes[SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS] == 650 + assert attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 8 + assert attributes[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 1662 + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_converse_error_is_unchanged(init_bedrock, capture_items, trace_mode): + client = init_bedrock(trace_mode) + params = { + "modelId": MODEL_ID, + "messages": [{"role": "user", "content": [{"text": "Hello!"}]}], + } + errors = [] + + def call(): + with pytest.raises(ClientError) as exc_info: + client.converse(**params) + errors.append(exc_info.value) + + with Stubber(client) as stubber: + stubber.add_client_error( + "converse", + service_error_code="ThrottlingException", + http_status_code=429, + expected_params=params, + ) + spans = _capture_spans(capture_items, trace_mode, call) + + assert errors[0].response["Error"]["Code"] == "ThrottlingException" + + span = _get_bedrock_span(spans) + attributes = span["attributes"] + assert span["name"] == f"chat {MODEL_ID}" + assert span["status"] == "error" + assert attributes[SPANDATA.SENTRY_OP] == OP.GEN_AI_CHAT + assert attributes[SPANDATA.GEN_AI_REQUEST_MODEL] == MODEL_ID + assert attributes[SPANDATA.ERROR_TYPE] == "ThrottlingException" + assert attributes[SPANDATA.HTTP_STATUS_CODE] == 429 + assert SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS not in attributes + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_other_operations_keep_generic_span(init_bedrock, capture_items, trace_mode): + client = init_bedrock(trace_mode, send_default_pii=True) + + with Stubber(client) as stubber: + stubber.add_response("list_async_invokes", {"asyncInvokeSummaries": []}, {}) + spans = _capture_spans( + capture_items, trace_mode, lambda: client.list_async_invokes() + ) + + span = _get_bedrock_span(spans, origin=ORIGIN) + assert span["name"] == "aws.bedrock-runtime.ListAsyncInvokes" + assert span["attributes"][SPANDATA.SENTRY_OP] == OP.HTTP_CLIENT + assert not any(key.startswith("gen_ai.") for key in span["attributes"])