ref(pydantic-ai): Adopt shared span-data helper for dual span APIs

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}]),
                 )

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,18 +54,19 @@ 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,
         )
 
     _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))

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

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
+        )

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