feat(pydantic-ai): Support span streaming (#6389)

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -15,7 +15,7 @@
     raise DidNotEnable("pydantic-ai not installed")
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Optional
+    from typing import Any, Callable, Optional, Union
 
 
 class _StreamingContextManagerWrapper:
@@ -37,7 +37,7 @@ def __init__(
         self.model_settings = model_settings
         self.is_streaming = is_streaming
         self._isolation_scope: "Any" = None
-        self._span: "Optional[sentry_sdk.tracing.Span]" = None
+        self._span: "Optional[Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]]" = None
         self._result: "Any" = None
 
     async def __aenter__(self) -> "Any":

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)
                 )
 
@@ -241,7 +261,7 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
 
 def ai_client_span(
     messages: "Any", agent: "Any", model: "Any", model_settings: "Any"
-) -> "sentry_sdk.tracing.Span":
+) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
     """Create a span for an AI client call (model request).
 
     Args:
@@ -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)
@@ -283,7 +314,7 @@ def ai_client_span(
 
 
 def update_ai_client_span(
-    span: "sentry_sdk.tracing.Span", model_response: "Any"
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", model_response: "Any"
 ) -> None:
     """Update the AI client span with response data."""
     if not span:

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
 
@@ -18,7 +20,7 @@ def execute_tool_span(
     tool_args: "Any",
     agent: "Any",
     tool_definition: "Optional[ToolDefinition]" = None,
-) -> "sentry_sdk.tracing.Span":
+) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
     """Create a span for tool execution.
 
     Args:
@@ -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))

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.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 ..consts import SPAN_ORIGIN
 from ..utils import (
@@ -22,7 +24,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Union
 
 try:
     from pydantic_ai.messages import BinaryContent, ImageUrl  # type: ignore
@@ -37,20 +39,31 @@ def invoke_agent_span(
     model: "Any",
     model_settings: "Any",
     is_streaming: bool = False,
-) -> "sentry_sdk.tracing.Span":
+) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
     """Create a span for invoking the agent."""
     # Determine agent name for span
     name = "agent"
     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)
@@ -135,7 +148,10 @@ def invoke_agent_span(
     return span
 
 
-def update_invoke_agent_span(span: "sentry_sdk.tracing.Span", result: "Any") -> None:
+def update_invoke_agent_span(
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
+    result: "Any",
+) -> None:
     """Update and close the invoke agent span."""
     if not span or not result:
         return
@@ -154,7 +170,12 @@ def update_invoke_agent_span(span: "sentry_sdk.tracing.Span", result: "Any") ->
         try:
             response = result.response
             if hasattr(response, "model_name") and response.model_name:
-                span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, 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)
         except Exception:
             # If response access fails, continue without setting model name
             pass