fix(anthropic): Add partial json support to streams (#3674)

Add `partial_json` for tool calling when streaming in Anthropic integrations. (This is an addition to #3615

---------

Co-authored-by: Anton Pirker <[email protected]>

Author: kwnath

Diff for: 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 for: 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