Skip to content
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
170 changes: 116 additions & 54 deletions sentry_sdk/integrations/anthropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,12 @@
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing import Span
from sentry_sdk.tracing_utils import (
has_span_streaming_enabled,
should_truncate_gen_ai_input,
)
from sentry_sdk.utils import (
capture_internal_exceptions,
event_from_exception,
Expand Down Expand Up @@ -78,7 +84,6 @@
)

from sentry_sdk._types import TextPart
from sentry_sdk.tracing import Span


class _RecordedUsage:
Expand Down Expand Up @@ -366,7 +371,7 @@ def _transform_system_instructions(


def _set_common_input_data(
span: "Span",
span: "Union[Span, StreamedSpan]",
integration: "AnthropicIntegration",
max_tokens: "int",
messages: "Iterable[MessageParam]",
Expand All @@ -380,16 +385,19 @@ def _set_common_input_data(
"""
Set input data for the span based on the provided keyword arguments for the anthropic message creation.
"""
span.set_data(SPANDATA.GEN_AI_SYSTEM, "anthropic")
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
set_on_span = (
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
)
set_on_span(SPANDATA.GEN_AI_SYSTEM, "anthropic")
set_on_span(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
if (
messages is not None
and len(messages) > 0 # type: ignore
and should_send_default_pii()
and integration.include_prompts
):
if isinstance(system, str) or isinstance(system, Iterable):
span.set_data(
set_on_span(
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
json.dumps(_transform_system_instructions(system)),
)
Expand Down Expand Up @@ -442,37 +450,44 @@ def _set_common_input_data(
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
role_normalized_messages
if client.options.get("stream_gen_ai_spans", False)
else truncate_and_annotate_messages(role_normalized_messages, span, scope)
truncate_and_annotate_messages(role_normalized_messages, span, scope)
if should_truncate_gen_ai_input(client.options)
else role_normalized_messages
)
if messages_data is not None:
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
)

if max_tokens is not None and _is_given(max_tokens):
span.set_data(SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
set_on_span(SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
if model is not None and _is_given(model):
span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model)
set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
if temperature is not None and _is_given(temperature):
span.set_data(SPANDATA.GEN_AI_REQUEST_TEMPERATURE, temperature)
set_on_span(SPANDATA.GEN_AI_REQUEST_TEMPERATURE, temperature)
if top_k is not None and _is_given(top_k):
span.set_data(SPANDATA.GEN_AI_REQUEST_TOP_K, top_k)
set_on_span(SPANDATA.GEN_AI_REQUEST_TOP_K, top_k)
if top_p is not None and _is_given(top_p):
span.set_data(SPANDATA.GEN_AI_REQUEST_TOP_P, top_p)
set_on_span(SPANDATA.GEN_AI_REQUEST_TOP_P, top_p)

if tools is not None and _is_given(tools) and len(tools) > 0: # type: ignore
span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools))
set_on_span(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools))


def _set_create_input_data(
span: "Span", kwargs: "dict[str, Any]", integration: "AnthropicIntegration"
span: "Union[Span, StreamedSpan]",
kwargs: "dict[str, Any]",
integration: "AnthropicIntegration",
) -> None:
"""
Set input data for the span based on the provided keyword arguments for the anthropic message creation.
"""
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, kwargs.get("stream", False))
if isinstance(span, StreamedSpan):
span.set_attribute(
SPANDATA.GEN_AI_RESPONSE_STREAMING, kwargs.get("stream", False)
)
else:
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, kwargs.get("stream", False))

_set_common_input_data(
span=span,
Expand Down Expand Up @@ -549,7 +564,7 @@ async def _wrap_asynchronous_message_iterator(


def _set_output_data(
span: "Span",
span: "Union[Span, StreamedSpan]",
integration: "AnthropicIntegration",
model: "str | None",
input_tokens: "int | None",
Expand All @@ -562,12 +577,15 @@ def _set_output_data(
) -> None:
"""
Set output data for the span based on the AI response."""
set_on_span = (
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
)
if model is not None:
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, model)
set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, model)
if response_id is not None:
span.set_data(SPANDATA.GEN_AI_RESPONSE_ID, response_id)
set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, response_id)
if finish_reason is not None:
span.set_data(SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason])
set_on_span(SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason])
if should_send_default_pii() and integration.include_prompts:
output_messages: "dict[str, list[Any]]" = {
"response": [],
Expand Down Expand Up @@ -620,12 +638,22 @@ def _sentry_patched_create_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any

model = kwargs.get("model", "")

span = get_start_span_function()(
op=OP.GEN_AI_CHAT,
name=f"chat {model}".strip(),
origin=AnthropicIntegration.origin,
)
span.__enter__()
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
if span_streaming:
span = sentry_sdk.traces.start_span(
name=f"chat {model}".strip(),
attributes={
"sentry.op": OP.GEN_AI_CHAT,
"sentry.origin": AnthropicIntegration.origin,
},
)
else:
span = get_start_span_function()(
op=OP.GEN_AI_CHAT,
name=f"chat {model}".strip(),
origin=AnthropicIntegration.origin,
)
span.__enter__()

_set_create_input_data(span, kwargs, integration)

Expand Down Expand Up @@ -680,10 +708,10 @@ def _sentry_patched_create_sync(f: "Any", *args: "Any", **kwargs: "Any") -> "Any
response_id=getattr(result, "id", None),
finish_reason=getattr(result, "stop_reason", None),
)
span.__exit__(None, None, None)
else:
elif isinstance(span, Span):
Comment thread
alexander-alderman-webb marked this conversation as resolved.
span.set_data("unknown_response", True)
span.__exit__(None, None, None)

span.__exit__(None, None, None)
Comment thread
alexander-alderman-webb marked this conversation as resolved.

return result

Expand All @@ -708,12 +736,22 @@ async def _sentry_patched_create_async(

model = kwargs.get("model", "")

span = get_start_span_function()(
op=OP.GEN_AI_CHAT,
name=f"chat {model}".strip(),
origin=AnthropicIntegration.origin,
)
span.__enter__()
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
if span_streaming:
span = sentry_sdk.traces.start_span(
name=f"chat {model}".strip(),
attributes={
"sentry.op": OP.GEN_AI_CHAT,
"sentry.origin": AnthropicIntegration.origin,
},
)
else:
span = get_start_span_function()(
op=OP.GEN_AI_CHAT,
name=f"chat {model}".strip(),
origin=AnthropicIntegration.origin,
)
span.__enter__()

_set_create_input_data(span, kwargs, integration)

Expand Down Expand Up @@ -768,10 +806,10 @@ async def _sentry_patched_create_async(
response_id=getattr(result, "id", None),
finish_reason=getattr(result, "stop_reason", None),
)
span.__exit__(None, None, None)
else:
elif isinstance(span, Span):
span.set_data("unknown_response", True)
span.__exit__(None, None, None)

span.__exit__(None, None, None)

return result

Expand Down Expand Up @@ -929,7 +967,8 @@ def _sentry_patched_enter(self: "MessageStreamManager") -> "MessageStream":
if not hasattr(self, "_max_tokens"):
return f(self)

integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
client = sentry_sdk.get_client()
integration = client.get_integration(AnthropicIntegration)

if integration is None:
return f(self)
Expand All @@ -942,14 +981,25 @@ def _sentry_patched_enter(self: "MessageStreamManager") -> "MessageStream":
except TypeError:
return f(self)

span = get_start_span_function()(
op=OP.GEN_AI_CHAT,
name="chat" if self._model is None else f"chat {self._model}".strip(),
origin=AnthropicIntegration.origin,
)
span.__enter__()
if has_span_streaming_enabled(client.options):
span = sentry_sdk.traces.start_span(
name="chat" if self._model is None else f"chat {self._model}".strip(),
attributes={
"sentry.op": OP.GEN_AI_CHAT,
"sentry.origin": AnthropicIntegration.origin,
SPANDATA.GEN_AI_RESPONSE_STREAMING: True,
},
)
else:
span = get_start_span_function()(
op=OP.GEN_AI_CHAT,
name="chat" if self._model is None else f"chat {self._model}".strip(),
origin=AnthropicIntegration.origin,
)
span.__enter__()

span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)

span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
_set_common_input_data(
span=span,
integration=integration,
Expand Down Expand Up @@ -1024,7 +1074,8 @@ async def _sentry_patched_aenter(
if not hasattr(self, "_max_tokens"):
return await f(self)

integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
client = sentry_sdk.get_client()
integration = client.get_integration(AnthropicIntegration)

if integration is None:
return await f(self)
Expand All @@ -1037,14 +1088,25 @@ async def _sentry_patched_aenter(
except TypeError:
return await f(self)

span = get_start_span_function()(
op=OP.GEN_AI_CHAT,
name="chat" if self._model is None else f"chat {self._model}".strip(),
origin=AnthropicIntegration.origin,
)
span.__enter__()
if has_span_streaming_enabled(client.options):
span = sentry_sdk.traces.start_span(
name="chat" if self._model is None else f"chat {self._model}".strip(),
attributes={
"sentry.op": OP.GEN_AI_CHAT,
"sentry.origin": AnthropicIntegration.origin,
SPANDATA.GEN_AI_RESPONSE_STREAMING: True,
},
)
else:
span = get_start_span_function()(
op=OP.GEN_AI_CHAT,
name="chat" if self._model is None else f"chat {self._model}".strip(),
origin=AnthropicIntegration.origin,
)
span.__enter__()

span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)

span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
_set_common_input_data(
span=span,
integration=integration,
Expand Down
9 changes: 9 additions & 0 deletions sentry_sdk/tracing_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,15 @@ def has_span_streaming_enabled(options: "Optional[dict[str, Any]]") -> bool:
return (options.get("_experiments") or {}).get("trace_lifecycle") == "stream"


def should_truncate_gen_ai_input(options: "Optional[dict[str, Any]]") -> bool:
if options is None:
return True

return not options.get(
"stream_gen_ai_spans", False
) and not has_span_streaming_enabled(options)


@contextlib.contextmanager
def record_sql_queries(
cursor: "Any",
Expand Down
Loading
Loading