Skip to content

Commit

Permalink
fix(anthropic): Add partial json support to streams (#3674)
Browse files Browse the repository at this point in the history
Add `partial_json` for tool calling when streaming in Anthropic integrations. (This is an addition to #3615

---------

Co-authored-by: Anton Pirker <anton.pirker@sentry.io>
  • Loading branch information
kwnath and antonpirker authored Feb 20, 2025
1 parent 2423299 commit 48ebd73
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 7 deletions.
2 changes: 2 additions & 0 deletions sentry_sdk/integrations/anthropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down
71 changes: 64 additions & 7 deletions tests/integrations/anthropic/test_anthropic.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from unittest import mock


try:
from unittest.mock import AsyncMock
except ImportError:
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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

Expand Down Expand Up @@ -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"]
Expand Down Expand Up @@ -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"]
Expand Down Expand Up @@ -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

0 comments on commit 48ebd73

Please sign in to comment.