From 1a1f370bffda84cd4d16cdc920d9daf86de6a6ea Mon Sep 17 00:00:00 2001 From: "sentry-junior[bot]" <264270552+sentry-junior[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 22:36:57 +0000 Subject: [PATCH 1/2] feat(boto3): Instrument Bedrock Runtime InvokeModel calls as gen_ai spans Record InvokeModel calls as `gen_ai.generate_content` spans. Request and response bodies are parsed per model family (Anthropic Claude, Amazon Nova and Titan, Meta Llama, Cohere Command and Command R, Mistral), matching the JavaScript SDK. The response body is recorded as the caller reads it, so it isn't consumed or replaced, and the span finishes once it's read or closed. Token usage falls back to Bedrock's usage response headers. Refs #5470 Co-Authored-By: Sergiy Dybskiy <3760543+sergical@users.noreply.github.com> --- sentry_sdk/consts.py | 1 + sentry_sdk/integrations/boto3/_client.py | 11 +- .../integrations/boto3/_instrumentation.py | 58 +- .../integrations/boto3/_services/base.py | 13 + .../boto3/_services/bedrock_runtime.py | 595 +++++++++++++++-- .../boto3/test_bedrock_runtime.py | 612 ++++++++++++++++++ 6 files changed, 1214 insertions(+), 76 deletions(-) diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index db212a83c3..c396b37302 100644 --- a/sentry_sdk/consts.py +++ b/sentry_sdk/consts.py @@ -1321,6 +1321,7 @@ class OP: GEN_AI_EMBEDDINGS = "gen_ai.embeddings" GEN_AI_EVALUATE = "gen_ai.evaluate" GEN_AI_EXECUTE_TOOL = "gen_ai.execute_tool" + GEN_AI_GENERATE_CONTENT = "gen_ai.generate_content" GEN_AI_TEXT_COMPLETION = "gen_ai.text_completion" GEN_AI_HANDOFF = "gen_ai.handoff" GEN_AI_INVOKE_AGENT = "gen_ai.invoke_agent" diff --git a/sentry_sdk/integrations/boto3/_client.py b/sentry_sdk/integrations/boto3/_client.py index 32015c802e..07f5f0620a 100644 --- a/sentry_sdk/integrations/boto3/_client.py +++ b/sentry_sdk/integrations/boto3/_client.py @@ -144,9 +144,14 @@ def sentry_patched_make_api_call( streaming_body_instrumented = False with capture_internal_exceptions(): - streaming_body_instrumented = _instrument_streaming_body(span, parsed) - - # `StreamingBody`s finish their span when consumed or closed. + if service_ext is not None: + streaming_body_instrumented = service_ext.instrument_response( + ctx, span, parsed + ) + if not streaming_body_instrumented: + streaming_body_instrumented = _instrument_streaming_body(span, parsed) + + # streaming responses finish their span when consumed or closed. if not streaming_body_instrumented: _finish_span(span) return parsed diff --git a/sentry_sdk/integrations/boto3/_instrumentation.py b/sentry_sdk/integrations/boto3/_instrumentation.py index ac70b01464..d947b3e635 100644 --- a/sentry_sdk/integrations/boto3/_instrumentation.py +++ b/sentry_sdk/integrations/boto3/_instrumentation.py @@ -25,7 +25,7 @@ ) if TYPE_CHECKING: - from typing import Any, Dict, Mapping, Optional, Union + from typing import Any, Callable, Dict, Mapping, Optional, Union from sentry_sdk._types import Attributes from sentry_sdk.integrations.boto3._context import AwsCallContext @@ -294,17 +294,8 @@ def _instrument_streaming_body( origin=ORIGIN, ) - finished = False - read_in_progress = False - def finish_span(error: "Optional[BaseException]" = None) -> None: - nonlocal finished - if finished: - return - - finished = True - # finish the stream span before the client span, and only once across - # read and close. + # finish the stream span before the client span. if error is not None: with capture_internal_exceptions(): attributes = _get_error_attributes(error) @@ -314,6 +305,32 @@ def finish_span(error: "Optional[BaseException]" = None) -> None: _finish_span(streaming_span, error) _finish_span(span, error) + _patch_streaming_body(body, finish_span) + return True + + +def _patch_streaming_body( + body: "StreamingBody", + on_finish: "Callable[[Optional[BaseException]], None]", + on_read: "Optional[Callable[[bytes], None]]" = None, +) -> None: + """ + Call `on_finish` once, when `body` is consumed, closed, or fails to be read. + `on_read` receives the bytes returned by each `read()`. + """ + finished = False + read_in_progress = False + + def finish(error: "Optional[BaseException]" = None) -> None: + nonlocal finished + if finished: + return + + # only finish once across read and close. + finished = True + with capture_internal_exceptions(): + on_finish(error) + def content_length_reached() -> bool: content_length = body._content_length # type: ignore[attr-defined] return content_length is not None and body._amount_read >= int(content_length) # type: ignore[attr-defined] @@ -323,6 +340,9 @@ def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes: read_in_progress = True try: read_return_value = orig_read(*args, **kwargs) + if on_read is not None: + with capture_internal_exceptions(): + on_read(read_return_value) with capture_internal_exceptions(): amount_of_bytes_requested = args[0] if args else kwargs.get("amt") # detect read-to-end, eof, or the known content length being consumed. @@ -332,10 +352,10 @@ def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes: or (amount_of_bytes_requested > 0 and not read_return_value) or content_length_reached() ): - finish_span() + finish() return read_return_value except BaseException as error: - finish_span(error) + finish(error) raise finally: read_in_progress = False @@ -343,18 +363,18 @@ def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes: def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None: try: orig_close(*args, **kwargs) - finish_span() + finish() except BaseException as error: - finish_span(error) + finish(error) raise def sentry_raw_stream_close(*args: "Any", **kwargs: "Any") -> None: try: orig_raw_close(*args, **kwargs) if not read_in_progress: - finish_span() + finish() except BaseException as error: - finish_span(error) + finish(error) raise try: @@ -367,11 +387,9 @@ def sentry_raw_stream_close(*args: "Any", **kwargs: "Any") -> None: body.read = sentry_streaming_body_read # type: ignore body.close = sentry_streaming_body_close # type: ignore except Exception: - finish_span() + finish() raise - return True - def _set_request_attributes( span: "Union[Span, StreamedSpan]", diff --git a/sentry_sdk/integrations/boto3/_services/base.py b/sentry_sdk/integrations/boto3/_services/base.py index be3f8add7b..bb41755aa3 100644 --- a/sentry_sdk/integrations/boto3/_services/base.py +++ b/sentry_sdk/integrations/boto3/_services/base.py @@ -50,3 +50,16 @@ def get_response_attributes( ) -> "Attributes": """Return response attributes to add after the AWS request is made.""" return {} + + def instrument_response( + self, + ctx: "AwsCallContext", + span: "Union[Span, StreamedSpan]", + response: "Any", + ) -> bool: + """ + Instrument a response that is read after the AWS request is made, e.g. a + streaming body. Return `True` if the extension finishes `span` once the + response is read, or `False` to finish it when the request is made. + """ + return False diff --git a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py index 516971a670..3f1e75550e 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py @@ -8,6 +8,13 @@ truncate_and_annotate_messages, ) from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import DidNotEnable +from sentry_sdk.integrations.boto3._instrumentation import ( + _finish_span, + _get_error_attributes, + _patch_streaming_body, + _set_span_attributes, +) 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 @@ -16,7 +23,7 @@ from sentry_sdk.utils import capture_internal_exceptions, has_data_collection_enabled if TYPE_CHECKING: - from typing import Any, Literal, Optional, Union + from typing import Any, Callable, Iterable, Literal, Optional, Union from sentry_sdk._types import ( Attributes, @@ -35,6 +42,11 @@ ContentPart = Union[TextPart, ReasoningPart, ToolCallPart, ToolCallResponsePart] +try: + from botocore.response import StreamingBody +except ImportError: + raise DidNotEnable("botocore not installed") + # `gen_ai.provider.name` of Amazon Bedrock. # https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#gen-ai-provider-name @@ -43,8 +55,93 @@ # `gen_ai.operation.name` and span op of each instrumented operation. GEN_AI_OPERATIONS = { "Converse": ("chat", OP.GEN_AI_CHAT), + "InvokeModel": ("generate_content", OP.GEN_AI_GENERATE_CONTENT), +} + +# `InvokeModel` request and response bodies have a model-specific format. For each +# model family: the object that holds the inference parameters in the request body, +# or `None` for the top level, and the attribute of each parameter. +# The model family is matched in the model ID, which can also be an inference profile +# ID or an ARN, e.g. `us.anthropic.claude-3-5-sonnet-20240620-v1:0`, in this order; +# `cohere.command-r` has to be matched before `cohere.command`. +# https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html +INVOKE_MODEL_PARAMETERS = { + "amazon.titan": ( + "textGenerationConfig", + { + "maxTokenCount": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "topP": SPANDATA.GEN_AI_REQUEST_TOP_P, + "stopSequences": SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES, + }, + ), + "amazon.nova": ( + "inferenceConfig", + { + "max_new_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "stopSequences": SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES, + }, + ), + "anthropic.claude": ( + None, + { + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "stop_sequences": SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES, + }, + ), + "meta.llama": ( + None, + { + "max_gen_len": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + }, + ), + "cohere.command-r": ( + None, + { + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "stop_sequences": SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES, + }, + ), + "cohere.command": ( + None, + { + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "stop_sequences": SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES, + }, + ), + "mistral": ( + None, + { + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "stop": SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES, + }, + ), +} + +# the prompt of model families that take a single text prompt in `InvokeModel`. +INVOKE_MODEL_PROMPTS = { + "amazon.titan": "inputText", + "meta.llama": "prompt", + "cohere.command-r": "message", + "cohere.command": "prompt", + "mistral": "prompt", } +# larger `InvokeModel` response bodies aren't parsed, e.g. image generation results. +MAX_INVOKE_MODEL_RESPONSE_BYTES = 1_000_000 + class _BedrockRuntimeExtension(_ServiceExtension): """ @@ -90,34 +187,111 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": if model_id: attributes[SPANDATA.GEN_AI_REQUEST_MODEL] = model_id - with capture_internal_exceptions(): - attributes.update(_get_converse_request_attributes(ctx.params)) + if ctx.operation_name != "InvokeModel": + 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") + # e.g. a span dropped by `ignore_spans`. + if ctx.operation_name not in GEN_AI_OPERATIONS or isinstance( + span, NoOpStreamedSpan ): return {} + # the request body is parsed once, for its parameters and input messages. + if ctx.operation_name == "InvokeModel": + return _get_invoke_model_request_attributes( + span, + _get_model_family(ctx.params.get("modelId")), + ctx.params.get("body"), + ) + + if not _should_record_gen_ai("inputs"): + return {} + return _get_input_messages_attributes( - span, _transform_input_messages(ctx.params.get("messages")) + span, + _transform_input_messages( + ctx.params.get("messages"), _transform_content_blocks + ), ) def get_response_attributes( self, ctx: "AwsCallContext", response: "Any" ) -> "Attributes": - if ctx.operation_name not in GEN_AI_OPERATIONS: + if ctx.operation_name != "Converse": return {} return _get_converse_response_attributes(response) + def instrument_response( + self, + ctx: "AwsCallContext", + span: "Union[Span, StreamedSpan]", + response: "Any", + ) -> bool: + """ + The `InvokeModel` response body is a `StreamingBody`. Its bytes are recorded + as the caller reads them, and the span finishes once it's read or closed. + """ + if ctx.operation_name != "InvokeModel" or isinstance(span, NoOpStreamedSpan): + return False + + body = response.get("body") if isinstance(response, Mapping) else None + if not isinstance(body, StreamingBody): + return False + + model_family = _get_model_family(ctx.params.get("modelId")) + headers = response.get("ResponseMetadata", {}).get("HTTPHeaders", {}) + chunks: "list[bytes]" = [] + size = 0 + + def on_read(data: bytes) -> None: + nonlocal size + if model_family is None or size > MAX_INVOKE_MODEL_RESPONSE_BYTES: + return + + size += len(data) + if size > MAX_INVOKE_MODEL_RESPONSE_BYTES: + chunks.clear() + else: + chunks.append(data) + + def on_finish(error: "Optional[BaseException]") -> None: + attributes: "Attributes" = {} + if error is not None: + with capture_internal_exceptions(): + attributes.update(_get_error_attributes(error)) + else: + with capture_internal_exceptions(): + attributes.update( + _get_invoke_model_response_attributes( + model_family, b"".join(chunks), headers + ) + ) + + with capture_internal_exceptions(): + _set_span_attributes(span, attributes) + _finish_span(span, error) + + _patch_streaming_body(body, on_finish, on_read) + return True + + +def _get_model_family(model_id: "Any") -> "Optional[str]": + if not isinstance(model_id, str): + return None + + for model_family in INVOKE_MODEL_PARAMETERS: + if model_family in model_id: + return model_family + + return None + def _should_record_gen_ai(kind: 'Literal["inputs", "outputs"]') -> bool: """ @@ -161,24 +335,107 @@ def _get_converse_request_attributes(params: "Mapping[str, Any]") -> "Attributes inference_config["stopSequences"] ) + if _should_record_gen_ai("inputs"): + attributes.update(_get_converse_input_attributes(params)) + + return attributes + + +def _get_converse_input_attributes(request: "Mapping[str, Any]") -> "Attributes": + """ + Get the input attributes, except the input messages, of a request in the Converse + format, which `InvokeModel` requests to Amazon Nova models also use. + """ + tool_config = request.get("toolConfig") + return _get_input_attributes( + _transform_system_instructions(request.get("system")), + _transform_tool_definitions( + tool_config.get("tools") if isinstance(tool_config, Mapping) else None + ), + ) + + +def _get_invoke_model_request_attributes( + span: "Union[Span, StreamedSpan]", model_family: "Optional[str]", body: "Any" +) -> "Attributes": + """ + A request body that is a file-like object isn't read, as that would consume it. + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html + """ + if model_family is None or not isinstance(body, (bytes, bytearray, str)): + return {} + + request = json.loads(body) + + if not isinstance(request, Mapping): + return {} + + attributes: "Attributes" = {} + + parameters_key, parameters = INVOKE_MODEL_PARAMETERS[model_family] + request_parameters = ( + request.get(parameters_key) if parameters_key is not None else request + ) + if isinstance(request_parameters, Mapping): + for key, attribute in parameters.items(): + value = request_parameters.get(key) + if isinstance(value, (int, float, str, list)): + attributes[attribute] = value + if not _should_record_gen_ai("inputs"): return attributes + input_messages: "list[InputMessage]" = [] + if model_family == "amazon.nova": + attributes.update(_get_converse_input_attributes(request)) + input_messages = _transform_input_messages( + request.get("messages"), _transform_content_blocks + ) + elif model_family == "anthropic.claude": + attributes.update( + _get_input_attributes( + _transform_system_instructions(request.get("system")), + _transform_anthropic_tool_definitions(request.get("tools")), + ) + ) + input_messages = _transform_input_messages( + request.get("messages"), _transform_anthropic_content + ) + else: + # Cohere Command R's system prompt. + attributes.update( + _get_input_attributes( + _transform_system_instructions(request.get("preamble")), [] + ) + ) + prompt = request.get(INVOKE_MODEL_PROMPTS[model_family]) + if isinstance(prompt, str): + input_messages = [ + {"role": "user", "parts": [{"type": "text", "content": prompt}]} + ] + + with capture_internal_exceptions(): + attributes.update(_get_input_messages_attributes(span, input_messages)) + + return attributes + + +def _get_input_attributes( + system_instructions: "list[TextPart]", tool_definitions: "list[ToolDefinition]" +) -> "Attributes": + attributes: "Attributes" = {} + + # serialize each attribute separately so that an unserializable value only drops + # the attribute it's in. 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 - ) + if tool_definitions: + attributes[SPANDATA.GEN_AI_TOOL_DEFINITIONS] = json.dumps(tool_definitions) return attributes @@ -246,39 +503,163 @@ def _get_converse_response_attributes(response: "Any") -> "Attributes": usage = response.get("usage") if isinstance(usage, Mapping): - attributes.update(_get_usage_attributes(usage)) + attributes.update( + _get_usage_attributes( + usage.get("inputTokens"), + usage.get("outputTokens"), + usage.get("cacheReadInputTokens"), + usage.get("cacheWriteInputTokens"), + usage.get("totalTokens"), + ) + ) - if not _should_record_gen_ai("outputs"): - return attributes + output = response.get("output") + message = output.get("message") if isinstance(output, Mapping) else None + if isinstance(message, Mapping): + attributes.update( + _get_output_attributes( + [_to_output_message(_transform_content_blocks(message.get("content")))] + ) + ) - with capture_internal_exceptions(): - output = response.get("output") - message = output.get("message") if isinstance(output, Mapping) else None + return attributes + + +def _get_invoke_model_response_attributes( + model_family: "Optional[str]", body: bytes, headers: "Mapping[str, str]" +) -> "Attributes": + """ + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html + """ + attributes: "Attributes" = {} + + response = None + if model_family is not None and body: + try: + response = json.loads(body) + except ValueError: + # e.g. a body that was closed before being read completely. + pass + + if isinstance(response, Mapping): + with capture_internal_exceptions(): + attributes.update(_parse_invoke_model_response(model_family, response)) + + # Bedrock reports the token usage of every model in response headers. + for attribute, header in ( + (SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, "x-amzn-bedrock-input-token-count"), + (SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, "x-amzn-bedrock-output-token-count"), + ): + if attribute not in attributes and header in headers: + with capture_internal_exceptions(): + attributes[attribute] = int(headers[header]) + + return attributes + + +def _parse_invoke_model_response( + model_family: "Optional[str]", response: "Mapping[str, Any]" +) -> "Attributes": + attributes: "Attributes" = {} + finish_reasons: "list[Any]" = [] + output_messages: "list[OutputMessage]" = [] + + if model_family == "amazon.nova": + usage = response.get("usage") or {} + attributes.update( + _get_usage_attributes( + usage.get("inputTokens"), + usage.get("outputTokens"), + usage.get("cacheReadInputTokenCount"), + usage.get("cacheWriteInputTokenCount"), + usage.get("totalTokens"), + ) + ) + finish_reasons = [response.get("stopReason")] + message = (response.get("output") or {}).get("message") 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]) + output_messages = [ + _to_output_message(_transform_content_blocks(message.get("content"))) + ] + elif model_family == "anthropic.claude": + usage = response.get("usage") or {} + attributes.update( + _get_usage_attributes( + usage.get("input_tokens"), + usage.get("output_tokens"), + usage.get("cache_read_input_tokens"), + usage.get("cache_creation_input_tokens"), + ) + ) + finish_reasons = [response.get("stop_reason")] + output_messages = [ + _to_output_message(_transform_anthropic_content(response.get("content"))) + ] + + elif model_family == "amazon.titan": + results = response.get("results") or [] + attributes.update( + _get_usage_attributes( + response.get("inputTextTokenCount"), + results[0].get("tokenCount") if results else None, + ) + ) + finish_reasons = [result.get("completionReason") for result in results] + output_messages = _to_text_output_messages( + result.get("outputText") for result in results + ) + + elif model_family == "meta.llama": + attributes.update( + _get_usage_attributes( + response.get("prompt_token_count"), + response.get("generation_token_count"), + ) + ) + finish_reasons = [response.get("stop_reason")] + output_messages = _to_text_output_messages([response.get("generation")]) + + elif model_family == "cohere.command-r": + finish_reasons = [response.get("finish_reason")] + output_messages = _to_text_output_messages([response.get("text")]) + + elif model_family == "cohere.command": + generations = response.get("generations") or [] + finish_reasons = [generation.get("finish_reason") for generation in generations] + output_messages = _to_text_output_messages( + generation.get("text") for generation in generations + ) + + elif model_family == "mistral": + outputs = response.get("outputs") or [] + finish_reasons = [output.get("stop_reason") for output in outputs] + output_messages = _to_text_output_messages( + output.get("text") for output in outputs + ) + + finish_reasons = [reason for reason in finish_reasons if isinstance(reason, str)] + if finish_reasons: + attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] = finish_reasons + + attributes.update(_get_output_attributes(output_messages)) return attributes -def _get_usage_attributes(usage: "Mapping[str, Any]") -> "Attributes": +def _get_usage_attributes( + input_tokens: "Any", + output_tokens: "Any", + cache_read_input_tokens: "Any" = None, + cache_write_input_tokens: "Any" = None, + total_tokens: "Any" = None, +) -> "Attributes": """ - Bedrock's `inputTokens` excludes the tokens read from or written to the prompt + Bedrock's input tokens exclude 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 @@ -301,7 +682,42 @@ def _get_usage_attributes(usage: "Mapping[str, Any]") -> "Attributes": return attributes +def _get_output_attributes(output_messages: "list[OutputMessage]") -> "Attributes": + attributes: "Attributes" = {} + if not output_messages or not _should_record_gen_ai("outputs"): + return attributes + + with capture_internal_exceptions(): + attributes[SPANDATA.GEN_AI_OUTPUT_MESSAGES] = json.dumps(output_messages) + + return attributes + + +def _to_output_message(parts: "list[ContentPart]") -> "OutputMessage": + output_message: "OutputMessage" = {"role": "assistant", "parts": []} + for part in parts: + # models don't respond with tool results. + if part["type"] != "tool_call_response": + output_message["parts"].append(part) + + return output_message + + +def _to_text_output_messages(texts: "Iterable[Any]") -> "list[OutputMessage]": + return [ + {"role": "assistant", "parts": [{"type": "text", "content": text}]} + for text in texts + if isinstance(text, str) + ] + + def _transform_system_instructions(system: "Any") -> "list[TextPart]": + """ + Transform a system prompt, or a list of Converse or Anthropic system content blocks. + """ + if isinstance(system, str): + return [{"type": "text", "content": system}] + if not isinstance(system, (list, tuple)): return [] @@ -312,10 +728,12 @@ def _transform_system_instructions(system: "Any") -> "list[TextPart]": ] -def _transform_input_messages(messages: "Any") -> "list[InputMessage]": +def _transform_input_messages( + messages: "Any", transform_content: "Callable[[Any], list[ContentPart]]" +) -> "list[InputMessage]": """ - Bedrock sends tool results in `user` messages; they're recorded as `tool` - messages, which is how `gen_ai.input.messages` represents them. + Tool results are sent 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 [] @@ -331,7 +749,7 @@ def _transform_input_messages(messages: "Any") -> "list[InputMessage]": parts: "list[ContentPart]" = [] tool_parts: "list[ContentPart]" = [] - for part in _transform_content_blocks(message.get("content")): + for part in transform_content(message.get("content")): if part["type"] == "tool_call_response": tool_parts.append(part) else: @@ -402,6 +820,57 @@ def _transform_content_blocks(content: "Any") -> "list[ContentPart]": return parts +def _transform_anthropic_content(content: "Any") -> "list[ContentPart]": + """ + Transform the text, thinking, tool use, and tool result content of an Anthropic + Claude message. Media content, e.g. images and documents, is not recorded. + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + """ + if isinstance(content, str): + return [{"type": "text", "content": content}] + + if not isinstance(content, (list, tuple)): + return [] + + parts: "list[ContentPart]" = [] + for block in content: + if not isinstance(block, Mapping): + continue + + block_type = block.get("type") + if block_type == "text" and isinstance(block.get("text"), str): + parts.append({"type": "text", "content": block["text"]}) + + elif block_type == "thinking" and isinstance(block.get("thinking"), str): + parts.append({"type": "reasoning", "content": block["thinking"]}) + + elif block_type == "tool_use": + tool_call: "ToolCallPart" = {"type": "tool_call"} + if "id" in block: + tool_call["id"] = block["id"] + if "name" in block: + tool_call["name"] = block["name"] + if "input" in block: + tool_call["arguments"] = block["input"] + parts.append(tool_call) + + elif block_type == "tool_result": + tool_response: "ToolCallResponsePart" = {"type": "tool_call_response"} + if "tool_use_id" in block: + tool_response["id"] = block["tool_use_id"] + + result = [ + part["content"] + for part in _transform_anthropic_content(block.get("content")) + if part["type"] == "text" + ] + 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 @@ -415,17 +884,37 @@ def _transform_tool_definitions(tools: "Any") -> "list[ToolDefinition]": 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) + tool_definitions.append( + _to_tool_definition( + tool_spec, + input_schema.get("json") if isinstance(input_schema, Mapping) else None, + ) + ) return tool_definitions + + +def _transform_anthropic_tool_definitions(tools: "Any") -> "list[ToolDefinition]": + if not isinstance(tools, (list, tuple)): + return [] + + return [ + _to_tool_definition(tool, tool.get("input_schema")) + for tool in tools + if isinstance(tool, Mapping) + ] + + +def _to_tool_definition( + tool: "Mapping[str, Any]", parameters: "Any" +) -> "ToolDefinition": + tool_definition: "ToolDefinition" = {"type": "function"} + if "name" in tool: + tool_definition["name"] = tool["name"] + if "description" in tool: + tool_definition["description"] = tool["description"] + if isinstance(parameters, Mapping): + tool_definition["parameters"] = dict(parameters) + + return tool_definition diff --git a/tests/integrations/boto3/test_bedrock_runtime.py b/tests/integrations/boto3/test_bedrock_runtime.py index 257d350f08..29bb8bce94 100644 --- a/tests/integrations/boto3/test_bedrock_runtime.py +++ b/tests/integrations/boto3/test_bedrock_runtime.py @@ -1,8 +1,10 @@ import copy import json +from io import BytesIO import boto3 import pytest +from botocore.awsrequest import AWSResponse from botocore.exceptions import ClientError from botocore.stub import Stubber @@ -11,6 +13,7 @@ from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.boto3 import Boto3Integration from sentry_sdk.integrations.boto3.consts import AI_ORIGIN, ORIGIN +from tests.integrations.boto3.aws_mock import Body, MockResponse session = boto3.Session( # type: ignore[attr-defined] aws_access_key_id="-", @@ -170,6 +173,18 @@ def _expected_input_messages(trace_mode, input_messages): return input_messages +def _expected_gen_ai_attributes(trace_mode, attributes): + if SPANDATA.GEN_AI_INPUT_MESSAGES not in attributes: + return attributes + + return { + **attributes, + SPANDATA.GEN_AI_INPUT_MESSAGES: _expected_input_messages( + trace_mode, attributes[SPANDATA.GEN_AI_INPUT_MESSAGES] + ), + } + + def _get_bedrock_span(spans, origin=AI_ORIGIN): (span,) = ( span @@ -742,3 +757,600 @@ def test_other_operations_keep_generic_span(init_bedrock, capture_items, trace_m 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"]) + + +INVOKE_MODEL_CASES = [ + pytest.param( + "us.anthropic.claude-3-5-sonnet-20240620-v1:0", + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 100, + "temperature": 0.5, + "top_p": 0.9, + "stop_sequences": ["STOP"], + "system": "Be brief.", + "messages": [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "I need the tool."}, + { + "type": "tool_use", + "id": "tool-1", + "name": "get_weather", + "input": {"city": "Paris"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool-1", + "content": "rainy", + } + ], + }, + ], + "tools": [ + { + "name": "get_weather", + "description": "Get the weather for a city.", + "input_schema": {"type": "object"}, + } + ], + }, + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Hi!"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "cache_read_input_tokens": 100, + "cache_creation_input_tokens": 50, + }, + }, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_REQUEST_TEMPERATURE: 0.5, + SPANDATA.GEN_AI_REQUEST_TOP_P: 0.9, + SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES: ["STOP"], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["end_turn"], + # Anthropic's `input_tokens` excludes cached tokens. + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 160, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 5, + SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS: 100, + SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS: 50, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: [ + {"type": "text", "content": "Be brief."} + ], + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + { + "role": "user", + "parts": [{"type": "text", "content": "Weather in Paris?"}], + }, + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "I need the tool."}, + { + "type": "tool_call", + "id": "tool-1", + "name": "get_weather", + "arguments": {"city": "Paris"}, + }, + ], + }, + { + "role": "tool", + "parts": [ + { + "type": "tool_call_response", + "id": "tool-1", + "result": "rainy", + } + ], + }, + ], + SPANDATA.GEN_AI_TOOL_DEFINITIONS: [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": {"type": "object"}, + } + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ], + }, + id="anthropic_claude", + ), + pytest.param( + "amazon.nova-lite-v1:0", + { + "schemaVersion": "messages-v1", + "system": [{"text": "Be brief."}], + "messages": [{"role": "user", "content": [{"text": "Hello!"}]}], + "inferenceConfig": { + "max_new_tokens": 100, + "temperature": 0.5, + "top_p": 0.9, + "stopSequences": ["STOP"], + }, + }, + { + "output": {"message": {"role": "assistant", "content": [{"text": "Hi!"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + }, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_REQUEST_TEMPERATURE: 0.5, + SPANDATA.GEN_AI_REQUEST_TOP_P: 0.9, + SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES: ["STOP"], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["end_turn"], + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 10, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 5, + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS: 15, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: [ + {"type": "text", "content": "Be brief."} + ], + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + {"role": "user", "parts": [{"type": "text", "content": "Hello!"}]} + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ], + }, + id="amazon_nova", + ), + pytest.param( + "amazon.titan-text-express-v1", + { + "inputText": "Hello!", + "textGenerationConfig": { + "maxTokenCount": 100, + "temperature": 0.5, + "topP": 0.9, + "stopSequences": ["User:"], + }, + }, + { + "inputTextTokenCount": 10, + "results": [ + {"tokenCount": 5, "outputText": "Hi!", "completionReason": "FINISH"} + ], + }, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_REQUEST_TEMPERATURE: 0.5, + SPANDATA.GEN_AI_REQUEST_TOP_P: 0.9, + SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES: ["User:"], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["FINISH"], + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 10, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 5, + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + {"role": "user", "parts": [{"type": "text", "content": "Hello!"}]} + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ], + }, + id="amazon_titan", + ), + pytest.param( + "meta.llama3-8b-instruct-v1:0", + {"prompt": "Hello!", "max_gen_len": 100, "temperature": 0.5, "top_p": 0.9}, + { + "generation": "Hi!", + "prompt_token_count": 10, + "generation_token_count": 5, + "stop_reason": "stop", + }, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_REQUEST_TEMPERATURE: 0.5, + SPANDATA.GEN_AI_REQUEST_TOP_P: 0.9, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["stop"], + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 10, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 5, + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + {"role": "user", "parts": [{"type": "text", "content": "Hello!"}]} + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ], + }, + id="meta_llama", + ), + pytest.param( + "cohere.command-r-v1:0", + { + "message": "Hello!", + "preamble": "Be brief.", + "max_tokens": 100, + "temperature": 0.5, + "p": 0.9, + "stop_sequences": ["STOP"], + }, + {"response_id": "1", "text": "Hi!", "finish_reason": "COMPLETE"}, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_REQUEST_TEMPERATURE: 0.5, + SPANDATA.GEN_AI_REQUEST_TOP_P: 0.9, + SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES: ["STOP"], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["COMPLETE"], + # the response body has no usage; it's in the response headers. + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 7, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 3, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: [ + {"type": "text", "content": "Be brief."} + ], + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + {"role": "user", "parts": [{"type": "text", "content": "Hello!"}]} + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ], + }, + id="cohere_command_r", + ), + pytest.param( + "cohere.command-text-v14", + { + "prompt": "Hello!", + "max_tokens": 100, + "temperature": 0.5, + "p": 0.9, + "stop_sequences": ["STOP"], + }, + { + "id": "1", + "prompt": "Hello!", + "generations": [{"id": "1", "text": "Hi!", "finish_reason": "COMPLETE"}], + }, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_REQUEST_TEMPERATURE: 0.5, + SPANDATA.GEN_AI_REQUEST_TOP_P: 0.9, + SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES: ["STOP"], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["COMPLETE"], + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 7, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 3, + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + {"role": "user", "parts": [{"type": "text", "content": "Hello!"}]} + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ], + }, + id="cohere_command", + ), + pytest.param( + "mistral.mistral-7b-instruct-v0:2", + { + "prompt": "[INST] Hello! [/INST]", + "max_tokens": 100, + "temperature": 0.5, + "top_p": 0.9, + "stop": ["STOP"], + }, + {"outputs": [{"text": "Hi!", "stop_reason": "stop"}]}, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_REQUEST_TEMPERATURE: 0.5, + SPANDATA.GEN_AI_REQUEST_TOP_P: 0.9, + SPANDATA.GEN_AI_REQUEST_STOP_SEQUENCES: ["STOP"], + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["stop"], + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 7, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 3, + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + { + "role": "user", + "parts": [{"type": "text", "content": "[INST] Hello! [/INST]"}], + } + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ], + }, + id="mistral", + ), +] + +# attributes that hold JSON-serialized prompts and responses. +GEN_AI_MESSAGE_ATTRIBUTES = ( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + SPANDATA.GEN_AI_INPUT_MESSAGES, + SPANDATA.GEN_AI_TOOL_DEFINITIONS, + SPANDATA.GEN_AI_OUTPUT_MESSAGES, +) + + +def _invoke_model_response(client, body, input_tokens=7, output_tokens=3): + """ + Respond to `InvokeModel` with `body`, and its token usage in the headers. + """ + return MockResponse( + client, + 200, + { + "content-type": "application/json", + "content-length": str(len(body)), + "x-amzn-requestid": "request-id", + "x-amzn-bedrock-input-token-count": str(input_tokens), + "x-amzn-bedrock-output-token-count": str(output_tokens), + }, + body, + ) + + +def _get_attributes(span): + """ + Return the span's attributes, with the gen_ai message attributes deserialized. + """ + return { + key: json.loads(value) if key in GEN_AI_MESSAGE_ATTRIBUTES else value + for key, value in span["attributes"].items() + } + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +@pytest.mark.parametrize( + "model_id, request_body, response_body, expected_attributes", INVOKE_MODEL_CASES +) +def test_invoke_model( + init_bedrock, + capture_items, + trace_mode, + model_id, + request_body, + response_body, + expected_attributes, +): + client = init_bedrock(trace_mode, send_default_pii=True) + response_bytes = json.dumps(response_body).encode() + bodies = [] + + def call(): + response = client.invoke_model(modelId=model_id, body=json.dumps(request_body)) + bodies.append(response["body"].read()) + + with _invoke_model_response(client, response_bytes): + spans = _capture_spans(capture_items, trace_mode, call) + + assert bodies == [response_bytes] + + span = _get_bedrock_span(spans) + attributes = _get_attributes(span) + assert span["name"] == f"generate_content {model_id}" + assert span["status"] == "ok" + assert attributes[SPANDATA.SENTRY_OP] == OP.GEN_AI_GENERATE_CONTENT + assert attributes[SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert attributes[SPANDATA.GEN_AI_OPERATION_NAME] == "generate_content" + assert attributes[SPANDATA.GEN_AI_REQUEST_MODEL] == model_id + assert attributes[SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + assert attributes[SPANDATA.RPC_METHOD] == "InvokeModel" + assert attributes[SPANDATA.AWS_REQUEST_ID] == "request-id" + gen_ai_attributes = { + key: value + for key, value in attributes.items() + if key.startswith("gen_ai.") + and key + not in ( + SPANDATA.GEN_AI_PROVIDER_NAME, + SPANDATA.GEN_AI_OPERATION_NAME, + SPANDATA.GEN_AI_REQUEST_MODEL, + SPANDATA.GEN_AI_RESPONSE_STREAMING, + ) + } + assert gen_ai_attributes == _expected_gen_ai_attributes( + trace_mode, expected_attributes + ) + # only the gen_ai span is recorded, not an `http.client.stream` span. + assert not any( + span["attributes"].get(SPANDATA.SENTRY_OP) == OP.HTTP_CLIENT_STREAM + for span in spans + ) + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +@pytest.mark.parametrize( + "init_kwargs, include_prompts, record_inputs, record_outputs", + GEN_AI_DATA_COLLECTION_CASES, +) +def test_invoke_model_data_collection( + init_bedrock, + capture_items, + trace_mode, + init_kwargs, + include_prompts, + record_inputs, + record_outputs, +): + client = init_bedrock(trace_mode, include_prompts=include_prompts, **init_kwargs) + model_id, request_body, response_body, expected_attributes = INVOKE_MODEL_CASES[ + 0 + ].values + + def call(): + response = client.invoke_model(modelId=model_id, body=json.dumps(request_body)) + response["body"].read() + + with _invoke_model_response(client, json.dumps(response_body).encode()): + spans = _capture_spans(capture_items, trace_mode, call) + + attributes = _get_attributes(_get_bedrock_span(spans)) + expected_attributes = _expected_gen_ai_attributes(trace_mode, expected_attributes) + for attribute, recorded in ( + (SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, record_inputs), + (SPANDATA.GEN_AI_INPUT_MESSAGES, record_inputs), + (SPANDATA.GEN_AI_TOOL_DEFINITIONS, record_inputs), + (SPANDATA.GEN_AI_OUTPUT_MESSAGES, record_outputs), + ): + if recorded: + assert attributes[attribute] == expected_attributes[attribute] + else: + assert attribute not in attributes + # the other attributes are always recorded. + assert attributes[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == ["end_turn"] + assert attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 160 + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +@pytest.mark.parametrize( + "read_body, parsed", + [ + pytest.param(lambda body: body.read(), True, id="read"), + pytest.param( + lambda body: b"".join(iter(lambda: body.read(4), b"")), + True, + id="read_amt", + ), + pytest.param( + lambda body: b"".join(body.iter_chunks(chunk_size=4)), + True, + id="iter_chunks", + ), + pytest.param( + lambda body: (body.read(4), body.close())[0], False, id="close_early" + ), + ], +) +def test_invoke_model_body_consumption( + init_bedrock, capture_items, trace_mode, read_body, parsed +): + client = init_bedrock(trace_mode, send_default_pii=True) + response_bytes = json.dumps( + {"outputs": [{"text": "Hi!", "stop_reason": "stop"}]} + ).encode() + bodies = [] + + def call(): + response = client.invoke_model( + modelId="mistral.mistral-7b-instruct-v0:2", + body=json.dumps({"prompt": "Hello!"}), + ) + bodies.append(read_body(response["body"])) + + with _invoke_model_response(client, response_bytes): + spans = _capture_spans(capture_items, trace_mode, call) + + assert bodies == [response_bytes if parsed else response_bytes[:4]] + + span = _get_bedrock_span(spans) + attributes = span["attributes"] + assert span["status"] == "ok" + assert attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 7 + assert attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 3 + if parsed: + assert attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == ["stop"] + assert SPANDATA.GEN_AI_OUTPUT_MESSAGES in attributes + else: + assert SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS not in attributes + assert SPANDATA.GEN_AI_OUTPUT_MESSAGES not in attributes + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_invoke_model_body_read_error_is_unchanged( + init_bedrock, capture_items, trace_mode +): + client = init_bedrock(trace_mode) + original_exception = OSError("stream read failed") + + class _FailingBody(Body): + def read(self, *args, **kwargs): + raise original_exception + + def respond(request, **kwargs): + return AWSResponse( + request.url, 200, {"content-length": "10"}, _FailingBody(b"") + ) + + client.meta.events.register("before-send", respond) + + def call(): + response = client.invoke_model( + modelId=MODEL_ID, body=json.dumps({"max_tokens": 100}) + ) + with pytest.raises(OSError) as exc_info: + response["body"].read() + assert exc_info.value is original_exception + + spans = _capture_spans(capture_items, trace_mode, call) + + span = _get_bedrock_span(spans) + assert span["status"] == "error" + assert span["attributes"][SPANDATA.ERROR_TYPE] == "OSError" + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_invoke_model_unknown_model(init_bedrock, capture_items, trace_mode): + client = init_bedrock(trace_mode, send_default_pii=True) + model_id = "ai21.jamba-1-5-mini-v1:0" + + def call(): + response = client.invoke_model( + modelId=model_id, + body=json.dumps( + {"messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 1} + ), + ) + response["body"].read() + + with _invoke_model_response(client, b'{"choices": []}'): + spans = _capture_spans(capture_items, trace_mode, call) + + span = _get_bedrock_span(spans) + assert span["name"] == f"generate_content {model_id}" + assert { + key: value + for key, value in span["attributes"].items() + if key.startswith("gen_ai.") + } == { + SPANDATA.GEN_AI_PROVIDER_NAME: "aws.bedrock", + SPANDATA.GEN_AI_OPERATION_NAME: "generate_content", + SPANDATA.GEN_AI_REQUEST_MODEL: model_id, + SPANDATA.GEN_AI_RESPONSE_STREAMING: False, + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 7, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 3, + } + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_invoke_model_bodies_that_are_not_parsed( + init_bedrock, capture_items, trace_mode +): + client = init_bedrock(trace_mode, send_default_pii=True) + # a file-like request body isn't read, and a large response body isn't parsed. + response_bytes = json.dumps( + {"generation": "x" * 1_000_001, "stop_reason": "length"} + ).encode() + bodies = [] + + def call(): + response = client.invoke_model( + modelId="meta.llama3-8b-instruct-v1:0", + body=BytesIO(json.dumps({"prompt": "Hello!", "max_gen_len": 1}).encode()), + ) + bodies.append(response["body"].read()) + + with _invoke_model_response(client, response_bytes): + spans = _capture_spans(capture_items, trace_mode, call) + + assert bodies == [response_bytes] + + attributes = _get_bedrock_span(spans)["attributes"] + assert SPANDATA.GEN_AI_REQUEST_MAX_TOKENS not in attributes + assert SPANDATA.GEN_AI_INPUT_MESSAGES not in attributes + assert SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS not in attributes + assert SPANDATA.GEN_AI_OUTPUT_MESSAGES not in attributes + assert attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 7 From b8c26346e3836d54e2105ac5644953a751ac831a Mon Sep 17 00:00:00 2001 From: "sentry-junior[bot]" <264270552+sentry-junior[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 23:44:34 +0000 Subject: [PATCH 2/2] fix(boto3): Skip large or malformed Bedrock InvokeModel request bodies Don't parse `InvokeModel` request bodies over 1 MB, e.g. requests with images, and ignore malformed ones quietly, like response bodies. Test that media content of Claude and Nova requests isn't recorded. Refs #5470 Co-Authored-By: Sergiy Dybskiy <3760543+sergical@users.noreply.github.com> --- .../boto3/_services/bedrock_runtime.py | 21 ++- .../boto3/test_bedrock_runtime.py | 170 +++++++++++++++++- 2 files changed, 181 insertions(+), 10 deletions(-) diff --git a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py index 3f1e75550e..302833ad72 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py @@ -139,8 +139,9 @@ "mistral": "prompt", } -# larger `InvokeModel` response bodies aren't parsed, e.g. image generation results. -MAX_INVOKE_MODEL_RESPONSE_BYTES = 1_000_000 +# larger `InvokeModel` request and response bodies aren't parsed, e.g. requests with +# images or image generation results. +MAX_INVOKE_MODEL_BODY_BYTES = 1_000_000 class _BedrockRuntimeExtension(_ServiceExtension): @@ -252,11 +253,11 @@ def instrument_response( def on_read(data: bytes) -> None: nonlocal size - if model_family is None or size > MAX_INVOKE_MODEL_RESPONSE_BYTES: + if model_family is None or size > MAX_INVOKE_MODEL_BODY_BYTES: return size += len(data) - if size > MAX_INVOKE_MODEL_RESPONSE_BYTES: + if size > MAX_INVOKE_MODEL_BODY_BYTES: chunks.clear() else: chunks.append(data) @@ -362,10 +363,18 @@ def _get_invoke_model_request_attributes( A request body that is a file-like object isn't read, as that would consume it. https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html """ - if model_family is None or not isinstance(body, (bytes, bytearray, str)): + if ( + model_family is None + or not isinstance(body, (bytes, bytearray, str)) + or len(body) > MAX_INVOKE_MODEL_BODY_BYTES + ): return {} - request = json.loads(body) + try: + request = json.loads(body) + except ValueError: + # the request fails, as Bedrock rejects it. + return {} if not isinstance(request, Mapping): return {} diff --git a/tests/integrations/boto3/test_bedrock_runtime.py b/tests/integrations/boto3/test_bedrock_runtime.py index 29bb8bce94..7d4a92be25 100644 --- a/tests/integrations/boto3/test_bedrock_runtime.py +++ b/tests/integrations/boto3/test_bedrock_runtime.py @@ -1,3 +1,4 @@ +import base64 import copy import json from io import BytesIO @@ -759,6 +760,14 @@ def test_other_operations_keep_generic_span(init_bedrock, capture_items, trace_m assert not any(key.startswith("gen_ai.") for key in span["attributes"]) +# base64-encoded media content, which isn't recorded. +MEDIA_BASE64 = base64.b64encode(b"\x89PNG media bytes").decode() +ANTHROPIC_MEDIA_SOURCE = { + "type": "base64", + "media_type": "image/png", + "data": MEDIA_BASE64, +} + INVOKE_MODEL_CASES = [ pytest.param( "us.anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -871,6 +880,101 @@ def test_other_operations_keep_generic_span(init_bedrock, capture_items, trace_m }, id="anthropic_claude", ), + pytest.param( + "anthropic.claude-3-5-sonnet-20240620-v1:0", + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 100, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in these?"}, + {"type": "image", "source": ANTHROPIC_MEDIA_SOURCE}, + { + "type": "document", + "source": { + **ANTHROPIC_MEDIA_SOURCE, + "media_type": "application/pdf", + }, + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-1", + "name": "screenshot", + "input": {}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool-1", + "content": [ + {"type": "text", "text": "Done."}, + {"type": "image", "source": ANTHROPIC_MEDIA_SOURCE}, + ], + } + ], + }, + ], + }, + { + "id": "msg_2", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "A cat."}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["end_turn"], + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 10, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 5, + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + { + "role": "user", + "parts": [{"type": "text", "content": "What's in these?"}], + }, + { + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "id": "tool-1", + "name": "screenshot", + "arguments": {}, + }, + ], + }, + { + "role": "tool", + "parts": [ + { + "type": "tool_call_response", + "id": "tool-1", + "result": "Done.", + } + ], + }, + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + { + "role": "assistant", + "parts": [{"type": "text", "content": "A cat."}], + } + ], + }, + id="anthropic_claude_media", + ), pytest.param( "amazon.nova-lite-v1:0", { @@ -910,6 +1014,46 @@ def test_other_operations_keep_generic_span(init_bedrock, capture_items, trace_m }, id="amazon_nova", ), + pytest.param( + "amazon.nova-lite-v1:0", + { + "schemaVersion": "messages-v1", + "messages": [ + { + "role": "user", + "content": [ + {"text": "What's this?"}, + {"image": {"format": "png", "source": {"bytes": MEDIA_BASE64}}}, + ], + } + ], + "inferenceConfig": {"max_new_tokens": 100}, + }, + { + "output": { + "message": {"role": "assistant", "content": [{"text": "A cat."}]} + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + }, + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["end_turn"], + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 10, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 5, + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS: 15, + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + {"role": "user", "parts": [{"type": "text", "content": "What's this?"}]} + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + { + "role": "assistant", + "parts": [{"type": "text", "content": "A cat."}], + } + ], + }, + id="amazon_nova_media", + ), pytest.param( "amazon.titan-text-express-v1", { @@ -1153,6 +1297,8 @@ def call(): assert gen_ai_attributes == _expected_gen_ai_attributes( trace_mode, expected_attributes ) + # media content isn't recorded. + assert MEDIA_BASE64 not in json.dumps(span["attributes"]) # only the gen_ai span is recorded, not an `http.client.stream` span. assert not any( span["attributes"].get(SPANDATA.SENTRY_OP) == OP.HTTP_CLIENT_STREAM @@ -1326,11 +1472,28 @@ def call(): @pytest.mark.parametrize("trace_mode", TRACE_MODES) +@pytest.mark.parametrize( + "make_request_body", + [ + pytest.param( + lambda: BytesIO( + json.dumps({"prompt": "Hello!", "max_gen_len": 1}).encode() + ), + id="file_like", + ), + pytest.param( + lambda: json.dumps({"prompt": "x" * 1_000_001, "max_gen_len": 1}).encode(), + id="large", + ), + pytest.param(lambda: b'{"prompt": "Hello!", "max_gen_len": 1', id="malformed"), + ], +) def test_invoke_model_bodies_that_are_not_parsed( - init_bedrock, capture_items, trace_mode + init_bedrock, capture_items, trace_mode, make_request_body ): client = init_bedrock(trace_mode, send_default_pii=True) - # a file-like request body isn't read, and a large response body isn't parsed. + # a file-like request body isn't read, and large or malformed request bodies and + # a large response body aren't parsed. The response is unchanged. response_bytes = json.dumps( {"generation": "x" * 1_000_001, "stop_reason": "length"} ).encode() @@ -1338,8 +1501,7 @@ def test_invoke_model_bodies_that_are_not_parsed( def call(): response = client.invoke_model( - modelId="meta.llama3-8b-instruct-v1:0", - body=BytesIO(json.dumps({"prompt": "Hello!", "max_gen_len": 1}).encode()), + modelId="meta.llama3-8b-instruct-v1:0", body=make_request_body() ) bodies.append(response["body"].read())