diff --git a/sentry_sdk/integrations/boto3/__init__.py b/sentry_sdk/integrations/boto3/__init__.py index d36b44d271..d44d864c48 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._services.bedrock import _patch_event_stream from sentry_sdk.integrations.boto3.consts import IDENTIFIER, ORIGIN from sentry_sdk.utils import parse_version @@ -19,3 +20,4 @@ def setup_once() -> None: _check_minimum_version(Boto3Integration, version, "botocore") _patch_botocore_client() + _patch_event_stream() diff --git a/sentry_sdk/integrations/boto3/_client.py b/sentry_sdk/integrations/boto3/_client.py index 280a570d45..e58b6bcded 100644 --- a/sentry_sdk/integrations/boto3/_client.py +++ b/sentry_sdk/integrations/boto3/_client.py @@ -27,7 +27,6 @@ try: from botocore.client import BaseClient - from botocore.exceptions import ClientError except ImportError: raise DidNotEnable("botocore not installed") @@ -121,21 +120,11 @@ def sentry_patched_make_api_call( try: parsed = orig_make_api_call(self, operation_name, api_params) except BaseException as error: - if service_ext is not None and isinstance(error, ClientError): - with capture_internal_exceptions(): - attributes.update( - service_ext.get_response_attributes(ctx, error.response) - ) with capture_internal_exceptions(): attributes.update(_get_error_attributes(error)) with capture_internal_exceptions(): _set_span_attributes(span, attributes) raise - if service_ext is not None: - with capture_internal_exceptions(): - attributes.update( - service_ext.get_response_attributes(ctx, parsed) - ) with capture_internal_exceptions(): attributes.update(_get_response_attributes(parsed)) with capture_internal_exceptions(): @@ -144,6 +133,10 @@ def sentry_patched_make_api_call( _finish_span(span, error) raise + if service_ext is not None: + service_ext.on_success_response(ctx, parsed, span) + return parsed + streaming_body_instrumented = False with capture_internal_exceptions(): streaming_body_instrumented = _instrument_streaming_body(span, parsed) diff --git a/sentry_sdk/integrations/boto3/_services/base.py b/sentry_sdk/integrations/boto3/_services/base.py index 8cc29b59e7..5c7dd257b5 100644 --- a/sentry_sdk/integrations/boto3/_services/base.py +++ b/sentry_sdk/integrations/boto3/_services/base.py @@ -1,10 +1,12 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Optional + from typing import Any, Optional, Union from sentry_sdk._types import Attributes from sentry_sdk.integrations.boto3._context import AwsCallContext + from sentry_sdk.traces import StreamedSpan + from sentry_sdk.tracing import Span class _ServiceExtension: @@ -30,8 +32,11 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": """Return request attributes to add before the AWS request is made.""" return {} - def get_response_attributes( - self, ctx: "AwsCallContext", response: "Any" - ) -> "Attributes": - """Return response attributes to add after the AWS request is made.""" - return {} + def on_success_response( + self, + ctx: "AwsCallContext", + response: "Any", + span: "Union[Span, StreamedSpan]", + ) -> None: + """Hook that fires after a successful response. Responsible for closing the span.""" + return None diff --git a/sentry_sdk/integrations/boto3/_services/bedrock.py b/sentry_sdk/integrations/boto3/_services/bedrock.py index 690121fbfe..bc8ad1748b 100644 --- a/sentry_sdk/integrations/boto3/_services/bedrock.py +++ b/sentry_sdk/integrations/boto3/_services/bedrock.py @@ -1,18 +1,149 @@ import json +import sys from typing import TYPE_CHECKING, cast import sentry_sdk from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import DidNotEnable from sentry_sdk.integrations.boto3._services.base import _ServiceExtension -from sentry_sdk.utils import has_data_collection_enabled +from sentry_sdk.traces import SpanStatus, StreamedSpan +from sentry_sdk.tracing import SPANSTATUS +from sentry_sdk.utils import ( + capture_internal_exceptions, + has_data_collection_enabled, + reraise, + should_be_treated_as_error, +) + +try: + from botocore.eventstream import EventStream +except ImportError: + raise DidNotEnable("botocore is not installed") + if TYPE_CHECKING: - from typing import Any, Optional + from typing import Any, Iterator, Optional, Union - from mypy_boto3_bedrock_runtime.type_defs import ConverseResponseTypeDef + from mypy_boto3_bedrock_runtime.type_defs import ( + ConverseResponseTypeDef, + ConverseStreamResponseTypeDef, + ) from sentry_sdk._types import Attributes, ToolDefinition from sentry_sdk.integrations.boto3._context import AwsCallContext + from sentry_sdk.tracing import Span + + +def _finish_span( + stream: "EventStream[Any]", + exc_type: "Optional[type[BaseException]]", + exc_val: "Optional[BaseException]", + exc_tb: "Optional[Any]", +) -> None: + """ + Finishes the span attached to the stream. + Is a no-op if the stream has no span set, i.e., when the span has already been finished. + """ + if not hasattr(stream, "_span"): + return + + span = stream._span + del stream._span + + if exc_type is not None and should_be_treated_as_error(exc_type, exc_val): + if isinstance(span, StreamedSpan): + span.status = SpanStatus.ERROR + else: + span.status = SPANSTATUS.INTERNAL_ERROR + + if isinstance(span, StreamedSpan): + span.end() + else: + span.finish() + + +def _patch_event_stream() -> None: + original_iter = EventStream.__iter__ + original_close = EventStream.close + + def sentry_patched_iter( + self: "EventStream[Any]", + ) -> "Iterator[Any]": + """ + Sets information received while iterating the response stream on the AI Client Span. + Responsible for closing the AI Client Span unless the span has already been closed in the close() patch. + """ + span = getattr(self, "_span", None) + if span is None: + yield from original_iter(self) + return + + set_on_span = ( + span.set_attribute if isinstance(span, StreamedSpan) else span.set_data + ) + + try: + for event in original_iter(self): + if "metadata" not in event: + yield event + continue + + usage = event["metadata"]["usage"] + + input_tokens = usage["inputTokens"] + set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"]) + set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage["totalTokens"]) + + cache_read_tokens = usage.get("cacheReadInputTokens") + if cache_read_tokens is not None: + input_tokens += cache_read_tokens + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, cache_read_tokens + ) + + cache_write_tokens = usage.get("cacheWriteInputTokens") + if cache_write_tokens is not None: + input_tokens += cache_write_tokens + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, + cache_write_tokens, + ) + + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, + input_tokens, + ) + + yield event + except GeneratorExit: + raise # Leave span open since underlying stream is not necessarily closed. + except BaseException: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _finish_span(self, *exc_info) + reraise(*exc_info) + + _finish_span(self, None, None, None) + + def sentry_patched_close(self: "EventStream[Any]") -> None: + """ + Closes the AI Client Span unless `sentry_patched_iter()` closes the span first. + """ + if not hasattr(self, "_span"): + return original_close(self) + + try: + original_close(self) + except BaseException: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _finish_span(self, *exc_info) + reraise(*exc_info) + + _finish_span(self, None, None, None) + + EventStream.__iter__ = sentry_patched_iter # type: ignore[method-assign] + EventStream.close = sentry_patched_close # type: ignore[method-assign] def _transform_tool_definitions( @@ -119,34 +250,54 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": return attributes - def get_response_attributes( + def on_success_response( self, ctx: "AwsCallContext", response: "Any", - ) -> "Attributes": - if ctx.operation_name != "Converse" or "Error" in response: - return {} + span: "Union[Span, StreamedSpan]", + ) -> None: + set_on_span = ( + span.set_attribute if isinstance(span, StreamedSpan) else span.set_data + ) - usage = cast("ConverseResponseTypeDef", response)["usage"] + if ctx.operation_name == "Converse": + usage = cast("ConverseResponseTypeDef", response)["usage"] - input_tokens = usage["inputTokens"] - attributes: "Attributes" = { - SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS: usage["outputTokens"], - SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS: usage["totalTokens"], - } + set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"]) + set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage["totalTokens"]) - cache_read_tokens = usage.get("cacheReadInputTokens") - if cache_read_tokens is not None: - input_tokens += cache_read_tokens - attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] = cache_read_tokens + input_tokens = usage["inputTokens"] - cache_write_tokens = usage.get("cacheWriteInputTokens") - if cache_write_tokens is not None: - input_tokens += cache_write_tokens - attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] = ( - cache_write_tokens - ) + cache_read_tokens = usage.get("cacheReadInputTokens") + if cache_read_tokens is not None: + input_tokens += cache_read_tokens + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, cache_read_tokens + ) - attributes[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] = input_tokens + cache_write_tokens = usage.get("cacheWriteInputTokens") + if cache_write_tokens is not None: + input_tokens += cache_write_tokens + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, cache_write_tokens + ) - return attributes + set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + + if isinstance(span, StreamedSpan): + span.end() + return + + span.finish() + return + + if ctx.operation_name == "ConverseStream": + stream = cast("ConverseStreamResponseTypeDef", response)["stream"] + stream._span = span # type: ignore[attr-defined] + return + + if isinstance(span, StreamedSpan): + span.end() + return + + span.finish() diff --git a/tests/integrations/boto3/test_bedrock.py b/tests/integrations/boto3/test_bedrock.py index 4c17f73ef7..ab7934ecf4 100644 --- a/tests/integrations/boto3/test_bedrock.py +++ b/tests/integrations/boto3/test_bedrock.py @@ -1,10 +1,16 @@ +import io import json +from unittest.mock import Mock import boto3 import botocore +import botocore.session import pytest +import urllib3 +from botocore.eventstream import EventStream, EventStreamMessage from botocore.exceptions import ClientError from botocore.loaders import Loader +from botocore.parsers import EventStreamJSONParser from botocore.stub import Stubber from packaging.version import Version @@ -35,6 +41,118 @@ def converse_response(): } +@pytest.fixture +def converse_stream_response(): + usage = {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30} + if BOTOCORE_VERSION >= Version("1.37.24"): + usage["cacheReadInputTokens"] = 80 + usage["cacheWriteInputTokens"] = 20 + + stream_shape = ( + botocore.session.get_session() + .get_service_model("bedrock-runtime") + .operation_model("ConverseStream") + .output_shape.members["stream"] + ) + raw_stream = urllib3.HTTPResponse(body=io.BytesIO(b""), preload_content=False) + stream = EventStream( + raw_stream, stream_shape, EventStreamJSONParser(), "ConverseStream" + ) + stream._event_generator = iter( + [ + EventStreamMessage( + prelude=None, + headers={":message-type": "event", ":event-type": "messageStart"}, + payload=json.dumps({"role": "assistant"}).encode(), + crc=None, + ), + EventStreamMessage( + prelude=None, + headers={":message-type": "event", ":event-type": "contentBlockDelta"}, + payload=json.dumps( + {"contentBlockIndex": 0, "delta": {"text": "Hello"}} + ).encode(), + crc=None, + ), + EventStreamMessage( + prelude=None, + headers={":message-type": "event", ":event-type": "messageStop"}, + payload=json.dumps({"stopReason": "end_turn"}).encode(), + crc=None, + ), + EventStreamMessage( + prelude=None, + headers={":message-type": "event", ":event-type": "metadata"}, + payload=json.dumps( + {"usage": usage, "metrics": {"latencyMs": 500}} + ).encode(), + crc=None, + ), + ] + ) + + return { + "stream": stream, + "ResponseMetadata": { + "RequestId": "test-id", + "HTTPStatusCode": 200, + "HTTPHeaders": {}, + "RetryAttempts": 0, + }, + } + + +@pytest.fixture +def converse_stream_error_response(): + stream_shape = ( + botocore.session.get_session() + .get_service_model("bedrock-runtime") + .operation_model("ConverseStream") + .output_shape.members["stream"] + ) + raw_stream = urllib3.HTTPResponse(body=io.BytesIO(b""), preload_content=False) + stream = EventStream( + raw_stream, stream_shape, EventStreamJSONParser(), "ConverseStream" + ) + stream._event_generator = iter( + [ + EventStreamMessage( + prelude=None, + headers={":message-type": "event", ":event-type": "messageStart"}, + payload=json.dumps({"role": "assistant"}).encode(), + crc=None, + ), + EventStreamMessage( + prelude=None, + headers={":message-type": "event", ":event-type": "contentBlockDelta"}, + payload=json.dumps( + {"contentBlockIndex": 0, "delta": {"text": "Hi"}} + ).encode(), + crc=None, + ), + EventStreamMessage( + prelude=None, + headers={ + ":message-type": "exception", + ":exception-type": "modelStreamErrorException", + }, + payload=json.dumps({"message": "Stream failed"}).encode(), + crc=None, + ), + ] + ) + + return { + "stream": stream, + "ResponseMetadata": { + "RequestId": "test-id", + "HTTPStatusCode": 200, + "HTTPHeaders": {}, + "RetryAttempts": 0, + }, + } + + @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_converse( @@ -359,7 +477,12 @@ def test_converse_tool_definitions( @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_converse_stream( - sentry_init, capture_items, capture_events, span_streaming, stream_gen_ai_spans + sentry_init, + capture_items, + capture_events, + span_streaming, + stream_gen_ai_spans, + converse_stream_response, ): sentry_init( traces_sample_rate=1.0, @@ -372,14 +495,17 @@ def test_converse_stream( aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" ) client = session.client("bedrock-runtime") - stubber = Stubber(client) - stubber.add_response("converse_stream", {}) + mock_http = Mock(status_code=200) + client.meta.events.register( + "before-call.bedrock-runtime.ConverseStream", + lambda **kwargs: (mock_http, converse_stream_response), + ) if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with stubber, sentry_sdk.start_transaction(name="parent"): - client.converse_stream( + with sentry_sdk.start_transaction(name="parent"): + response = client.converse_stream( modelId="anthropic.claude-3-sonnet-20240229-v1:0", messages=[{"role": "user", "content": [{"text": "Hi"}]}], inferenceConfig={ @@ -388,6 +514,8 @@ def test_converse_stream( "topP": 0.9, }, ) + for _ in response["stream"]: + pass sentry_sdk.flush() spans = [item.payload for item in items] @@ -408,11 +536,20 @@ def test_converse_stream( assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 1024 assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + + assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + if BOTOCORE_VERSION < Version("1.37.24"): + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + return + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 110 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 else: events = capture_events() - with stubber, sentry_sdk.start_transaction(name="parent"): - client.converse_stream( + with sentry_sdk.start_transaction(name="parent"): + response = client.converse_stream( modelId="anthropic.claude-3-sonnet-20240229-v1:0", messages=[{"role": "user", "content": [{"text": "Hi"}]}], inferenceConfig={ @@ -421,6 +558,8 @@ def test_converse_stream( "topP": 0.9, }, ) + for _ in response["stream"]: + pass (transaction,) = events (span,) = transaction["spans"] @@ -437,11 +576,25 @@ def test_converse_stream( assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + if BOTOCORE_VERSION < Version("1.37.24"): + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + return + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 110 + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80 + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 + @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_converse_stream_error( - sentry_init, capture_items, capture_events, span_streaming, stream_gen_ai_spans +def test_converse_stream_error_during_iteration( + sentry_init, + capture_items, + capture_events, + span_streaming, + stream_gen_ai_spans, + converse_stream_error_response, ): sentry_init( traces_sample_rate=1.0, @@ -454,38 +607,23 @@ def test_converse_stream_error( aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" ) client = session.client("bedrock-runtime") - original_exception = ClientError( - error_response={ - "Error": {"Code": "AccessDeniedException", "Message": "Access denied"}, - "ResponseMetadata": { - "RequestId": "test-id", - "HostId": "", - "HTTPStatusCode": 403, - "HTTPHeaders": {}, - "RetryAttempts": 1, - }, - }, - operation_name="ConverseStream", + mock_http = Mock(status_code=200) + client.meta.events.register( + "before-call.bedrock-runtime.ConverseStream", + lambda **kwargs: (mock_http, converse_stream_error_response), ) - def raise_client_error(**kwargs): - raise original_exception - - client.meta.events.register("before-parameter-build", raise_client_error) - if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with sentry_sdk.start_transaction(name="parent"), pytest.raises(ClientError): - client.converse_stream( + with sentry_sdk.start_transaction(name="parent"): + response = client.converse_stream( modelId="anthropic.claude-3-sonnet-20240229-v1:0", messages=[{"role": "user", "content": [{"text": "Hi"}]}], - inferenceConfig={ - "maxTokens": 1024, - "temperature": 0.7, - "topP": 0.9, - }, ) + with pytest.raises(ClientError): + for _ in response["stream"]: + pass sentry_sdk.flush() spans = [item.payload for item in items] @@ -495,56 +633,22 @@ def raise_client_error(**kwargs): if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT ) - assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" - assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - assert ( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "anthropic.claude-3-sonnet-20240229-v1:0" - ) - - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 1024 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - - assert span["attributes"][SPANDATA.ERROR_TYPE] == "AccessDeniedException" - assert span["attributes"][SPANDATA.HTTP_STATUS_CODE] == 403 - assert span["attributes"][SPANDATA.AWS_REQUEST_ID] == "test-id" - assert span["attributes"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1 assert span["status"] == "error" else: events = capture_events() - with sentry_sdk.start_transaction(name="parent"), pytest.raises(ClientError): - client.converse_stream( + with sentry_sdk.start_transaction(name="parent"): + response = client.converse_stream( modelId="anthropic.claude-3-sonnet-20240229-v1:0", messages=[{"role": "user", "content": [{"text": "Hi"}]}], - inferenceConfig={ - "maxTokens": 1024, - "temperature": 0.7, - "topP": 0.9, - }, ) + with pytest.raises(ClientError): + for _ in response["stream"]: + pass (transaction,) = events (span,) = transaction["spans"] - assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" - assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - assert ( - span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "anthropic.claude-3-sonnet-20240229-v1:0" - ) - - assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 1024 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - - assert span["data"][SPANDATA.ERROR_TYPE] == "AccessDeniedException" - assert span["data"][SPANDATA.HTTP_STATUS_CODE] == 403 - assert span["data"][SPANDATA.AWS_REQUEST_ID] == "test-id" - assert span["data"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1 assert span["status"] == "internal_error" @@ -558,6 +662,7 @@ def test_converse_stream_tool_definitions( data_collection, span_streaming, stream_gen_ai_spans, + converse_stream_response, ): if data_collection: sentry_init( @@ -611,18 +716,23 @@ def test_converse_stream_tool_definitions( aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" ) client = session.client("bedrock-runtime") - stubber = Stubber(client) - stubber.add_response("converse_stream", {}) + mock_http = Mock(status_code=200) + client.meta.events.register( + "before-call.bedrock-runtime.ConverseStream", + lambda **kwargs: (mock_http, converse_stream_response), + ) if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with stubber, sentry_sdk.start_transaction(name="parent"): - client.converse_stream( + with sentry_sdk.start_transaction(name="parent"): + response = client.converse_stream( modelId="anthropic.claude-3-sonnet-20240229-v1:0", messages=[{"role": "user", "content": [{"text": "Hi"}]}], toolConfig={"tools": tools}, ) + for _ in response["stream"]: + pass sentry_sdk.flush() spans = [item.payload for item in items] @@ -639,12 +749,14 @@ def test_converse_stream_tool_definitions( else: events = capture_events() - with stubber, sentry_sdk.start_transaction(name="parent"): - client.converse_stream( + with sentry_sdk.start_transaction(name="parent"): + response = client.converse_stream( modelId="anthropic.claude-3-sonnet-20240229-v1:0", messages=[{"role": "user", "content": [{"text": "Hi"}]}], toolConfig={"tools": tools}, ) + for _ in response["stream"]: + pass (transaction,) = events (span,) = transaction["spans"]