Skip to content

Make Anthropic integration work with Potel #3686

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Nov 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions sentry_sdk/integrations/anthropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from sentry_sdk.integrations import DidNotEnable, Integration
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.utils import (
_serialize_span_attribute,
capture_internal_exceptions,
event_from_exception,
package_version,
Expand Down Expand Up @@ -126,7 +127,7 @@ def _add_ai_data_to_span(
complete_message = "".join(content_blocks)
span.set_data(
SPANDATA.AI_RESPONSES,
[{"type": "text", "text": complete_message}],
_serialize_span_attribute([{"type": "text", "text": complete_message}]),
)
total_tokens = input_tokens + output_tokens
record_token_usage(span, input_tokens, output_tokens, total_tokens)
Expand Down Expand Up @@ -165,11 +166,16 @@ def _sentry_patched_create_common(f, *args, **kwargs):
span.set_data(SPANDATA.AI_STREAMING, False)

if should_send_default_pii() and integration.include_prompts:
span.set_data(SPANDATA.AI_INPUT_MESSAGES, messages)
span.set_data(
SPANDATA.AI_INPUT_MESSAGES, _serialize_span_attribute(messages)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we maybe just call _serialize_span_attribute directly in POTelSpan.set_attribute? (which set_data calls under the hood if I'm not mistaken)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good idea! Will create a PR for that!

)

if hasattr(result, "content"):
if should_send_default_pii() and integration.include_prompts:
span.set_data(SPANDATA.AI_RESPONSES, _get_responses(result.content))
span.set_data(
SPANDATA.AI_RESPONSES,
_serialize_span_attribute(_get_responses(result.content)),
)
_calculate_token_usage(result, span)
span.__exit__(None, None, None)

Expand Down
63 changes: 38 additions & 25 deletions tests/integrations/anthropic/test_anthropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,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.utils import package_version
from sentry_sdk.utils import _serialize_span_attribute, package_version

try:
from anthropic.types import InputJSONDelta
Expand Down Expand Up @@ -115,10 +115,12 @@ def test_nonstreaming_create_message(
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] == [
{"type": "text", "text": "Hi, I'm Claude."}
]
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == _serialize_span_attribute(
messages
)
assert span["data"][SPANDATA.AI_RESPONSES] == _serialize_span_attribute(
[{"type": "text", "text": "Hi, I'm Claude."}]
)
else:
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
assert SPANDATA.AI_RESPONSES not in span["data"]
Expand Down Expand Up @@ -183,10 +185,12 @@ async def test_nonstreaming_create_message_async(
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] == [
{"type": "text", "text": "Hi, I'm Claude."}
]
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == _serialize_span_attribute(
messages
)
assert span["data"][SPANDATA.AI_RESPONSES] == _serialize_span_attribute(
[{"type": "text", "text": "Hi, I'm Claude."}]
)
else:
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
assert SPANDATA.AI_RESPONSES not in span["data"]
Expand Down Expand Up @@ -282,10 +286,12 @@ def test_streaming_create_message(
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] == [
{"type": "text", "text": "Hi! I'm Claude!"}
]
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == _serialize_span_attribute(
messages
)
assert span["data"][SPANDATA.AI_RESPONSES] == _serialize_span_attribute(
[{"type": "text", "text": "Hi! I'm Claude!"}]
)

else:
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
Expand Down Expand Up @@ -385,10 +391,12 @@ async def test_streaming_create_message_async(
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] == [
{"type": "text", "text": "Hi! I'm Claude!"}
]
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == _serialize_span_attribute(
messages
)
assert span["data"][SPANDATA.AI_RESPONSES] == _serialize_span_attribute(
[{"type": "text", "text": "Hi! I'm Claude!"}]
)

else:
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
Expand Down Expand Up @@ -515,10 +523,12 @@ def test_streaming_create_message_with_input_json_delta(
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
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == _serialize_span_attribute(
messages
)
assert span["data"][SPANDATA.AI_RESPONSES] == _serialize_span_attribute(
[{"type": "text", "text": ""}]
) # we do not record InputJSONDelta because it could contain PII

else:
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
Expand Down Expand Up @@ -652,10 +662,12 @@ async def test_streaming_create_message_with_input_json_delta_async(
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
assert span["data"][SPANDATA.AI_INPUT_MESSAGES] == _serialize_span_attribute(
messages
)
assert span["data"][SPANDATA.AI_RESPONSES] == _serialize_span_attribute(
[{"type": "text", "text": ""}]
) # we do not record InputJSONDelta because it could contain PII

else:
assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
Expand All @@ -667,6 +679,7 @@ async def test_streaming_create_message_with_input_json_delta_async(
assert span["data"]["ai.streaming"] is True


@pytest.mark.forked
def test_exception_message_create(sentry_init, capture_events):
sentry_init(integrations=[AnthropicIntegration()], traces_sample_rate=1.0)
events = capture_events()
Expand Down
Loading