From 0dae13b60009185bf7f9bddce3780dea071298ba 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:49:37 +0000 Subject: [PATCH 1/2] feat(boto3): Instrument Bedrock Runtime streaming calls as gen_ai spans Record ConverseStream and InvokeModelWithResponseStream calls as gen_ai spans that finish once the caller has received the stream: when it's exhausted, closed, iteration stops early, or it fails. The finish reason, token usage, and output messages are accumulated from the events the caller receives, which are unchanged. `EventStream.__iter__` and `close` are patched globally, since the iterator is a method; only streams of instrumented calls are recorded. Refs #5470 Co-Authored-By: Sergiy Dybskiy <3760543+sergical@users.noreply.github.com> --- sentry_sdk/integrations/boto3/__init__.py | 2 + .../integrations/boto3/_instrumentation.py | 89 ++- .../boto3/_services/bedrock_runtime.py | 409 +++++++++- .../boto3/test_bedrock_runtime.py | 749 +++++++++++++++++- 4 files changed, 1208 insertions(+), 41 deletions(-) diff --git a/sentry_sdk/integrations/boto3/__init__.py b/sentry_sdk/integrations/boto3/__init__.py index 3688211be9..3d4a19bcc3 100644 --- a/sentry_sdk/integrations/boto3/__init__.py +++ b/sentry_sdk/integrations/boto3/__init__.py @@ -1,5 +1,6 @@ from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version from sentry_sdk.integrations.boto3._client import _patch_botocore_client +from sentry_sdk.integrations.boto3._instrumentation import _patch_event_stream from sentry_sdk.integrations.boto3.consts import IDENTIFIER, ORIGIN from sentry_sdk.utils import parse_version @@ -24,3 +25,4 @@ def setup_once() -> None: _check_minimum_version(Boto3Integration, version, "botocore") _patch_botocore_client() + _patch_event_stream() diff --git a/sentry_sdk/integrations/boto3/_instrumentation.py b/sentry_sdk/integrations/boto3/_instrumentation.py index d947b3e635..c1d6a84abf 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, Callable, Dict, Mapping, Optional, Union + from typing import Any, Callable, Dict, Iterator, Mapping, Optional, Union from sentry_sdk._types import Attributes from sentry_sdk.integrations.boto3._context import AwsCallContext @@ -33,6 +33,7 @@ try: from botocore.awsrequest import AWSRequest + from botocore.eventstream import EventStream from botocore.exceptions import ClientError from botocore.response import StreamingBody except ImportError: @@ -391,6 +392,92 @@ def sentry_raw_stream_close(*args: "Any", **kwargs: "Any") -> None: raise +class _EventStreamHooks: + """ + Callbacks of an instrumented `EventStream`, stored on the stream. + """ + + __slots__ = ("on_event", "on_finish", "finished") + + def __init__( + self, + on_event: "Callable[[Any], None]", + on_finish: "Callable[[Optional[BaseException]], None]", + ) -> None: + self.on_event = on_event + self.on_finish = on_finish + self.finished = False + + def finish(self, error: "Optional[BaseException]" = None) -> None: + # only finish once across iteration and close. + if self.finished: + return + + self.finished = True + with capture_internal_exceptions(): + self.on_finish(error) + + +def _instrument_event_stream( + stream: "EventStream[Any]", + on_event: "Callable[[Any], None]", + on_finish: "Callable[[Optional[BaseException]], None]", +) -> None: + """ + Call `on_event` with each event the caller receives from `stream`, and + `on_finish` once, when `stream` is exhausted, closed, stops being iterated, + or fails. + """ + stream._sentry_hooks = _EventStreamHooks(on_event, on_finish) # type: ignore[attr-defined] + + +def _patch_event_stream() -> None: + """ + Patch `EventStream` globally, since its iterator is a method; streams without + hooks are unchanged. + """ + orig_iter = EventStream.__iter__ + orig_close = EventStream.close + + def sentry_event_stream_iter(self: "EventStream[Any]") -> "Iterator[Any]": + hooks: "Optional[_EventStreamHooks]" = getattr(self, "_sentry_hooks", None) + if hooks is None: + yield from orig_iter(self) + return + + try: + for event in orig_iter(self): + with capture_internal_exceptions(): + hooks.on_event(event) + yield event + except GeneratorExit: + # the caller stopped iterating, e.g. with `break`. + hooks.finish() + raise + except BaseException as error: + hooks.finish(error) + raise + + hooks.finish() + + def sentry_event_stream_close( + self: "EventStream[Any]", *args: "Any", **kwargs: "Any" + ) -> None: + hooks: "Optional[_EventStreamHooks]" = getattr(self, "_sentry_hooks", None) + try: + orig_close(self, *args, **kwargs) + except BaseException as error: + if hooks is not None: + hooks.finish(error) + raise + + if hooks is not None: + hooks.finish() + + EventStream.__iter__ = sentry_event_stream_iter # type: ignore[method-assign] + EventStream.close = sentry_event_stream_close # type: ignore[method-assign] + + def _set_request_attributes( span: "Union[Span, StreamedSpan]", request: "AWSRequest", diff --git a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py index 302833ad72..8533a74e2f 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py @@ -12,6 +12,7 @@ from sentry_sdk.integrations.boto3._instrumentation import ( _finish_span, _get_error_attributes, + _instrument_event_stream, _patch_streaming_body, _set_span_attributes, ) @@ -43,6 +44,7 @@ ContentPart = Union[TextPart, ReasoningPart, ToolCallPart, ToolCallResponsePart] try: + from botocore.eventstream import EventStream from botocore.response import StreamingBody except ImportError: raise DidNotEnable("botocore not installed") @@ -55,9 +57,16 @@ # `gen_ai.operation.name` and span op of each instrumented operation. GEN_AI_OPERATIONS = { "Converse": ("chat", OP.GEN_AI_CHAT), + "ConverseStream": ("chat", OP.GEN_AI_CHAT), "InvokeModel": ("generate_content", OP.GEN_AI_GENERATE_CONTENT), + "InvokeModelWithResponseStream": ("generate_content", OP.GEN_AI_GENERATE_CONTENT), } +# operations with a model-specific request body; the others use the Converse format. +INVOKE_MODEL_OPERATIONS = ("InvokeModel", "InvokeModelWithResponseStream") + +STREAMING_OPERATIONS = ("ConverseStream", "InvokeModelWithResponseStream") + # `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. @@ -181,14 +190,15 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": attributes: "Attributes" = { SPANDATA.GEN_AI_PROVIDER_NAME: PROVIDER_NAME, SPANDATA.GEN_AI_OPERATION_NAME: operation_name, - SPANDATA.GEN_AI_RESPONSE_STREAMING: False, + SPANDATA.GEN_AI_RESPONSE_STREAMING: ctx.operation_name + in STREAMING_OPERATIONS, } model_id = ctx.params.get("modelId") if model_id: attributes[SPANDATA.GEN_AI_REQUEST_MODEL] = model_id - if ctx.operation_name != "InvokeModel": + if ctx.operation_name not in INVOKE_MODEL_OPERATIONS: with capture_internal_exceptions(): attributes.update(_get_converse_request_attributes(ctx.params)) @@ -204,7 +214,7 @@ def get_span_start_attributes( return {} # the request body is parsed once, for its parameters and input messages. - if ctx.operation_name == "InvokeModel": + if ctx.operation_name in INVOKE_MODEL_OPERATIONS: return _get_invoke_model_request_attributes( span, _get_model_family(ctx.params.get("modelId")), @@ -236,51 +246,107 @@ def instrument_response( 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. + The `InvokeModel` response body and the streamed responses are read by the + caller. They're recorded as the caller reads them, and the span finishes once + they're consumed or closed. """ - if ctx.operation_name != "InvokeModel" or isinstance(span, NoOpStreamedSpan): + if isinstance(span, NoOpStreamedSpan) or not isinstance(response, Mapping): return False - body = response.get("body") if isinstance(response, Mapping) else None - if not isinstance(body, StreamingBody): + model_family = _get_model_family(ctx.params.get("modelId")) + + if ctx.operation_name == "InvokeModel": + body = response.get("body") + if not isinstance(body, StreamingBody): + return False + + _instrument_invoke_model_body(span, model_family, body, response) + return True + + if ctx.operation_name == "ConverseStream": + stream = response.get("stream") + elif ctx.operation_name == "InvokeModelWithResponseStream": + stream = response.get("body") + else: 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) + if not isinstance(stream, EventStream): + return False + + streamed_response = _StreamedResponse() - def on_finish(error: "Optional[BaseException]") -> None: - attributes: "Attributes" = {} - if error is not None: - with capture_internal_exceptions(): - attributes.update(_get_error_attributes(error)) + def on_event(event: "Any") -> None: + if ctx.operation_name == "ConverseStream": + _record_converse_stream_event(streamed_response, event) else: - with capture_internal_exceptions(): - attributes.update( - _get_invoke_model_response_attributes( - model_family, b"".join(chunks), headers - ) - ) + _record_invoke_model_stream_event( + streamed_response, model_family, event + ) + + _instrument_event_stream( + stream, + on_event, + lambda error: _finish_gen_ai_span( + span, error, streamed_response.get_attributes + ), + ) + return True - with capture_internal_exceptions(): - _set_span_attributes(span, attributes) - _finish_span(span, error) - _patch_streaming_body(body, on_finish, on_read) - return True +def _instrument_invoke_model_body( + span: "Union[Span, StreamedSpan]", + model_family: "Optional[str]", + body: "StreamingBody", + response: "Mapping[str, Any]", +) -> None: + 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) + + _patch_streaming_body( + body, + lambda error: _finish_gen_ai_span( + span, + error, + lambda: _get_invoke_model_response_attributes( + model_family, b"".join(chunks), headers + ), + ), + on_read, + ) + + +def _finish_gen_ai_span( + span: "Union[Span, StreamedSpan]", + error: "Optional[BaseException]", + get_response_attributes: "Callable[[], Attributes]", +) -> None: + """ + Finish a span once its response is read, with the attributes of the response + that was read. + """ + attributes: "Attributes" = {} + with capture_internal_exceptions(): + attributes.update(get_response_attributes()) + + if error is not None: + with capture_internal_exceptions(): + attributes.update(_get_error_attributes(error)) + + with capture_internal_exceptions(): + _set_span_attributes(span, attributes) + _finish_span(span, error) def _get_model_family(model_id: "Any") -> "Optional[str]": @@ -655,6 +721,271 @@ def _parse_invoke_model_response( return attributes +class _StreamedResponse: + """ + A `ConverseStream` or `InvokeModelWithResponseStream` response, accumulated from + the events the caller receives. + """ + + __slots__ = ("record_outputs", "finish_reason", "usage", "content_blocks") + + def __init__(self) -> None: + # only buffer the response content if it's recorded. + self.record_outputs = _should_record_gen_ai("outputs") + self.finish_reason: "Optional[str]" = None + # arguments of `_get_usage_attributes()`. + self.usage: "dict[str, Any]" = {} + # the part type, streamed fragments, and tool call ID and name of each content + # block, by index. + self.content_blocks: "dict[Any, dict[str, Any]]" = {} + + def set_finish_reason(self, finish_reason: "Any") -> None: + if isinstance(finish_reason, str): + self.finish_reason = finish_reason + + def set_usage(self, overwrite: bool = True, **usage: "Any") -> None: + for key, value in usage.items(): + if isinstance(value, int) and (overwrite or key not in self.usage): + self.usage[key] = value + + def start_tool_call(self, index: "Any", tool_call_id: "Any", name: "Any") -> None: + if self.record_outputs: + self.content_blocks[index] = { + "type": "tool_call", + "fragments": [], + "id": tool_call_id, + "name": name, + } + + def add_fragment( + self, + index: "Any", + part_type: 'Literal["text", "reasoning", "tool_call"]', + fragment: "Any", + ) -> None: + """ + Add streamed text, reasoning, or tool call arguments to a content block. + """ + if not self.record_outputs or not isinstance(fragment, str): + return + + content_block = self.content_blocks.setdefault( + index, {"type": part_type, "fragments": []} + ) + content_block["fragments"].append(fragment) + + def get_attributes(self) -> "Attributes": + attributes: "Attributes" = {} + if self.finish_reason is not None: + attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] = [self.finish_reason] + + attributes.update( + _get_usage_attributes( + self.usage.get("input_tokens"), + self.usage.get("output_tokens"), + self.usage.get("cache_read_input_tokens"), + self.usage.get("cache_write_input_tokens"), + self.usage.get("total_tokens"), + ) + ) + + parts: "list[ContentPart]" = [] + # content blocks start in order. + for content_block in self.content_blocks.values(): + content = "".join(content_block["fragments"]) + if content_block["type"] == "text": + parts.append({"type": "text", "content": content}) + elif content_block["type"] == "reasoning": + parts.append({"type": "reasoning", "content": content}) + else: + tool_call: "ToolCallPart" = {"type": "tool_call"} + if content_block.get("id") is not None: + tool_call["id"] = content_block["id"] + if content_block.get("name") is not None: + tool_call["name"] = content_block["name"] + if content: + try: + tool_call["arguments"] = json.loads(content) + except ValueError: + # e.g. a stream that stopped before the arguments were complete. + tool_call["arguments"] = content + parts.append(tool_call) + + if parts: + attributes.update(_get_output_attributes([_to_output_message(parts)])) + + return attributes + + +def _record_converse_stream_event( + response: "_StreamedResponse", event: "Mapping[str, Any]" +) -> None: + """ + Record a `ConverseStream` event, or a streamed Amazon Nova `InvokeModel` chunk, + which has the same format. + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStreamOutput.html + """ + content_block_start = event.get("contentBlockStart") + if isinstance(content_block_start, Mapping): + tool_use = (content_block_start.get("start") or {}).get("toolUse") + if isinstance(tool_use, Mapping): + response.start_tool_call( + content_block_start.get("contentBlockIndex"), + tool_use.get("toolUseId"), + tool_use.get("name"), + ) + + content_block_delta = event.get("contentBlockDelta") + if isinstance(content_block_delta, Mapping): + index = content_block_delta.get("contentBlockIndex") + delta = content_block_delta.get("delta") or {} + if "text" in delta: + response.add_fragment(index, "text", delta["text"]) + elif isinstance(delta.get("toolUse"), Mapping): + response.add_fragment(index, "tool_call", delta["toolUse"].get("input")) + elif isinstance(delta.get("reasoningContent"), Mapping): + response.add_fragment( + index, "reasoning", delta["reasoningContent"].get("text") + ) + + message_stop = event.get("messageStop") + if isinstance(message_stop, Mapping): + response.set_finish_reason(message_stop.get("stopReason")) + + metadata = event.get("metadata") + usage = metadata.get("usage") if isinstance(metadata, Mapping) else None + if isinstance(usage, Mapping): + response.set_usage( + input_tokens=usage.get("inputTokens"), + output_tokens=usage.get("outputTokens"), + total_tokens=usage.get("totalTokens"), + # Converse and Nova name the prompt cache usage differently. + cache_read_input_tokens=usage.get( + "cacheReadInputTokens", usage.get("cacheReadInputTokenCount") + ), + cache_write_input_tokens=usage.get( + "cacheWriteInputTokens", usage.get("cacheWriteInputTokenCount") + ), + ) + + +def _record_invoke_model_stream_event( + response: "_StreamedResponse", + model_family: "Optional[str]", + event: "Mapping[str, Any]", +) -> None: + """ + Record an `InvokeModelWithResponseStream` event. Its chunk has the format of the + model family's streamed response. + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html + """ + payload = (event.get("chunk") or {}).get("bytes") + if not payload: + return + + chunk = json.loads(payload) + if not isinstance(chunk, Mapping): + return + + if model_family == "amazon.nova": + _record_converse_stream_event(response, chunk) + elif model_family == "anthropic.claude": + _record_anthropic_stream_event(response, chunk) + elif model_family is not None: + _record_text_stream_chunk(response, model_family, chunk) + + # Bedrock adds the token usage of every model to the last chunk. + metrics = chunk.get("amazon-bedrock-invocationMetrics") + if isinstance(metrics, Mapping): + response.set_usage( + overwrite=False, + input_tokens=metrics.get("inputTokenCount"), + output_tokens=metrics.get("outputTokenCount"), + ) + + +def _record_anthropic_stream_event( + response: "_StreamedResponse", event: "Mapping[str, Any]" +) -> None: + """ + https://docs.anthropic.com/en/docs/build-with-claude/streaming + """ + event_type = event.get("type") + if event_type == "message_start": + usage = (event.get("message") or {}).get("usage") or {} + # the output tokens are reported by `message_delta`. + response.set_usage( + input_tokens=usage.get("input_tokens"), + cache_read_input_tokens=usage.get("cache_read_input_tokens"), + cache_write_input_tokens=usage.get("cache_creation_input_tokens"), + ) + + elif event_type == "content_block_start": + content_block = event.get("content_block") or {} + if content_block.get("type") == "tool_use": + response.start_tool_call( + event.get("index"), content_block.get("id"), content_block.get("name") + ) + + elif event_type == "content_block_delta": + index = event.get("index") + delta = event.get("delta") or {} + delta_type = delta.get("type") + if delta_type == "text_delta": + response.add_fragment(index, "text", delta.get("text")) + elif delta_type == "thinking_delta": + response.add_fragment(index, "reasoning", delta.get("thinking")) + elif delta_type == "input_json_delta": + response.add_fragment(index, "tool_call", delta.get("partial_json")) + + elif event_type == "message_delta": + response.set_finish_reason((event.get("delta") or {}).get("stop_reason")) + response.set_usage( + output_tokens=(event.get("usage") or {}).get("output_tokens") + ) + + +def _record_text_stream_chunk( + response: "_StreamedResponse", model_family: str, chunk: "Mapping[str, Any]" +) -> None: + """ + Record a streamed chunk of a model family that generates text from a prompt. + """ + text = finish_reason = None + + if model_family == "amazon.titan": + text = chunk.get("outputText") + finish_reason = chunk.get("completionReason") + response.set_usage( + input_tokens=chunk.get("inputTextTokenCount"), + output_tokens=chunk.get("totalOutputTextTokenCount"), + ) + + elif model_family == "meta.llama": + text = chunk.get("generation") + finish_reason = chunk.get("stop_reason") + response.set_usage( + input_tokens=chunk.get("prompt_token_count"), + output_tokens=chunk.get("generation_token_count"), + ) + + elif model_family in ("cohere.command-r", "cohere.command"): + # Cohere Command chunks may also have the format of its response. + generations = chunk.get("generations") + generation = generations[0] if generations else chunk + text = generation.get("text") + finish_reason = generation.get("finish_reason") + + elif model_family == "mistral": + outputs = chunk.get("outputs") + if outputs: + text = outputs[0].get("text") + finish_reason = outputs[0].get("stop_reason") + + response.add_fragment(0, "text", text) + response.set_finish_reason(finish_reason) + + def _get_usage_attributes( input_tokens: "Any", output_tokens: "Any", diff --git a/tests/integrations/boto3/test_bedrock_runtime.py b/tests/integrations/boto3/test_bedrock_runtime.py index 7d4a92be25..d8b342106a 100644 --- a/tests/integrations/boto3/test_bedrock_runtime.py +++ b/tests/integrations/boto3/test_bedrock_runtime.py @@ -1,12 +1,14 @@ import base64 +import binascii import copy import json +import struct from io import BytesIO import boto3 import pytest from botocore.awsrequest import AWSResponse -from botocore.exceptions import ClientError +from botocore.exceptions import ClientError, EventStreamError from botocore.stub import Stubber import sentry_sdk @@ -1516,3 +1518,748 @@ def call(): 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 + + +def _encode_event_stream_message(headers, payload): + """ + Encode a message in the `application/vnd.amazon.eventstream` format, with + string headers. + https://docs.aws.amazon.com/transcribe/latest/dg/streaming-setting-up.html#streaming-event-stream + """ + encoded_headers = b"".join( + struct.pack("B", len(name)) + + name.encode() + + b"\x07" + + struct.pack(">H", len(value)) + + value.encode() + for name, value in headers.items() + ) + prelude = struct.pack( + ">II", 16 + len(encoded_headers) + len(payload), len(encoded_headers) + ) + message = ( + prelude + struct.pack(">I", binascii.crc32(prelude)) + encoded_headers + payload + ) + return message + struct.pack(">I", binascii.crc32(message)) + + +def _event_stream_response(client, events): + """ + Respond with an event stream of `(event type, payload)` events. An event type of + `None` is a `modelStreamErrorException`. + """ + body = b"".join( + _encode_event_stream_message( + { + ":message-type": "event", + ":event-type": event_type, + ":content-type": "application/json", + } + if event_type is not None + else { + ":message-type": "exception", + ":exception-type": "modelStreamErrorException", + ":content-type": "application/json", + }, + json.dumps(payload).encode(), + ) + for event_type, payload in events + ) + return MockResponse( + client, + 200, + { + "content-type": "application/vnd.amazon.eventstream", + "x-amzn-requestid": "request-id", + }, + body, + ) + + +def _invoke_model_stream_chunks(*chunks): + return [ + ("chunk", {"bytes": base64.b64encode(json.dumps(chunk).encode()).decode()}) + for chunk in chunks + ] + + +CONVERSE_STREAM_EVENTS = [ + ("messageStart", {"role": "assistant"}), + ("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "Let me "}}), + ("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "check."}}), + ("contentBlockStop", {"contentBlockIndex": 0}), + ( + "contentBlockStart", + { + "contentBlockIndex": 1, + "start": {"toolUse": {"toolUseId": "tool-1", "name": "get_weather"}}, + }, + ), + ( + "contentBlockDelta", + {"contentBlockIndex": 1, "delta": {"toolUse": {"input": '{"city": '}}}, + ), + ( + "contentBlockDelta", + {"contentBlockIndex": 1, "delta": {"toolUse": {"input": '"Paris"}'}}}, + ), + ("contentBlockStop", {"contentBlockIndex": 1}), + ("messageStop", {"stopReason": "tool_use"}), + ( + "metadata", + { + "usage": {"inputTokens": 12, "outputTokens": 8, "totalTokens": 20}, + "metrics": {"latencyMs": 100}, + }, + ), +] + +CONVERSE_STREAM_PARAMS = { + "modelId": MODEL_ID, + "messages": [{"role": "user", "content": [{"text": "Weather in Paris?"}]}], + "inferenceConfig": {"maxTokens": 100}, +} + + +@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_converse_stream( + 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) + received_events = [] + + def call(): + response = client.converse_stream(**CONVERSE_STREAM_PARAMS) + received_events.extend(response["stream"]) + + with _event_stream_response(client, CONVERSE_STREAM_EVENTS): + spans = _capture_spans(capture_items, trace_mode, call) + + # the caller receives the events unchanged. + assert received_events == [ + {event_type: payload} for event_type, payload in CONVERSE_STREAM_EVENTS + ] + + span = _get_bedrock_span(spans) + attributes = _get_attributes(span) + assert span["name"] == f"chat {MODEL_ID}" + assert span["status"] == "ok" + assert attributes[SPANDATA.SENTRY_OP] == OP.GEN_AI_CHAT + assert attributes[SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert attributes[SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert attributes[SPANDATA.GEN_AI_REQUEST_MODEL] == MODEL_ID + assert attributes[SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + assert attributes[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == ["tool_use"] + assert attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 12 + assert attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 8 + assert attributes[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 20 + assert attributes[SPANDATA.RPC_METHOD] == "ConverseStream" + assert attributes[SPANDATA.AWS_REQUEST_ID] == "request-id" + + if record_inputs: + assert attributes[SPANDATA.GEN_AI_INPUT_MESSAGES] == [ + { + "role": "user", + "parts": [{"type": "text", "content": "Weather in Paris?"}], + } + ] + else: + assert SPANDATA.GEN_AI_INPUT_MESSAGES not in attributes + + if record_outputs: + assert attributes[SPANDATA.GEN_AI_OUTPUT_MESSAGES] == [ + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Let me check."}, + { + "type": "tool_call", + "id": "tool-1", + "name": "get_weather", + "arguments": {"city": "Paris"}, + }, + ], + } + ] + else: + assert SPANDATA.GEN_AI_OUTPUT_MESSAGES not in attributes + + +def _consume_until_error(stream): + events = [] + with pytest.raises(EventStreamError) as exc_info: + for event in stream: + events.append(event) + assert exc_info.value.response["Error"]["Code"] == "modelStreamErrorException" + return events + + +def _break_after_first_event(stream): + events = [] + for event in stream: + events.append(event) + break + return events + + +def _close_unread(stream): + stream.close() + return [] + + +def _close_after_first_event(stream): + events = [next(iter(stream))] + stream.close() + return events + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +@pytest.mark.parametrize( + "events, consume, received, status", + [ + pytest.param(CONVERSE_STREAM_EVENTS, list, 10, "ok", id="exhausted"), + pytest.param( + CONVERSE_STREAM_EVENTS, _break_after_first_event, 1, "ok", id="break" + ), + pytest.param( + CONVERSE_STREAM_EVENTS, _close_after_first_event, 1, "ok", id="close" + ), + pytest.param( + CONVERSE_STREAM_EVENTS, + _close_unread, + 0, + "ok", + id="close_unread", + ), + pytest.param( + CONVERSE_STREAM_EVENTS[:2] + [(None, {"message": "Model error"})], + _consume_until_error, + 2, + "error", + id="error", + ), + ], +) +def test_converse_stream_consumption( + init_bedrock, capture_items, trace_mode, events, consume, received, status +): + client = init_bedrock(trace_mode, send_default_pii=True) + received_events = [] + + def call(): + response = client.converse_stream(**CONVERSE_STREAM_PARAMS) + received_events.extend(consume(response["stream"])) + + with _event_stream_response(client, events): + spans = _capture_spans(capture_items, trace_mode, call) + + assert received_events == [ + {event_type: payload} for event_type, payload in events[:received] + ] + + # the span finishes once, with the response received so far. + span = _get_bedrock_span(spans) + attributes = _get_attributes(span) + assert span["status"] == status + if status == "error": + assert attributes[SPANDATA.ERROR_TYPE] == "modelStreamErrorException" + else: + assert SPANDATA.ERROR_TYPE not in attributes + + if received == len(CONVERSE_STREAM_EVENTS): + assert attributes[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == ["tool_use"] + else: + assert SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS not in attributes + assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in attributes + + if received >= 2: + assert attributes[SPANDATA.GEN_AI_OUTPUT_MESSAGES][0]["parts"][0] == { + "type": "text", + "content": "Let me check." if received > 2 else "Let me ", + } + else: + assert SPANDATA.GEN_AI_OUTPUT_MESSAGES not in attributes + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_converse_stream_without_parent_span(init_bedrock, capture_items, trace_mode): + client = init_bedrock(trace_mode, send_default_pii=True) + items = capture_items("transaction", "span") + + with _event_stream_response(client, CONVERSE_STREAM_EVENTS): + response = client.converse_stream(**CONVERSE_STREAM_PARAMS) + received_events = list(response["stream"]) + sentry_sdk.flush() + + assert received_events == [ + {event_type: payload} for event_type, payload in CONVERSE_STREAM_EVENTS + ] + assert not items + + +INVOKE_MODEL_STREAM_CASES = [ + pytest.param( + "us.anthropic.claude-3-5-sonnet-20240620-v1:0", + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Weather in Paris?"}], + }, + [ + { + "type": "message_start", + "message": { + "role": "assistant", + "content": [], + "usage": { + "input_tokens": 10, + "output_tokens": 1, + "cache_read_input_tokens": 100, + "cache_creation_input_tokens": 0, + }, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "thinking", "thinking": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": "I need the tool."}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "abc"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": "Let me "}, + }, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": "check."}, + }, + {"type": "content_block_stop", "index": 1}, + { + "type": "content_block_start", + "index": 2, + "content_block": { + "type": "tool_use", + "id": "tool-1", + "name": "get_weather", + "input": {}, + }, + }, + { + "type": "content_block_delta", + "index": 2, + "delta": {"type": "input_json_delta", "partial_json": '{"city": '}, + }, + { + "type": "content_block_delta", + "index": 2, + "delta": {"type": "input_json_delta", "partial_json": '"Paris"}'}, + }, + {"type": "content_block_stop", "index": 2}, + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use"}, + "usage": {"output_tokens": 25}, + }, + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 10, + "outputTokenCount": 25, + }, + }, + ], + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["tool_use"], + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 110, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 25, + SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS: 100, + SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS: 0, + SPANDATA.GEN_AI_INPUT_MESSAGES: [ + { + "role": "user", + "parts": [{"type": "text", "content": "Weather in Paris?"}], + } + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "I need the tool."}, + {"type": "text", "content": "Let me check."}, + { + "type": "tool_call", + "id": "tool-1", + "name": "get_weather", + "arguments": {"city": "Paris"}, + }, + ], + } + ], + }, + id="anthropic_claude", + ), + pytest.param( + "amazon.nova-lite-v1:0", + { + "messages": [{"role": "user", "content": [{"text": "Hello!"}]}], + "inferenceConfig": {"max_new_tokens": 100}, + }, + [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockDelta": {"delta": {"text": "Hi"}, "contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"text": "!"}, "contentBlockIndex": 0}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + { + "metadata": {"usage": {"inputTokens": 10, "outputTokens": 5}}, + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 10, + "outputTokenCount": 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": "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}}, + [ + { + "outputText": "Hi", + "index": 0, + "totalOutputTextTokenCount": 1, + "completionReason": None, + "inputTextTokenCount": 10, + }, + { + "outputText": "!", + "index": 0, + "totalOutputTextTokenCount": 5, + "completionReason": "FINISH", + "inputTextTokenCount": None, + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 10, + "outputTokenCount": 5, + }, + }, + ], + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + 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}, + [ + { + "generation": "Hi", + "prompt_token_count": 10, + "generation_token_count": 1, + "stop_reason": None, + }, + { + "generation": "!", + "prompt_token_count": None, + "generation_token_count": 5, + "stop_reason": "stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 10, + "outputTokenCount": 5, + }, + }, + ], + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + 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!", "max_tokens": 100}, + [ + {"is_finished": False, "event_type": "stream-start", "generation_id": "1"}, + {"is_finished": False, "event_type": "text-generation", "text": "Hi"}, + {"is_finished": False, "event_type": "text-generation", "text": "!"}, + { + "is_finished": True, + "event_type": "stream-end", + "finish_reason": "COMPLETE", + "response": {"text": "Hi!"}, + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 7, + "outputTokenCount": 3, + }, + }, + ], + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS: ["COMPLETE"], + # the chunks have no usage; it's in Bedrock's invocation metrics. + 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_r", + ), + pytest.param( + "cohere.command-text-v14", + {"prompt": "Hello!", "max_tokens": 100}, + [ + {"generations": [{"text": "Hi"}]}, + { + "generations": [{"text": "!", "finish_reason": "COMPLETE"}], + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 7, + "outputTokenCount": 3, + }, + }, + ], + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + 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": "Hello!", "max_tokens": 100}, + [ + {"outputs": [{"text": "Hi", "stop_reason": None}]}, + { + "outputs": [{"text": "!", "stop_reason": "stop"}], + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 7, + "outputTokenCount": 3, + }, + }, + ], + { + SPANDATA.GEN_AI_REQUEST_MAX_TOKENS: 100, + 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": "Hello!"}]} + ], + SPANDATA.GEN_AI_OUTPUT_MESSAGES: [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ], + }, + id="mistral", + ), + pytest.param( + "ai21.jamba-1-5-mini-v1:0", + {"messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 100}, + [ + {"choices": [{"delta": {"content": "Hi!"}}]}, + { + "choices": [{"finish_reason": "stop"}], + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 7, + "outputTokenCount": 3, + }, + }, + ], + { + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS: 7, + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: 3, + }, + id="unknown_model", + ), +] + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +@pytest.mark.parametrize( + "model_id, request_body, chunks, expected_attributes", INVOKE_MODEL_STREAM_CASES +) +def test_invoke_model_with_response_stream( + init_bedrock, + capture_items, + trace_mode, + model_id, + request_body, + chunks, + expected_attributes, +): + client = init_bedrock(trace_mode, send_default_pii=True) + received_chunks = [] + + def call(): + response = client.invoke_model_with_response_stream( + modelId=model_id, body=json.dumps(request_body) + ) + received_chunks.extend( + json.loads(event["chunk"]["bytes"]) for event in response["body"] + ) + + with _event_stream_response(client, _invoke_model_stream_chunks(*chunks)): + spans = _capture_spans(capture_items, trace_mode, call) + + assert received_chunks == chunks + + 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.RPC_METHOD] == "InvokeModelWithResponseStream" + assert { + 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: True, **expected_attributes} + 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 + + +@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_with_response_stream_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, chunks, expected_attributes = INVOKE_MODEL_STREAM_CASES[ + 0 + ].values + + def call(): + response = client.invoke_model_with_response_stream( + modelId=model_id, body=json.dumps(request_body) + ) + for _ in response["body"]: + pass + + with _event_stream_response(client, _invoke_model_stream_chunks(*chunks)): + spans = _capture_spans(capture_items, trace_mode, call) + + attributes = _get_attributes(_get_bedrock_span(spans)) + for attribute, recorded in ( + (SPANDATA.GEN_AI_INPUT_MESSAGES, record_inputs), + (SPANDATA.GEN_AI_OUTPUT_MESSAGES, record_outputs), + ): + if recorded: + assert attributes[attribute] == expected_attributes[attribute] + else: + assert attribute not in attributes + assert ( + attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] + == expected_attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] + ) + + +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_invoke_model_with_response_stream_error_is_unchanged( + init_bedrock, capture_items, trace_mode +): + client = init_bedrock(trace_mode) + events = _invoke_model_stream_chunks( + {"outputs": [{"text": "Hi", "stop_reason": None}]} + ) + [(None, {"message": "Model error"})] + received_chunks = [] + + def call(): + response = client.invoke_model_with_response_stream( + modelId="mistral.mistral-7b-instruct-v0:2", + body=json.dumps({"prompt": "Hello!"}), + ) + received_chunks.extend(_consume_until_error(response["body"])) + + with _event_stream_response(client, events): + spans = _capture_spans(capture_items, trace_mode, call) + + assert len(received_chunks) == 1 + + span = _get_bedrock_span(spans) + assert span["status"] == "error" + assert span["attributes"][SPANDATA.ERROR_TYPE] == "modelStreamErrorException" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True From 81e76ade572a3de1846c898961ca613448b55907 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:49:46 +0000 Subject: [PATCH 2/2] fix(boto3): Skip malformed Bedrock stream chunks and gate stream outputs once Malformed InvokeModelWithResponseStream chunks are skipped quietly instead of reporting an SDK internal error, like the non-streamed response path. The buffered stream content is already gated on recording outputs, so serialize it directly rather than checking the gate a second time. That makes the buffering gate the only gate, so tests cover it. Co-Authored-By: Sergiy Dybskiy <3760543+sergical@users.noreply.github.com> --- .../boto3/_services/bedrock_runtime.py | 12 +++++- .../boto3/test_bedrock_runtime.py | 37 +++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py index 8533a74e2f..b3e2f40b79 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock_runtime.py @@ -811,8 +811,12 @@ def get_attributes(self) -> "Attributes": tool_call["arguments"] = content parts.append(tool_call) + # `record_outputs` already gated the buffered content blocks. if parts: - attributes.update(_get_output_attributes([_to_output_message(parts)])) + with capture_internal_exceptions(): + attributes[SPANDATA.GEN_AI_OUTPUT_MESSAGES] = json.dumps( + [_to_output_message(parts)] + ) return attributes @@ -883,7 +887,11 @@ def _record_invoke_model_stream_event( if not payload: return - chunk = json.loads(payload) + try: + chunk = json.loads(payload) + except ValueError: + return + if not isinstance(chunk, Mapping): return diff --git a/tests/integrations/boto3/test_bedrock_runtime.py b/tests/integrations/boto3/test_bedrock_runtime.py index d8b342106a..6161b14c11 100644 --- a/tests/integrations/boto3/test_bedrock_runtime.py +++ b/tests/integrations/boto3/test_bedrock_runtime.py @@ -2237,6 +2237,43 @@ def call(): ) +@pytest.mark.parametrize("trace_mode", TRACE_MODES) +def test_invoke_model_with_response_stream_skips_malformed_chunks( + init_bedrock, capture_items, trace_mode +): + client = init_bedrock(trace_mode, send_default_pii=True) + malformed_chunk = ("chunk", {"bytes": base64.b64encode(b"not json").decode()}) + events = [malformed_chunk] + _invoke_model_stream_chunks( + { + "outputs": [{"text": "Hi!", "stop_reason": "stop"}], + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 7, + "outputTokenCount": 3, + }, + } + ) + received_chunks = [] + + def call(): + response = client.invoke_model_with_response_stream( + modelId="mistral.mistral-7b-instruct-v0:2", + body=json.dumps({"prompt": "Hello!"}), + ) + received_chunks.extend(event["chunk"]["bytes"] for event in response["body"]) + + with _event_stream_response(client, events): + spans = _capture_spans(capture_items, trace_mode, call) + + assert received_chunks[0] == b"not json" + assert len(received_chunks) == 2 + + attributes = _get_attributes(_get_bedrock_span(spans)) + assert attributes[SPANDATA.GEN_AI_OUTPUT_MESSAGES] == [ + {"role": "assistant", "parts": [{"type": "text", "content": "Hi!"}]} + ] + assert attributes[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 3 + + @pytest.mark.parametrize("trace_mode", TRACE_MODES) def test_invoke_model_with_response_stream_error_is_unchanged( init_bedrock, capture_items, trace_mode