feat(pydantic-ai): Support span streaming

alexander-alderman-webb · alexander-alderman-webb · commit b94997495ce9 · 2026-05-22T10:13:34.000+02:00
diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py
@@ -8,6 +8,8 @@
     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
 from sentry_sdk.utils import safe_serialize
 
 from ..consts import SPAN_ORIGIN
@@ -27,7 +29,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Dict, List
+    from typing import Any, Dict, List, Union
 
     from pydantic_ai.messages import ModelMessage, SystemPromptPart  # type: ignore
 
@@ -97,7 +99,9 @@ def _get_system_instructions(
     return permanent_instructions, current_instructions
 
 
-def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> None:
+def _set_input_messages(
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", messages: "Any"
+) -> None:
     """Set input messages data on a span."""
     if not _should_send_prompts():
         return
@@ -107,14 +111,24 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non
 
     permanent_instructions, current_instructions = _get_system_instructions(messages)
     if len(permanent_instructions) > 0 or len(current_instructions) > 0:
-        span.set_data(
-            SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
-            json.dumps(
-                _transform_system_instructions(
-                    permanent_instructions, current_instructions
-                )
-            ),
-        )
+        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
+                    )
+                ),
+            )
 
     try:
         formatted_messages = []
@@ -198,15 +212,21 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non
         pass
 
 
-def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
+def _set_output_data(
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", response: "Any"
+) -> None:
     """Set output data on a span."""
     if not _should_send_prompts():
         return
 
     if not response:
         return
 
-    span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name)
+    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)
+
     try:
         # Extract text from ModelResponse
         if hasattr(response, "parts"):
@@ -230,7 +250,7 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
                 set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts)
 
             if tool_calls:
-                span.set_data(
+                set_on_span(
                     SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)
                 )
 
@@ -257,20 +277,31 @@ def ai_client_span(
 
     model_name = _get_model_name(model_obj) or "unknown"
 
-    span = sentry_sdk.start_span(
-        op=OP.GEN_AI_CHAT,
-        name=f"chat {model_name}",
-        origin=SPAN_ORIGIN,
-    )
+    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    if span_streaming:
+        span = sentry_sdk.traces.start_span(
+            name=f"chat {model_name}",
+            attributes={
+                "sentry.op": OP.GEN_AI_CHAT,
+                "sentry.origin": SPAN_ORIGIN,
+                SPANDATA.GEN_AI_OPERATION_NAME: "chat",
+                SPANDATA.GEN_AI_RESPONSE_STREAMING: get_is_streaming(),
+            },
+        )
+    else:
+        span = sentry_sdk.start_span(
+            op=OP.GEN_AI_CHAT,
+            name=f"chat {model_name}",
+            origin=SPAN_ORIGIN,
+        )
 
-    span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+        # Set streaming flag from contextvar
+        span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming())
 
     _set_agent_data(span, agent)
     _set_model_data(span, model, model_settings)
 
-    # Set streaming flag from contextvar
-    span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming())
-
     # Add available tools if agent is available
     agent_obj = agent or get_current_agent()
     _set_available_tools(span, agent_obj)
diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py
@@ -2,13 +2,15 @@
 
 import sentry_sdk
 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
 
 from ..consts import SPAN_ORIGIN
 from ..utils import _set_agent_data, _should_send_prompts
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Optional, Union
 
     from pydantic_ai._tool_manager import ToolDefinition  # type: ignore
 
@@ -27,33 +29,56 @@ def execute_tool_span(
         agent: The agent executing the tool
         tool_definition: The definition of the tool, if available
     """
-    span = sentry_sdk.start_span(
-        op=OP.GEN_AI_EXECUTE_TOOL,
-        name=f"execute_tool {tool_name}",
-        origin=SPAN_ORIGIN,
-    )
+    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    if span_streaming:
+        span = sentry_sdk.traces.start_span(
+            name=f"execute_tool {tool_name}",
+            attributes={
+                "sentry.op": OP.GEN_AI_EXECUTE_TOOL,
+                "sentry.origin": SPAN_ORIGIN,
+                SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool",
+                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,
+            name=f"execute_tool {tool_name}",
+            origin=SPAN_ORIGIN,
+        )
 
-    span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
-    span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name)
+        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"):
-        span.set_data(
+        set_on_span(
             SPANDATA.GEN_AI_TOOL_DESCRIPTION,
             tool_definition.description,
         )
 
     _set_agent_data(span, agent)
 
     if _should_send_prompts() and tool_args is not None:
-        span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args))
+        set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args))
 
     return span
 
 
-def update_execute_tool_span(span: "sentry_sdk.tracing.Span", result: "Any") -> None:
+def update_execute_tool_span(
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", result: "Any"
+) -> None:
     """Update the execute tool span with the result."""
     if not span:
         return
 
-    if _should_send_prompts() and result is not None:
+    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))
diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py
@@ -8,6 +8,7 @@
     truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 from ..consts import SPAN_ORIGIN
 from ..utils import (
@@ -44,13 +45,24 @@ def invoke_agent_span(
     if agent and getattr(agent, "name", None):
         name = agent.name
 
-    span = get_start_span_function()(
-        op=OP.GEN_AI_INVOKE_AGENT,
-        name=f"invoke_agent {name}",
-        origin=SPAN_ORIGIN,
-    )
+    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    if span_streaming:
+        span = sentry_sdk.traces.start_span(
+            name=f"invoke_agent {name}",
+            attributes={
+                "sentry.op": OP.GEN_AI_INVOKE_AGENT,
+                "sentry.origin": SPAN_ORIGIN,
+                SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent",
+            },
+        )
+    else:
+        span = get_start_span_function()(
+            op=OP.GEN_AI_INVOKE_AGENT,
+            name=f"invoke_agent {name}",
+            origin=SPAN_ORIGIN,
+        )
 
-    span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
 
     _set_agent_data(span, agent)
     _set_model_data(span, model, model_settings)
diff --git a/sentry_sdk/integrations/pydantic_ai/spans/utils.py b/sentry_sdk/integrations/pydantic_ai/spans/utils.py
@@ -7,6 +7,7 @@
 from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX
 from sentry_sdk.ai.utils import 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
@@ -46,7 +47,8 @@ def _serialize_binary_content_item(item: "Any") -> "Dict[str, Any]":
 
 
 def _set_usage_data(
-    span: "sentry_sdk.tracing.Span", usage: "Union[RequestUsage, RunUsage]"
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
+    usage: "Union[RequestUsage, RunUsage]",
 ) -> None:
     """Set token usage data on a span.
 
@@ -60,24 +62,26 @@ 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:
-        span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
+        set_on_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:
-        span.set_data(
-            SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens
-        )
+        set_on_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:
-        span.set_data(
+        set_on_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:
-        span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens)
+        set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens)
 
     if hasattr(usage, "total_tokens") and usage.total_tokens is not None:
-        span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)
+        set_on_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
diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py
