diff --git a/sentry_sdk/integrations/anthropic.py b/sentry_sdk/integrations/anthropic.py index f06d8a14db..4cb54309c8 100644 --- a/sentry_sdk/integrations/anthropic.py +++ b/sentry_sdk/integrations/anthropic.py @@ -101,6 +101,8 @@ def _collect_ai_data(event, input_tokens, output_tokens, content_blocks): 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) elif event.type == "content_block_stop": pass elif event.type == "message_delta": diff --git a/tests/integrations/anthropic/test_anthropic.py b/tests/integrations/anthropic/test_anthropic.py index 8ce12e70f5..7f6622a1ba 100644 --- a/tests/integrations/anthropic/test_anthropic.py +++ b/tests/integrations/anthropic/test_anthropic.py @@ -1,5 +1,6 @@ from unittest import mock + try: from unittest.mock import AsyncMock except ImportError: @@ -10,7 +11,7 @@ async def __call__(self, *args, **kwargs): import pytest -from anthropic import AsyncAnthropic, Anthropic, AnthropicError, AsyncStream, Stream +from anthropic import Anthropic, AnthropicError, AsyncAnthropic, AsyncStream, Stream from anthropic.types import MessageDeltaUsage, TextDelta, Usage from anthropic.types.content_block_delta_event import ContentBlockDeltaEvent from anthropic.types.content_block_start_event import ContentBlockStartEvent @@ -19,6 +20,7 @@ async def __call__(self, *args, **kwargs): from anthropic.types.message_delta_event import MessageDeltaEvent from anthropic.types.message_start_event import MessageStartEvent +from sentry_sdk.integrations.anthropic import _add_ai_data_to_span, _collect_ai_data from sentry_sdk.utils import package_version try: @@ -42,7 +44,7 @@ async def __call__(self, *args, **kwargs): except ImportError: from anthropic.types.content_block import ContentBlock as TextBlock -from sentry_sdk import start_transaction +from sentry_sdk import start_transaction, start_span from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.anthropic import AnthropicIntegration @@ -517,9 +519,8 @@ def test_streaming_create_message_with_input_json_delta( if send_default_pii and include_prompts: assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages assert span["data"][SPANDATA.AI_RESPONSES] == [ - {"text": "", "type": "text"} - ] # we do not record InputJSONDelta because it could contain PII - + {"text": "{'location': 'San Francisco, CA'}", "type": "text"} + ] else: assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] assert SPANDATA.AI_RESPONSES not in span["data"] @@ -654,8 +655,8 @@ async def test_streaming_create_message_with_input_json_delta_async( if send_default_pii and include_prompts: assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == messages assert span["data"][SPANDATA.AI_RESPONSES] == [ - {"text": "", "type": "text"} - ] # we do not record InputJSONDelta because it could contain PII + {"text": "{'location': 'San Francisco, CA'}", "type": "text"} + ] else: assert SPANDATA.AI_INPUT_MESSAGES not in span["data"] @@ -757,3 +758,59 @@ async def test_span_origin_async(sentry_init, capture_events): assert event["contexts"]["trace"]["origin"] == "manual" assert event["spans"][0]["origin"] == "auto.ai.anthropic" + + +@pytest.mark.skipif( + ANTHROPIC_VERSION < (0, 27), + reason="Versions <0.27.0 do not include InputJSONDelta.", +) +def test_collect_ai_data_with_input_json_delta(): + event = ContentBlockDeltaEvent( + delta=InputJSONDelta(partial_json="test", type="input_json_delta"), + index=0, + type="content_block_delta", + ) + + input_tokens = 10 + output_tokens = 20 + content_blocks = [] + + new_input_tokens, new_output_tokens, new_content_blocks = _collect_ai_data( + event, input_tokens, output_tokens, content_blocks + ) + + assert new_input_tokens == input_tokens + assert new_output_tokens == output_tokens + assert new_content_blocks == ["test"] + + +@pytest.mark.skipif( + ANTHROPIC_VERSION < (0, 27), + reason="Versions <0.27.0 do not include InputJSONDelta.", +) +def test_add_ai_data_to_span_with_input_json_delta(sentry_init): + sentry_init( + integrations=[AnthropicIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with start_transaction(name="test"): + span = start_span() + integration = AnthropicIntegration() + + _add_ai_data_to_span( + span, + integration, + input_tokens=10, + output_tokens=20, + content_blocks=["{'test': 'data',", "'more': 'json'}"], + ) + + assert span._data.get(SPANDATA.AI_RESPONSES) == [ + {"type": "text", "text": "{'test': 'data','more': 'json'}"} + ] + assert span._data.get("ai.streaming") is True + assert span._measurements.get("ai_prompt_tokens_used")["value"] == 10 + assert span._measurements.get("ai_completion_tokens_used")["value"] == 20 + assert span._measurements.get("ai_total_tokens_used")["value"] == 30