diff --git a/sentry_sdk/integrations/anthropic.py b/sentry_sdk/integrations/anthropic.py index 4b04bc1cfc..c34b28ff1f 100644 --- a/sentry_sdk/integrations/anthropic.py +++ b/sentry_sdk/integrations/anthropic.py @@ -93,7 +93,7 @@ class _PatchedRawMessageStream(Stream[RawMessageStreamEvent]): _model: Optional[ModelParam] _usage: "_RecordedUsage" - _content_blocks: list[Any] + _content_blocks: dict[int, dict[str, Any]] _response_id: Optional[str] _finish_reason: Optional[str] @@ -103,7 +103,7 @@ class _PatchedMessageStream(MessageStream): _model: Optional[ModelParam] _usage: "_RecordedUsage" - _content_blocks: list[Any] + _content_blocks: dict[int, dict[str, Any]] _response_id: Optional[str] _finish_reason: Optional[str] @@ -113,7 +113,7 @@ class _PatchedRawAsyncMessageStream(AsyncStream[RawMessageStreamEvent]): _model: Optional[ModelParam] _usage: "_RecordedUsage" - _content_blocks: list[Any] + _content_blocks: dict[int, dict[str, Any]] _response_id: Optional[str] _finish_reason: Optional[str] @@ -123,7 +123,7 @@ class _PatchedAsyncMessageStream(AsyncMessageStream): _model: Optional[ModelParam] _usage: "_RecordedUsage" - _content_blocks: list[Any] + _content_blocks: dict[int, dict[str, Any]] _response_id: Optional[str] _finish_reason: Optional[str] @@ -323,24 +323,47 @@ def _collect_ai_data( event: "MessageStreamEvent", model: "str | None", usage: "_RecordedUsage", - content_blocks: "list[str]", + content_blocks: "dict[int, dict[str, Any]]", response_id: "str | None" = None, finish_reason: "str | None" = None, -) -> "tuple[str | None, _RecordedUsage, list[str], str | None, str | None]": +) -> "tuple[str | None, _RecordedUsage, dict[int, dict[str, Any]], str | None, str | None]": """ Collect model information, token usage, and collect content blocks from the AI streaming response. """ with capture_internal_exceptions(): if hasattr(event, "type"): if event.type == "content_block_start": - pass + content_block = event.content_block + if content_block.type == "text": + content_blocks[event.index] = { + "type": "text", + "text": content_block.text, + } + elif content_block.type == "tool_use": + content_blocks[event.index] = { + "type": "tool_use", + "id": content_block.id, + "name": content_block.name, + "input": "", + } elif event.type == "content_block_delta": - if hasattr(event.delta, "text"): - content_blocks.append(event.delta.text) - elif hasattr(event.delta, "partial_json"): - content_blocks.append(event.delta.partial_json) + block = content_blocks.get(event.index) + if block is not None: + if block["type"] == "text" and hasattr(event.delta, "text"): + block["text"] += event.delta.text + elif block["type"] == "tool_use" and hasattr( + event.delta, "partial_json" + ): + block["input"] += event.delta.partial_json elif event.type == "content_block_stop": - pass + block = content_blocks.get(event.index) + if block is not None and block["type"] == "tool_use": + try: + block["input"] = ( + json.loads(block["input"]) if block["input"] else {} + ) + except ValueError: + pass # Token counting logic mirrors anthropic SDK, which also extracts already accumulated tokens. # https://github.com/anthropics/anthropic-sdk-python/blob/9c485f6966e10ae0ea9eabb3a921d2ea8145a25b/src/anthropic/lib/streaming/_messages.py#L433-L518 @@ -946,7 +969,7 @@ def _initialize_data_accumulation_state( if not hasattr(stream, "_model"): stream._model = None # type: ignore[union-attr] stream._usage = _RecordedUsage() # type: ignore[union-attr] - stream._content_blocks = [] # type: ignore[union-attr] + stream._content_blocks = {} # type: ignore[union-attr] stream._response_id = None # type: ignore[union-attr] stream._finish_reason = None # type: ignore[union-attr] @@ -979,7 +1002,7 @@ def _set_streaming_output_data( integration: "AnthropicIntegration", model: "Optional[str]", usage: "_RecordedUsage", - content_blocks: "list[str]", + content_blocks: "dict[int, dict[str, Any]]", response_id: "Optional[str]", finish_reason: "Optional[str]", ) -> None: @@ -1002,7 +1025,7 @@ def _set_streaming_output_data( output_tokens=usage.output_tokens, cache_read_input_tokens=usage.cache_read_input_tokens, cache_write_input_tokens=usage.cache_write_input_tokens, - content_blocks=[{"text": "".join(content_blocks), "type": "text"}], + content_blocks=list(content_blocks.values()), response_id=response_id, finish_reason=finish_reason, ) diff --git a/tests/integrations/anthropic/test_anthropic.py b/tests/integrations/anthropic/test_anthropic.py index b3835c4b72..0fe25b10a0 100644 --- a/tests/integrations/anthropic/test_anthropic.py +++ b/tests/integrations/anthropic/test_anthropic.py @@ -22,6 +22,7 @@ async def __call__(self, *args, **kwargs): from anthropic.types.message import Message from anthropic.types.message_delta_event import MessageDeltaEvent from anthropic.types.message_start_event import MessageStartEvent +from anthropic.types.message_stop_event import MessageStopEvent try: from anthropic import APIStatusError @@ -4657,13 +4658,21 @@ def test_streaming_create_message_with_input_json_delta( span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -4704,13 +4713,21 @@ def test_streaming_create_message_with_input_json_delta( span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -4752,13 +4769,19 @@ def test_streaming_create_message_with_input_json_delta( span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert json.loads(span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["data"] assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -4910,13 +4933,21 @@ def test_stream_messages_with_input_json_delta( span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -4957,13 +4988,21 @@ def test_stream_messages_with_input_json_delta( span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -5004,13 +5043,19 @@ def test_stream_messages_with_input_json_delta( span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert json.loads(span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["data"] assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -5169,14 +5214,22 @@ async def test_streaming_create_message_with_input_json_delta_async( span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -5217,14 +5270,22 @@ async def test_streaming_create_message_with_input_json_delta_async( span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -5265,14 +5326,20 @@ async def test_streaming_create_message_with_input_json_delta_async( span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert json.loads(span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["data"] assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -5432,13 +5499,21 @@ async def test_stream_message_with_input_json_delta_async( span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -5480,13 +5555,21 @@ async def test_stream_message_with_input_json_delta_async( span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -5528,14 +5611,20 @@ async def test_stream_message_with_input_json_delta_async( span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] == '[{"role": "user", "content": "What is the weather like in San Francisco?"}]' ) - assert ( - span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == '{"location": "San Francisco, CA"}' - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert json.loads(span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == [ + { + "id": "toolu_0", + "input": {"location": "San Francisco, CA"}, + "name": "get_weather", + "type": "tool_use", + } + ] else: assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["data"] assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 366 assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 41 @@ -5543,6 +5632,343 @@ async def test_stream_message_with_input_json_delta_async( assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True +@pytest.mark.skipif( + ANTHROPIC_VERSION < (0, 27), + reason="Versions <0.27.0 do not include InputJSONDelta, which was introduced in >=0.27.0 along with a new message delta type for tool calling.", +) +@pytest.mark.parametrize("api", ["create", "stream"]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,outputs_collected", + [ + pytest.param( + {"gen_ai": {"outputs": True}}, + False, + False, + True, + id="gen-ai-outputs-enabled-overrides-pii-and-include-prompts", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + True, + True, + False, + id="gen-ai-outputs-disabled-overrides-pii-and-include-prompts", + ), + pytest.param( + None, + True, + True, + True, + id="legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + False, + id="legacy-pii-disabled", + ), + ], +) +def test_streaming_text_and_tool_use_output( + sentry_init, + capture_items, + get_model_response, + server_side_event_chunks, + api, + data_collection, + send_default_pii, + include_prompts, + outputs_collected, +): + sentry_init_kwargs = dict( + integrations=[AnthropicIntegration(include_prompts=include_prompts)], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + sentry_init(**sentry_init_kwargs) + items = capture_items("span") + + client = Anthropic(api_key="z") + + response = get_model_response( + server_side_event_chunks( + [ + MessageStartEvent( + message=Message( + id="msg_1", + content=[], + model="claude-sonnet-4-5", + role="assistant", + stop_reason=None, + stop_sequence=None, + type="message", + usage=Usage(input_tokens=20, output_tokens=1), + ), + type="message_start", + ), + ContentBlockStartEvent( + type="content_block_start", + index=0, + content_block=TextBlock(type="text", text=""), + ), + ContentBlockDeltaEvent( + delta=TextDelta(text="Checking ", type="text_delta"), + index=0, + type="content_block_delta", + ), + ContentBlockDeltaEvent( + delta=TextDelta(text="the weather.", type="text_delta"), + index=0, + type="content_block_delta", + ), + ContentBlockStopEvent(type="content_block_stop", index=0), + ContentBlockStartEvent( + type="content_block_start", + index=1, + content_block=ToolUseBlock( + id="toolu_1", input={}, name="get_weather", type="tool_use" + ), + ), + ContentBlockDeltaEvent( + delta=InputJSONDelta( + partial_json='{"city"', type="input_json_delta" + ), + index=1, + type="content_block_delta", + ), + ContentBlockDeltaEvent( + delta=InputJSONDelta( + partial_json=': "Paris"}', type="input_json_delta" + ), + index=1, + type="content_block_delta", + ), + ContentBlockStopEvent(type="content_block_stop", index=1), + MessageDeltaEvent( + delta=Delta(stop_reason="tool_use", stop_sequence=None), + usage=MessageDeltaUsage(output_tokens=30), + type="message_delta", + ), + MessageStopEvent(type="message_stop"), + ] + ) + ) + + messages = [{"role": "user", "content": "What is the weather in Paris?"}] + + with mock.patch.object( + client._client, + "send", + return_value=response, + ), start_transaction(name="anthropic"): + if api == "create": + for _ in client.messages.create( + max_tokens=1024, messages=messages, model="model", stream=True + ): + pass + else: + with client.messages.stream( + max_tokens=1024, messages=messages, model="model" + ) as stream: + for _ in stream: + pass + + sentry_sdk.flush() + (span,) = (item.payload for item in items) + + if outputs_collected: + assert ( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] == "Checking the weather." + ) + assert json.loads(span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == [ + { + "id": "toolu_1", + "input": {"city": "Paris"}, + "name": "get_weather", + "type": "tool_use", + } + ] + else: + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] + + +@pytest.mark.asyncio +@pytest.mark.skipif( + ANTHROPIC_VERSION < (0, 27), + reason="Versions <0.27.0 do not include InputJSONDelta, which was introduced in >=0.27.0 along with a new message delta type for tool calling.", +) +@pytest.mark.parametrize("api", ["create", "stream"]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,outputs_collected", + [ + pytest.param( + {"gen_ai": {"outputs": True}}, + False, + False, + True, + id="gen-ai-outputs-enabled-overrides-pii-and-include-prompts", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + True, + True, + False, + id="gen-ai-outputs-disabled-overrides-pii-and-include-prompts", + ), + pytest.param( + None, + True, + True, + True, + id="legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + False, + id="legacy-pii-disabled", + ), + ], +) +async def test_streaming_text_and_tool_use_output_async( + sentry_init, + capture_items, + get_model_response, + async_iterator, + server_side_event_chunks, + api, + data_collection, + send_default_pii, + include_prompts, + outputs_collected, +): + sentry_init_kwargs = dict( + integrations=[AnthropicIntegration(include_prompts=include_prompts)], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + sentry_init(**sentry_init_kwargs) + items = capture_items("span") + + client = AsyncAnthropic(api_key="z") + + response = get_model_response( + async_iterator( + server_side_event_chunks( + [ + MessageStartEvent( + message=Message( + id="msg_1", + content=[], + model="claude-sonnet-4-5", + role="assistant", + stop_reason=None, + stop_sequence=None, + type="message", + usage=Usage(input_tokens=20, output_tokens=1), + ), + type="message_start", + ), + ContentBlockStartEvent( + type="content_block_start", + index=0, + content_block=TextBlock(type="text", text=""), + ), + ContentBlockDeltaEvent( + delta=TextDelta(text="Checking ", type="text_delta"), + index=0, + type="content_block_delta", + ), + ContentBlockDeltaEvent( + delta=TextDelta(text="the weather.", type="text_delta"), + index=0, + type="content_block_delta", + ), + ContentBlockStopEvent(type="content_block_stop", index=0), + ContentBlockStartEvent( + type="content_block_start", + index=1, + content_block=ToolUseBlock( + id="toolu_1", + input={}, + name="get_weather", + type="tool_use", + ), + ), + ContentBlockDeltaEvent( + delta=InputJSONDelta( + partial_json='{"city"', type="input_json_delta" + ), + index=1, + type="content_block_delta", + ), + ContentBlockDeltaEvent( + delta=InputJSONDelta( + partial_json=': "Paris"}', type="input_json_delta" + ), + index=1, + type="content_block_delta", + ), + ContentBlockStopEvent(type="content_block_stop", index=1), + MessageDeltaEvent( + delta=Delta(stop_reason="tool_use", stop_sequence=None), + usage=MessageDeltaUsage(output_tokens=30), + type="message_delta", + ), + MessageStopEvent(type="message_stop"), + ] + ) + ) + ) + + messages = [{"role": "user", "content": "What is the weather in Paris?"}] + + with mock.patch.object( + client._client, + "send", + return_value=response, + ), start_transaction(name="anthropic"): + if api == "create": + async for _ in await client.messages.create( + max_tokens=1024, messages=messages, model="model", stream=True + ): + pass + else: + async with client.messages.stream( + max_tokens=1024, messages=messages, model="model" + ) as stream: + async for _ in stream: + pass + + sentry_sdk.flush() + (span,) = (item.payload for item in items) + + if outputs_collected: + assert ( + span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] == "Checking the weather." + ) + assert json.loads(span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) == [ + { + "id": "toolu_1", + "input": {"city": "Paris"}, + "name": "get_weather", + "type": "tool_use", + } + ] + else: + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span["attributes"] + + @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_exception_message_create( @@ -6004,26 +6430,39 @@ async def test_span_origin_async( ) @pytest.mark.parametrize("span_streaming", [True, False]) def test_collect_ai_data_with_input_json_delta(span_streaming): - event = ContentBlockDeltaEvent( - delta=InputJSONDelta(partial_json="test", type="input_json_delta"), - index=0, - type="content_block_delta", - ) model = None usage = _RecordedUsage() usage.output_tokens = 20 usage.input_tokens = 10 - content_blocks = [] + content_blocks = {} + + for event in [ + ContentBlockStartEvent( + type="content_block_start", + index=0, + content_block=ToolUseBlock( + id="toolu_0", input={}, name="get_weather", type="tool_use" + ), + ), + ContentBlockDeltaEvent( + delta=InputJSONDelta(partial_json="test", type="input_json_delta"), + index=0, + type="content_block_delta", + ), + ContentBlockStopEvent(type="content_block_stop", index=0), + ]: + model, new_usage, content_blocks, response_id, finish_reason = _collect_ai_data( + event, model, usage, content_blocks + ) - model, new_usage, new_content_blocks, response_id, finish_reason = _collect_ai_data( - event, model, usage, content_blocks - ) assert model is None assert new_usage.input_tokens == usage.input_tokens assert new_usage.output_tokens == usage.output_tokens - assert new_content_blocks == ["test"] + assert content_blocks == { + 0: {"type": "tool_use", "id": "toolu_0", "name": "get_weather", "input": "test"} + } assert response_id is None assert finish_reason is None