Fix Anthropic integration when using tool calls

sentry_sdk/integrations/anthropic.py

@@ -1,4 +1,5 @@
 from functools import wraps
+from typing import TYPE_CHECKING
 
 import sentry_sdk
 from sentry_sdk.ai.monitoring import record_token_usage
@@ -11,8 +12,6 @@
     package_version,
 )
 
-from typing import TYPE_CHECKING
-
 try:
     from anthropic.resources import Messages
 
@@ -74,6 +73,21 @@ def _calculate_token_usage(result, span):
     record_token_usage(span, input_tokens, output_tokens, total_tokens)
 
 
+def _get_responses(content):
+    # type: (list[Any]) -> list[dict[str, Any]]
+    """Get JSON of a Anthropic responses."""
+    responses = []
+    for item in content:
+        if hasattr(item, "text"):
+            responses.append(
+                {
+                    "type": item.type,
+                    "text": item.text,
+                }
+            )
+    return responses
+
+
 def _wrap_message_create(f):
     # type: (Any) -> Any
     @wraps(f)
@@ -113,18 +127,7 @@ def _sentry_patched_create(*args, **kwargs):
                 span.set_data(SPANDATA.AI_INPUT_MESSAGES, messages)
             if hasattr(result, "content"):
                 if should_send_default_pii() and integration.include_prompts:
-                    span.set_data(
-                        SPANDATA.AI_RESPONSES,
-                        list(
-                            map(
-                                lambda message: {
-                                    "type": message.type,
-                                    "text": message.text,
-                                },
-                                result.content,
-                            )
-                        ),
-                    )
+                    span.set_data(SPANDATA.AI_RESPONSES, _get_responses(result.content))
                 _calculate_token_usage(result, span)
                 span.__exit__(None, None, None)
             elif hasattr(result, "_iterator"):
@@ -145,7 +148,8 @@ def new_iterator():
                                 elif event.type == "content_block_start":
                                     pass
                                 elif event.type == "content_block_delta":
-                                    content_blocks.append(event.delta.text)
+                                    if hasattr(event.delta, "text"):
+                                        content_blocks.append(event.delta.text)
                                 elif event.type == "content_block_stop":
                                     pass
                                 elif event.type == "message_delta":

tests/integrations/anthropic/test_anthropic.py

@@ -1,17 +1,29 @@
-import pytest
 from unittest import mock
-from anthropic import Anthropic, Stream, AnthropicError
-from anthropic.types import Usage, MessageDeltaUsage, TextDelta
+
+import pytest
+from anthropic import Anthropic, AnthropicError, 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
+from anthropic.types.content_block_stop_event import ContentBlockStopEvent
 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.content_block_start_event import ContentBlockStartEvent
-from anthropic.types.content_block_delta_event import ContentBlockDeltaEvent
-from anthropic.types.content_block_stop_event import ContentBlockStopEvent
+
+from sentry_sdk.utils import package_version
+
+try:
+    from anthropic.types import InputJSONDelta
+except ImportError:
+    try:
+        from anthropic.types import InputJsonDelta as InputJSONDelta
+    except ImportError:
+        pass
 
 try:
     # 0.27+
     from anthropic.types.raw_message_delta_event import Delta
+    from anthropic.types.tool_use_block import ToolUseBlock
 except ImportError:
     # pre 0.27
     from anthropic.types.message_delta_event import Delta
@@ -25,7 +37,7 @@
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.anthropic import AnthropicIntegration
 
-
+ANTHROPIC_VERSION = package_version("anthropic")
 EXAMPLE_MESSAGE = Message(
     id="id",
     model="model",
@@ -203,6 +215,136 @@ def test_streaming_create_message(
     assert span["data"]["ai.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(
+    "send_default_pii, include_prompts",
+    [
+        (True, True),
+        (True, False),
+        (False, True),
+        (False, False),
+    ],
+)
+def test_streaming_create_message_with_input_json_delta(
+    sentry_init, capture_events, send_default_pii, include_prompts
+):
+    client = Anthropic(api_key="z")
+    returned_stream = Stream(cast_to=None, response=None, client=client)
+    returned_stream._iterator = [
+        MessageStartEvent(
+            message=Message(
+                id="msg_0",
+                content=[],
+                model="claude-3-5-sonnet-20240620",
+                role="assistant",
+                stop_reason=None,
+                stop_sequence=None,
+                type="message",
+                usage=Usage(input_tokens=366, output_tokens=10),
+            ),
+            type="message_start",
+        ),
+        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="", type="input_json_delta"),
+            index=0,
+            type="content_block_delta",
+        ),
+        ContentBlockDeltaEvent(
+            delta=InputJSONDelta(partial_json="{'location':", type="input_json_delta"),
+            index=0,
+            type="content_block_delta",
+        ),
+        ContentBlockDeltaEvent(
+            delta=InputJSONDelta(partial_json=" 'S", type="input_json_delta"),
+            index=0,
+            type="content_block_delta",
+        ),
+        ContentBlockDeltaEvent(
+            delta=InputJSONDelta(partial_json="an ", type="input_json_delta"),
+            index=0,
+            type="content_block_delta",
+        ),
+        ContentBlockDeltaEvent(
+            delta=InputJSONDelta(partial_json="Francisco, C", type="input_json_delta"),
+            index=0,
+            type="content_block_delta",
+        ),
+        ContentBlockDeltaEvent(
+            delta=InputJSONDelta(partial_json="A'}", type="input_json_delta"),
+            index=0,
+            type="content_block_delta",
+        ),
+        ContentBlockStopEvent(type="content_block_stop", index=0),
+        MessageDeltaEvent(
+            delta=Delta(stop_reason="tool_use", stop_sequence=None),
+            usage=MessageDeltaUsage(output_tokens=41),
+            type="message_delta",
+        ),
+    ]
+
+    sentry_init(
+        integrations=[AnthropicIntegration(include_prompts=include_prompts)],
+        traces_sample_rate=1.0,
+        send_default_pii=send_default_pii,
+    )
+    events = capture_events()
+    client.messages._post = mock.Mock(return_value=returned_stream)
+
+    messages = [
+        {
+            "role": "user",
+            "content": "What is the weather like in San Francisco?",
+        }
+    ]
+
+    with start_transaction(name="anthropic"):
+        message = client.messages.create(
+            max_tokens=1024, messages=messages, model="model", stream=True
+        )
+
+        for _ in message:
+            pass
+
+    assert message == returned_stream
+    assert len(events) == 1
+    (event,) = events
+
+    assert event["type"] == "transaction"
+    assert event["transaction"] == "anthropic"
+
+    assert len(event["spans"]) == 1
+    (span,) = event["spans"]
+
+    assert span["op"] == OP.ANTHROPIC_MESSAGES_CREATE
+    assert span["description"] == "Anthropic messages create"
+    assert span["data"][SPANDATA.AI_MODEL_ID] == "model"
+
+    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
+
+    else:
+        assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
+        assert SPANDATA.AI_RESPONSES not in span["data"]
+
+    assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 366
+    assert span["measurements"]["ai_completion_tokens_used"]["value"] == 51
+    assert span["measurements"]["ai_total_tokens_used"]["value"] == 417
+    assert span["data"]["ai.streaming"] is True
+
+
 def test_exception_message_create(sentry_init, capture_events):
     sentry_init(integrations=[AnthropicIntegration()], traces_sample_rate=1.0)
     events = capture_events()