sentry-sdk 0.18.0__py2.py3-none-any.whl → 2.46.0__py2.py3-none-any.whl

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -0,0 +1,215 @@
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk.integrations import DidNotEnable
+
+from ..spans import invoke_agent_span, update_invoke_agent_span
+from ..utils import _capture_exception, pop_agent, push_agent
+
+from typing import TYPE_CHECKING
+
+try:
+    from pydantic_ai.agent import Agent  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Optional
+
+
+class _StreamingContextManagerWrapper:
+    """Wrapper for streaming methods that return async context managers."""
+
+    def __init__(
+        self,
+        agent,
+        original_ctx_manager,
+        user_prompt,
+        model,
+        model_settings,
+        is_streaming=True,
+    ):
+        # type: (Any, Any, Any, Any, Any, bool) -> None
+        self.agent = agent
+        self.original_ctx_manager = original_ctx_manager
+        self.user_prompt = user_prompt
+        self.model = model
+        self.model_settings = model_settings
+        self.is_streaming = is_streaming
+        self._isolation_scope = None  # type: Any
+        self._span = None  # type: Optional[sentry_sdk.tracing.Span]
+        self._result = None  # type: Any
+
+    async def __aenter__(self):
+        # type: () -> Any
+        # Set up isolation scope and invoke_agent span
+        self._isolation_scope = sentry_sdk.isolation_scope()
+        self._isolation_scope.__enter__()
+
+        # Create invoke_agent span (will be closed in __aexit__)
+        self._span = invoke_agent_span(
+            self.user_prompt,
+            self.agent,
+            self.model,
+            self.model_settings,
+            self.is_streaming,
+        )
+        self._span.__enter__()
+
+        # Push agent to contextvar stack after span is successfully created and entered
+        # This ensures proper pairing with pop_agent() in __aexit__ even if exceptions occur
+        push_agent(self.agent, self.is_streaming)
+
+        # Enter the original context manager
+        result = await self.original_ctx_manager.__aenter__()
+        self._result = result
+        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)
+
+            # Update span with output if successful
+            if exc_type is None and self._result and hasattr(self._result, "output"):
+                output = (
+                    self._result.output if hasattr(self._result, "output") else None
+                )
+                if self._span is not None:
+                    update_invoke_agent_span(self._span, output)
+        finally:
+            # Pop agent from contextvar stack
+            pop_agent()
+
+            # Clean up invoke span
+            if self._span:
+                self._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 an invoke_agent span.
+
+    Args:
+        original_func: The original run method
+        is_streaming: Whether this is a streaming method (for future use)
+    """
+
+    @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():
+            # Extract parameters for the span
+            user_prompt = kwargs.get("user_prompt") or (args[0] if args else None)
+            model = kwargs.get("model")
+            model_settings = kwargs.get("model_settings")
+
+            # Create invoke_agent span
+            with invoke_agent_span(
+                user_prompt, self, model, model_settings, is_streaming
+            ) as span:
+                # Push agent to contextvar stack after span is successfully created and entered
+                # This ensures proper pairing with pop_agent() in finally even if exceptions occur
+                push_agent(self, is_streaming)
+
+                try:
+                    result = await original_func(self, *args, **kwargs)
+
+                    # Update span with output
+                    output = result.output if hasattr(result, "output") else None
+                    update_invoke_agent_span(span, output)
+
+                    return result
+                except Exception as exc:
+                    _capture_exception(exc)
+                    raise exc from None
+                finally:
+                    # Pop agent from contextvar stack
+                    pop_agent()
+
+    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
+        # Extract parameters for the span
+        user_prompt = kwargs.get("user_prompt") or (args[0] if args else None)
+        model = kwargs.get("model")
+        model_settings = kwargs.get("model_settings")
+
+        # 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,
+            user_prompt=user_prompt,
+            model=model,
+            model_settings=model_settings,
+            is_streaming=True,
+        )
+
+    return wrapper
+
+
+def _create_streaming_events_wrapper(original_func):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+    """
+    Wraps run_stream_events method - no span needed as it delegates to run().
+
+    Note: run_stream_events internally calls self.run() with an event_stream_handler,
+    so the invoke_agent span will be created by the run() wrapper.
+    """
+
+    @wraps(original_func)
+    async def wrapper(self, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        # Just call the original generator - it will call run() which has the instrumentation
+        try:
+            async for event in original_func(self, *args, **kwargs):
+                yield event
+        except Exception as exc:
+            _capture_exception(exc)
+            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.
+    """
+
+    # Store original methods
+    original_run = Agent.run
+    original_run_stream = Agent.run_stream
+    original_run_stream_events = Agent.run_stream_events
+
+    # Wrap and apply patches for non-streaming methods
+    Agent.run = _create_run_wrapper(original_run, is_streaming=False)
+
+    # Wrap and apply patches for streaming methods
+    Agent.run_stream = _create_streaming_wrapper(original_run_stream)
+    Agent.run_stream_events = _create_streaming_events_wrapper(
+        original_run_stream_events
+    )

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -0,0 +1,110 @@
+from contextlib import asynccontextmanager
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk.integrations import DidNotEnable
+
+from ..spans import (
+    ai_client_span,
+    update_ai_client_span,
+)
+
+try:
+    from pydantic_ai._agent_graph import ModelRequestNode  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable
+
+
+def _extract_span_data(node, ctx):
+    # type: (Any, Any) -> tuple[list[Any], Any, Any]
+    """Extract common data needed for creating chat spans.
+
+    Returns:
+        Tuple of (messages, model, model_settings)
+    """
+    # Extract model and settings 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 = []
+    if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"):
+        messages.extend(ctx.state.message_history)
+
+    current_request = getattr(node, "request", None)
+    if current_request:
+        messages.append(current_request)
+
+    return messages, model, model_settings
+
+
+def _patch_graph_nodes():
+    # type: () -> None
+    """
+    Patches the graph node execution to create appropriate spans.
+
+    ModelRequestNode -> Creates ai_client span for model requests
+    CallToolsNode -> Handles tool calls (spans created in tool patching)
+    """
+
+    # Patch ModelRequestNode to create ai_client spans
+    original_model_request_run = ModelRequestNode.run
+
+    @wraps(original_model_request_run)
+    async def wrapped_model_request_run(self, ctx):
+        # type: (Any, Any) -> Any
+        messages, model, model_settings = _extract_span_data(self, ctx)
+
+        with ai_client_span(messages, None, model, model_settings) as span:
+            result = await original_model_request_run(self, ctx)
+
+            # Extract response from result if available
+            model_response = None
+            if hasattr(result, "model_response"):
+                model_response = result.model_response
+
+            update_ai_client_span(span, model_response)
+            return result
+
+    ModelRequestNode.run = wrapped_model_request_run
+
+    # 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."""
+
+        @asynccontextmanager
+        @wraps(original_stream_method)
+        async def wrapped_model_request_stream(self, ctx):
+            # type: (Any, Any) -> Any
+            messages, model, model_settings = _extract_span_data(self, ctx)
+
+            # Create chat span for streaming request
+            with ai_client_span(messages, None, model, model_settings) as span:
+                # 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)
+
+        return wrapped_model_request_stream
+
+    ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream)

sentry_sdk/integrations/pydantic_ai/patches/model_request.py

@@ -0,0 +1,40 @@
+from functools import wraps
+from typing import TYPE_CHECKING
+
+from sentry_sdk.integrations import DidNotEnable
+
+try:
+    from pydantic_ai import models  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+from ..spans import ai_client_span, update_ai_client_span
+
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+def _patch_model_request():
+    # type: () -> None
+    """
+    Patches model request execution to create AI client spans.
+
+    In pydantic-ai, model requests are handled through the Model interface.
+    We need to patch the request method on models to create spans.
+    """
+
+    # Patch the base Model class's request method
+    if hasattr(models, "Model"):
+        original_request = models.Model.request
+
+        @wraps(original_request)
+        async def wrapped_request(self, messages, *args, **kwargs):
+            # type: (Any, Any, *Any, **Any) -> Any
+            # Pass all messages (full conversation history)
+            with ai_client_span(messages, None, self, None) as span:
+                result = await original_request(self, messages, *args, **kwargs)
+                update_ai_client_span(span, result)
+                return result
+
+        models.Model.request = wrapped_request

sentry_sdk/integrations/pydantic_ai/patches/tools.py

@@ -0,0 +1,98 @@
+from functools import wraps
+
+from sentry_sdk.integrations import DidNotEnable
+import sentry_sdk
+
+from ..spans import execute_tool_span, update_execute_tool_span
+from ..utils import (
+    _capture_exception,
+    get_current_agent,
+)
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+try:
+    from pydantic_ai.mcp import MCPServer  # type: ignore
+
+    HAS_MCP = True
+except ImportError:
+    HAS_MCP = False
+
+try:
+    from pydantic_ai._tool_manager import ToolManager  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+
+def _patch_tool_execution():
+    # type: () -> None
+    """
+    Patch ToolManager._call_tool to create execute_tool spans.
+
+    This is the single point where ALL tool calls flow through in pydantic_ai,
+    regardless of toolset type (function, MCP, combined, wrapper, etc.).
+
+    By patching here, we avoid:
+    - Patching multiple toolset classes
+    - Dealing with signature mismatches from instrumented MCP servers
+    - Complex nested toolset handling
+    """
+
+    original_call_tool = ToolManager._call_tool
+
+    @wraps(original_call_tool)
+    async def wrapped_call_tool(self, call, *args, **kwargs):
+        # type: (Any, Any, *Any, **Any) -> Any
+
+        # Extract tool info before calling original
+        name = call.tool_name
+        tool = self.tools.get(name) if self.tools else None
+
+        # Determine tool type by checking tool.toolset
+        tool_type = "function"  # default
+        if tool and HAS_MCP and isinstance(tool.toolset, MCPServer):
+            tool_type = "mcp"
+
+        # Get agent from contextvar
+        agent = get_current_agent()
+
+        if agent and tool:
+            try:
+                args_dict = call.args_as_dict()
+            except Exception:
+                args_dict = call.args if isinstance(call.args, dict) else {}
+
+            # Create execute_tool span
+            # Nesting is handled by isolation_scope() to ensure proper parent-child relationships
+            with sentry_sdk.isolation_scope():
+                with execute_tool_span(
+                    name,
+                    args_dict,
+                    agent,
+                    tool_type=tool_type,
+                ) as span:
+                    try:
+                        result = await original_call_tool(
+                            self,
+                            call,
+                            *args,
+                            **kwargs,
+                        )
+                        update_execute_tool_span(span, result)
+                        return result
+                    except Exception as exc:
+                        _capture_exception(exc)
+                        raise exc from None
+
+        # No span context - just call original
+        return await original_call_tool(
+            self,
+            call,
+            *args,
+            **kwargs,
+        )
+
+    ToolManager._call_tool = wrapped_call_tool

sentry_sdk/integrations/pydantic_ai/spans/__init__.py

@@ -0,0 +1,3 @@
+from .ai_client import ai_client_span, update_ai_client_span  # noqa: F401
+from .execute_tool import execute_tool_span, update_execute_tool_span  # noqa: F401
+from .invoke_agent import invoke_agent_span, update_invoke_agent_span  # noqa: F401