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..302833ad72 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,94 @@ # `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` 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): """ @@ -90,34 +188,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_BODY_BYTES: + return + + size += len(data) + if size > MAX_INVOKE_MODEL_BODY_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 +336,115 @@ 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)) + or len(body) > MAX_INVOKE_MODEL_BODY_BYTES + ): + return {} + + try: + request = json.loads(body) + except ValueError: + # the request fails, as Bedrock rejects it. + return {} + + 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 +512,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 +691,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 +737,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 +758,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 +829,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 +893,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..7d4a92be25 100644 --- a/tests/integrations/boto3/test_bedrock_runtime.py +++ b/tests/integrations/boto3/test_bedrock_runtime.py @@ -1,8 +1,11 @@ +import base64 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 +14,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 +174,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 +758,761 @@ 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"]) + + +# 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", + { + "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( + "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", + { + "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.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", + { + "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 + ) + # 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 + 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) +@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, make_request_body +): + client = init_bedrock(trace_mode, send_default_pii=True) + # 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() + bodies = [] + + def call(): + response = client.invoke_model( + modelId="meta.llama3-8b-instruct-v1:0", body=make_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] + + 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