feat(pydantic-ai): Support span streaming

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