sentry-sdk 2.42.0__py2.py3-none-any.whl → 2.43.0__py2.py3-none-any.whl

sentry_sdk/integrations/openai_agents/patches/error_tracing.py

@@ -0,0 +1,77 @@
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk.consts import SPANSTATUS
+from sentry_sdk.tracing_utils import set_span_errored
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Optional
+
+
+def _patch_error_tracing():
+    # type: () -> None
+    """
+    Patches agents error tracing function to inject our span error logic
+    when a tool execution fails.
+
+    In newer versions, the function is at: agents.util._error_tracing.attach_error_to_current_span
+    In older versions, it was at: agents._utils.attach_error_to_current_span
+
+    This works even when the module or function doesn't exist.
+    """
+    error_tracing_module = None
+
+    # Try newer location first (agents.util._error_tracing)
+    try:
+        from agents.util import _error_tracing
+
+        error_tracing_module = _error_tracing
+    except (ImportError, AttributeError):
+        pass
+
+    # Try older location (agents._utils)
+    if error_tracing_module is None:
+        try:
+            import agents._utils
+
+            error_tracing_module = agents._utils
+        except (ImportError, AttributeError):
+            # Module doesn't exist in either location, nothing to patch
+            return
+
+    # Check if the function exists
+    if not hasattr(error_tracing_module, "attach_error_to_current_span"):
+        return
+
+    original_attach_error = error_tracing_module.attach_error_to_current_span
+
+    @wraps(original_attach_error)
+    def sentry_attach_error_to_current_span(error, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        """
+        Wraps agents' error attachment to also set Sentry span status to error.
+        This allows us to properly track tool execution errors even though
+        the agents library swallows exceptions.
+        """
+        # Set the current Sentry span to errored
+        current_span = sentry_sdk.get_current_span()
+        if current_span is not None:
+            set_span_errored(current_span)
+            current_span.set_data("span.status", "error")
+
+            # Optionally capture the error details if we have them
+            if hasattr(error, "__class__"):
+                current_span.set_data("error.type", error.__class__.__name__)
+            if hasattr(error, "__str__"):
+                error_message = str(error)
+                if error_message:
+                    current_span.set_data("error.message", error_message)
+
+        # Call the original function
+        return original_attach_error(error, *args, **kwargs)
+
+    error_tracing_module.attach_error_to_current_span = (
+        sentry_attach_error_to_current_span
+    )

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -0,0 +1,47 @@
+from sentry_sdk.integrations import DidNotEnable, Integration
+
+
+try:
+    import pydantic_ai  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+
+from .patches import (
+    _patch_agent_run,
+    _patch_graph_nodes,
+    _patch_model_request,
+    _patch_tool_execution,
+)
+
+
+class PydanticAIIntegration(Integration):
+    identifier = "pydantic_ai"
+    origin = f"auto.ai.{identifier}"
+
+    def __init__(self, include_prompts=True):
+        # type: (bool) -> None
+        """
+        Initialize the Pydantic AI integration.
+
+        Args:
+            include_prompts: Whether to include prompts and messages in span data.
+                Requires send_default_pii=True. Defaults to True.
+        """
+        self.include_prompts = include_prompts
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        """
+        Set up the pydantic-ai integration.
+
+        This patches the key methods in pydantic-ai to create Sentry spans for:
+        - Agent invocations (Agent.run methods)
+        - Model requests (AI client calls)
+        - Tool executions
+        """
+        _patch_agent_run()
+        _patch_graph_nodes()
+        _patch_model_request()
+        _patch_tool_execution()

sentry_sdk/integrations/pydantic_ai/consts.py

@@ -0,0 +1 @@
+SPAN_ORIGIN = "auto.ai.pydantic_ai"

sentry_sdk/integrations/pydantic_ai/patches/__init__.py

@@ -0,0 +1,4 @@
+from .agent_run import _patch_agent_run  # noqa: F401
+from .graph_nodes import _patch_graph_nodes  # noqa: F401
+from .model_request import _patch_model_request  # noqa: F401
+from .tools import _patch_tool_execution  # noqa: F401

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -0,0 +1,217 @@
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk.tracing_utils import set_span_errored
+from sentry_sdk.utils import event_from_exception
+
+from ..spans import invoke_agent_span, update_invoke_agent_span
+
+from typing import TYPE_CHECKING
+from pydantic_ai.agent import Agent  # type: ignore
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Optional
+
+
+def _capture_exception(exc):
+    # type: (Any) -> None
+    set_span_errored()
+
+    event, hint = event_from_exception(
+        exc,
+        client_options=sentry_sdk.get_client().options,
+        mechanism={"type": "pydantic_ai", "handled": False},
+    )
+    sentry_sdk.capture_event(event, hint=hint)
+
+
+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__()
+
+        # Store agent reference and streaming flag
+        sentry_sdk.get_current_scope().set_context(
+            "pydantic_ai_agent", {"_agent": self.agent, "_streaming": self.is_streaming}
+        )
+
+        # 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._span.__enter__()
+
+        # 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:
+            sentry_sdk.get_current_scope().remove_context("pydantic_ai_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():
+            # 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, "_streaming": is_streaming}
+            )
+
+            # 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) as span:
+                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:
+                    sentry_sdk.get_current_scope().remove_context("pydantic_ai_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,105 @@
+from contextlib import asynccontextmanager
+from functools import wraps
+
+import sentry_sdk
+
+from ..spans import (
+    ai_client_span,
+    update_ai_client_span,
+)
+from pydantic_ai._agent_graph import ModelRequestNode  # type: ignore
+
+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,35 @@
+from functools import wraps
+from typing import TYPE_CHECKING
+
+from pydantic_ai import models  # type: ignore
+
+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,75 @@
+from functools import wraps
+
+from pydantic_ai._tool_manager import ToolManager  # type: ignore
+
+import sentry_sdk
+
+from ..spans import execute_tool_span, update_execute_tool_span
+
+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
+
+
+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, allow_partial, wrap_validation_errors):
+        # type: (Any, Any, bool, bool) -> 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 Sentry scope
+        current_span = sentry_sdk.get_current_span()
+        if current_span and tool:
+            agent_data = (
+                sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {}
+            )
+            agent = agent_data.get("_agent")
+
+            # Get args for span (before validation)
+            # call.args can be a string (JSON) or dict
+            args_dict = call.args if isinstance(call.args, dict) else {}
+
+            with execute_tool_span(name, args_dict, agent, tool_type=tool_type) as span:
+                result = await original_call_tool(
+                    self, call, allow_partial, wrap_validation_errors
+                )
+                update_execute_tool_span(span, result)
+                return result
+
+        # No span context - just call original
+        return await original_call_tool(
+            self, call, allow_partial, wrap_validation_errors
+        )
+
+    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