diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index c89b8609d8..2991889059 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -3,12 +3,12 @@ import sentry_sdk from sentry_sdk.ai.utils import ( + _set_span_data_attribute, normalize_message_roles, set_data_normalized, truncate_and_annotate_messages, ) from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing_utils import ( has_span_streaming_enabled, should_truncate_gen_ai_input, @@ -37,6 +37,7 @@ from pydantic_ai.messages import ModelMessage, ModelResponse, SystemPromptPart from sentry_sdk import _types + from sentry_sdk.traces import StreamedSpan try: from pydantic_ai.messages import ( @@ -115,24 +116,15 @@ def _set_input_messages( permanent_instructions, current_instructions = _get_system_instructions(messages) if len(permanent_instructions) > 0 or len(current_instructions) > 0: - if isinstance(span, StreamedSpan): - span.set_attribute( - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - json.dumps( - _transform_system_instructions( - permanent_instructions, current_instructions - ) - ), - ) - else: - span.set_data( - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - json.dumps( - _transform_system_instructions( - permanent_instructions, current_instructions - ) - ), - ) + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps( + _transform_system_instructions( + permanent_instructions, current_instructions + ) + ), + ) try: formatted_messages = [] @@ -242,10 +234,10 @@ def _set_output_data( if not response: return - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) # type: ignore[arg-type] + if response.model_name: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name + ) try: if hasattr(response, "parts"): @@ -278,7 +270,8 @@ def _set_output_data( parts.append(tool_part) if parts: - set_on_span( + _set_span_data_attribute( + span, SPANDATA.GEN_AI_OUTPUT_MESSAGES, json.dumps([{"role": "assistant", "parts": parts}]), ) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py index 7648c1418a..bb4cfd8cc1 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py @@ -1,8 +1,8 @@ from typing import TYPE_CHECKING import sentry_sdk +from sentry_sdk.ai.utils import _set_span_data_attribute from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import safe_serialize @@ -14,6 +14,8 @@ from pydantic_ai._tool_manager import ToolDefinition # type: ignore + from sentry_sdk.traces import StreamedSpan + def execute_tool_span( tool_name: str, @@ -31,6 +33,8 @@ def execute_tool_span( """ span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) if span_streaming: + # Both keys must be present at span start so that attribute-based + # ignore_spans / traces_sampler rules can match this span. span = sentry_sdk.traces.start_span( name=f"execute_tool {tool_name}", attributes={ @@ -40,8 +44,6 @@ def execute_tool_span( SPANDATA.GEN_AI_TOOL_NAME: tool_name, }, ) - - set_on_span = span.set_attribute else: span = sentry_sdk.start_span( op=OP.GEN_AI_EXECUTE_TOOL, @@ -52,10 +54,9 @@ def execute_tool_span( span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool") span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) - set_on_span = span.set_data - if tool_definition is not None and hasattr(tool_definition, "description"): - set_on_span( + _set_span_data_attribute( + span, SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_definition.description, ) @@ -63,7 +64,9 @@ def execute_tool_span( _set_agent_data(span, agent) if _should_send_prompts() and tool_args is not None: - set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args)) + _set_span_data_attribute( + span, SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args) + ) return span @@ -78,7 +81,4 @@ def update_execute_tool_span( if not _should_send_prompts() or result is None: return - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) - else: - span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) + _set_span_data_attribute(span, SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index 876cb1ecea..f5cc5c4f97 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -2,13 +2,13 @@ import sentry_sdk from sentry_sdk.ai.utils import ( + _set_span_data_attribute, get_start_span_function, normalize_message_roles, set_data_normalized, truncate_and_annotate_messages, ) from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing_utils import ( has_span_streaming_enabled, should_truncate_gen_ai_input, @@ -29,6 +29,8 @@ if TYPE_CHECKING: from typing import Any, Union + from sentry_sdk.traces import StreamedSpan + try: from pydantic_ai.messages import BinaryContent, ImageUrl except ImportError: @@ -173,12 +175,9 @@ def update_invoke_agent_span( try: response = result.response if hasattr(response, "model_name") and response.model_name: - if isinstance(span, StreamedSpan): - span.set_attribute( - SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name - ) - else: - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) + _set_span_data_attribute( + span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name + ) except Exception: # If response access fails, continue without setting model name pass diff --git a/sentry_sdk/integrations/pydantic_ai/spans/utils.py b/sentry_sdk/integrations/pydantic_ai/spans/utils.py index cbce0c0c77..ae4a899f4b 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/utils.py @@ -5,15 +5,16 @@ import sentry_sdk from sentry_sdk._types import BLOB_DATA_SUBSTITUTE from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX -from sentry_sdk.ai.utils import get_modality_from_mime_type +from sentry_sdk.ai.utils import _set_span_data_attribute, get_modality_from_mime_type from sentry_sdk.consts import SPANDATA -from sentry_sdk.traces import StreamedSpan if TYPE_CHECKING: from typing import Any, Dict, Union from pydantic_ai.usage import RequestUsage, RunUsage + from sentry_sdk.traces import StreamedSpan + def _serialize_image_url_item(item: "Any") -> "Dict[str, Any]": """Serialize an ImageUrl content item for span data. @@ -62,26 +63,31 @@ def _set_usage_data( if usage is None: return - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - if hasattr(usage, "input_tokens") and usage.input_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) + _set_span_data_attribute( + span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens + ) # Pydantic AI uses cache_read_tokens (not input_tokens_cached) if hasattr(usage, "cache_read_tokens") and usage.cache_read_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens) + _set_span_data_attribute( + span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens + ) # Pydantic AI uses cache_write_tokens (not input_tokens_cache_write) if hasattr(usage, "cache_write_tokens") and usage.cache_write_tokens is not None: - set_on_span( + _set_span_data_attribute( + span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, usage.cache_write_tokens, ) if hasattr(usage, "output_tokens") and usage.output_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) + _set_span_data_attribute( + span, SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens + ) if hasattr(usage, "total_tokens") and usage.total_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens) + _set_span_data_attribute( + span, SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens + ) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 340dcf8953..560e4715fb 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -2,14 +2,16 @@ from typing import TYPE_CHECKING import sentry_sdk +from sentry_sdk.ai.utils import _set_span_data_attribute from sentry_sdk.consts import SPANDATA from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.traces import StreamedSpan from sentry_sdk.utils import event_from_exception, safe_serialize if TYPE_CHECKING: from typing import Any, Optional, Union + from sentry_sdk.traces import StreamedSpan + # Store the current agent context in a contextvar for re-entrant safety # Using a list as a stack to support nested agent calls @@ -85,10 +87,7 @@ def _set_agent_data( agent_obj = get_current_agent() if agent_obj and hasattr(agent_obj, "name") and agent_obj.name: - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) - else: - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) + _set_span_data_attribute(span, SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) def _get_model_name(model_obj: "Any") -> "Optional[str]": @@ -136,19 +135,15 @@ def _set_model_data( if not model_obj and agent_obj and hasattr(agent_obj, "model"): model_obj = agent_obj.model - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - if model_obj: # Set system from model if hasattr(model_obj, "system"): - set_on_span(SPANDATA.GEN_AI_SYSTEM, model_obj.system) + _set_span_data_attribute(span, SPANDATA.GEN_AI_SYSTEM, model_obj.system) # Set model name model_name = _get_model_name(model_obj) if model_name: - set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model_name) # Extract model settings settings = model_settings @@ -169,14 +164,14 @@ def _set_model_data( for setting_name, spandata_key in settings_map.items(): value = settings.get(setting_name) if value is not None: - set_on_span(spandata_key, value) + _set_span_data_attribute(span, spandata_key, value) else: # Fallback for object-style settings for setting_name, spandata_key in settings_map.items(): if hasattr(settings, setting_name): value = getattr(settings, setting_name) if value is not None: - set_on_span(spandata_key, value) + _set_span_data_attribute(span, spandata_key, value) def _set_available_tools( @@ -211,14 +206,9 @@ def _set_available_tools( tools.append(tool_info) if tools: - if isinstance(span, StreamedSpan): - span.set_attribute( - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) - ) - else: - span.set_data( - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) - ) + _set_span_data_attribute( + span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) except Exception: # If we can't extract tools, just skip it pass