feat: add support for run_stream

Author: constantinius

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -17,10 +17,59 @@
  raise DidNotEnable("pydantic-ai not installed")
 
 
-def _create_run_wrapper(original_func):
- # type: (Callable[..., Any]) -> Callable[..., Any]
+class _StreamingContextManagerWrapper:
+ """Wrapper for streaming methods that return async context managers."""
+
+ def __init__(self, agent, original_ctx_manager, is_streaming=True):
+ # type: (Any, Any, bool) -> None
+ self.agent = agent
+ self.original_ctx_manager = original_ctx_manager
+ self.is_streaming = is_streaming
+ self._isolation_scope = None
+ self._workflow_span = None
+
+ async def __aenter__(self):
+ # type: () -> Any
+ # Set up isolation scope and workflow span
+ self._isolation_scope = sentry_sdk.isolation_scope()
+ self._isolation_scope.__enter__()
+
+ # Store agent reference and streaming flag
+ sentry_sdk.get_current_scope().set_context(
+ "pydantic_ai_agent", {"_agent": self.agent, "_streaming": self.is_streaming}
+ )
+
+ # Create workflow span
+ self._workflow_span = agent_workflow_span(self.agent)
+ self._workflow_span.__enter__()
+
+ # Enter the original context manager
+ result = await self.original_ctx_manager.__aenter__()
+ return result
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
+ # type: (Any, Any, Any) -> None
+ try:
+ # Exit the original context manager first
+ await self.original_ctx_manager.__aexit__(exc_type, exc_val, exc_tb)
+ finally:
+ # Clean up workflow span
+ if self._workflow_span:
+ self._workflow_span.__exit__(exc_type, exc_val, exc_tb)
+
+ # Clean up isolation scope
+ if self._isolation_scope:
+ self._isolation_scope.__exit__(exc_type, exc_val, exc_tb)
+
+
+def _create_run_wrapper(original_func, is_streaming=False):
+ # type: (Callable[..., Any], bool) -> Callable[..., Any]
  """
  Wraps the Agent.run method to create a root span for the agent workflow.
+
+ Args:
+ original_func: The original run method
+ is_streaming: Whether this is a streaming method (for future use)
  """
 
  @wraps(original_func)
@@ -29,10 +78,10 @@ async def wrapper(self, *args, **kwargs):
  # Isolate each workflow so that when agents are run in asyncio tasks they
  # don't touch each other's scopes
  with sentry_sdk.isolation_scope():
- # Store agent reference in Sentry scope for access in nested spans
+ # Store agent reference and streaming flag in Sentry scope for access in nested spans
  # We store the full agent to allow access to tools and system prompts
  sentry_sdk.get_current_scope().set_context(
- "pydantic_ai_agent", {"_agent": self}
+ "pydantic_ai_agent", {"_agent": self, "_streaming": is_streaming}
  )
 
  with agent_workflow_span(self):
@@ -57,6 +106,7 @@ def _create_run_sync_wrapper(original_func):
  # type: (Callable[..., Any]) -> Callable[..., Any]
  """
  Wraps the Agent.run_sync method to create a root span for the agent workflow.
+ Note: run_sync is always non-streaming.
  """
 
  @wraps(original_func)
@@ -65,10 +115,10 @@ def wrapper(self, *args, **kwargs):
  # Isolate each workflow so that when agents are run they
  # don't touch each other's scopes
  with sentry_sdk.isolation_scope():
- # Store agent reference in Sentry scope for access in nested spans
+ # Store agent reference and streaming flag in Sentry scope for access in nested spans
  # We store the full agent to allow access to tools and system prompts
  sentry_sdk.get_current_scope().set_context(
- "pydantic_ai_agent", {"_agent": self}
+ "pydantic_ai_agent", {"_agent": self, "_streaming": False}
  )
 
  with agent_workflow_span(self):
@@ -89,18 +139,84 @@ def wrapper(self, *args, **kwargs):
  return wrapper
 
 
+def _create_streaming_wrapper(original_func):
+ # type: (Callable[..., Any]) -> Callable[..., Any]
+ """
+ Wraps run_stream method that returns an async context manager.
+ """
+
+ @wraps(original_func)
+ def wrapper(self, *args, **kwargs):
+ # type: (Any, *Any, **Any) -> Any
+ # Call original function to get the context manager
+ original_ctx_manager = original_func(self, *args, **kwargs)
+
+ # Wrap it with our instrumentation
+ return _StreamingContextManagerWrapper(
+ agent=self, original_ctx_manager=original_ctx_manager, is_streaming=True
+ )
+
+ return wrapper
+
+
+def _create_streaming_events_wrapper(original_func):
+ # type: (Callable[..., Any]) -> Callable[..., Any]
+ """
+ Wraps run_stream_events method that returns an async generator/iterator.
+ """
+
+ @wraps(original_func)
+ async def wrapper(self, *args, **kwargs):
+ # type: (Any, *Any, **Any) -> Any
+ # Isolate each workflow so that when agents are run in asyncio tasks they
+ # don't touch each other's scopes
+ with sentry_sdk.isolation_scope():
+ # Store agent reference and streaming flag in Sentry scope for access in nested spans
+ sentry_sdk.get_current_scope().set_context(
+ "pydantic_ai_agent", {"_agent": self, "_streaming": True}
+ )
+
+ with agent_workflow_span(self):
+ try:
+ # Call the original generator and yield all events
+ async for event in original_func(self, *args, **kwargs):
+ yield event
+ except Exception as exc:
+ _capture_exception(exc)
+
+ # It could be that there is an "invoke agent" span still open
+ current_span = sentry_sdk.get_current_span()
+ if current_span is not None and current_span.timestamp is None:
+ current_span.__exit__(None, None, None)
+
+ raise exc from None
+
+ return wrapper
+
+
 def _patch_agent_run():
  # type: () -> None
  """
  Patches the Agent run methods to create spans for agent execution.
+
+ This patches both non-streaming (run, run_sync) and streaming
+ (run_stream, run_stream_events) methods.
  """
  # Import here to avoid circular imports
  from pydantic_ai.agent import Agent
 
  # Store original methods
  original_run = Agent.run
  original_run_sync = Agent.run_sync
+ original_run_stream = Agent.run_stream
+ original_run_stream_events = Agent.run_stream_events
 
- # Wrap and apply patches
- Agent.run = _create_run_wrapper(original_run) # type: ignore
+ # Wrap and apply patches for non-streaming methods
+ Agent.run = _create_run_wrapper(original_run, is_streaming=False) # type: ignore
  Agent.run_sync = _create_run_sync_wrapper(original_run_sync) # type: ignore
+
+ # Wrap and apply patches for streaming methods
+ Agent.run_stream = _create_streaming_wrapper(original_run_stream) # type: ignore
+ Agent.run_stream_events = _create_streaming_events_wrapper(
+ original_run_stream_events
+ ) # type: ignore

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -103,6 +103,65 @@ async def wrapped_model_request_run(self, ctx):
 
  ModelRequestNode.run = wrapped_model_request_run # type: ignore
 
+ # Patch ModelRequestNode.stream for streaming requests
+ original_model_request_stream = ModelRequestNode.stream
+
+ def create_wrapped_stream(original_stream_method):
+ # type: (Callable[..., Any]) -> Callable[..., Any]
+ """Create a wrapper for ModelRequestNode.stream that creates chat spans."""
+ from contextlib import asynccontextmanager
+
+ @asynccontextmanager
+ @wraps(original_stream_method)
+ async def wrapped_model_request_stream(self, ctx):
+ # type: (Any, Any) -> Any
+ # Extract data from context
+ model = None
+ model_settings = None
+ if hasattr(ctx, "deps"):
+ model = getattr(ctx.deps, "model", None)
+ model_settings = getattr(ctx.deps, "model_settings", None)
+
+ # Build full message list: history + current request
+ messages = []
+
+ # Add message history
+ if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"):
+ messages.extend(ctx.state.message_history)
+
+ # Add current request
+ current_request = getattr(self, "request", None)
+ if current_request:
+ messages.append(current_request)
+
+ # Create chat span for streaming request
+ import sentry_sdk
+
+ span = ai_client_span(messages, None, model, model_settings)
+ span.__enter__()
+
+ try:
+ # Call the original stream method
+ async with original_stream_method(self, ctx) as stream:
+ yield stream
+
+ # After streaming completes, update span with response data
+ # The ModelRequestNode stores the final response in _result
+ model_response = None
+ if hasattr(self, "_result") and self._result is not None:
+ # _result is a NextNode containing the model_response
+ if hasattr(self._result, "model_response"):
+ model_response = self._result.model_response
+
+ update_ai_client_span(span, model_response)
+ finally:
+ # Close the span after streaming completes
+ span.__exit__(None, None, None)
+
+ return wrapped_model_request_stream
+
+ ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream) # type: ignore
+
  # Patch CallToolsNode to close invoke_agent span when done
  original_call_tools_run = CallToolsNode.run
 

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -45,6 +45,11 @@ def ai_client_span(messages, agent, model, model_settings):
  _set_agent_data(span, agent)
  _set_model_data(span, model, model_settings)
 
+ # Set streaming flag
+ agent_data = sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {}
+ is_streaming = agent_data.get("_streaming", False)
+ span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, is_streaming)
+
  # Add available tools if agent is available
  agent_obj = agent
  if not agent_obj: