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

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

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -0,0 +1,253 @@
+import sentry_sdk
+from sentry_sdk.ai.utils import set_data_normalized
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.utils import safe_serialize
+
+from ..consts import SPAN_ORIGIN
+from ..utils import (
+    _set_agent_data,
+    _set_available_tools,
+    _set_model_data,
+    _should_send_prompts,
+    _get_model_name,
+)
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, List, Dict
+    from pydantic_ai.usage import RequestUsage  # type: ignore
+
+try:
+    from pydantic_ai.messages import (  # type: ignore
+        BaseToolCallPart,
+        BaseToolReturnPart,
+        SystemPromptPart,
+        UserPromptPart,
+        TextPart,
+        ThinkingPart,
+    )
+except ImportError:
+    # Fallback if these classes are not available
+    BaseToolCallPart = None
+    BaseToolReturnPart = None
+    SystemPromptPart = None
+    UserPromptPart = None
+    TextPart = None
+    ThinkingPart = None
+
+
+def _set_usage_data(span, usage):
+    # type: (sentry_sdk.tracing.Span, RequestUsage) -> None
+    """Set token usage data on a span."""
+    if usage is None:
+        return
+
+    if hasattr(usage, "input_tokens") and usage.input_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
+
+    if hasattr(usage, "output_tokens") and usage.output_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens)
+
+    if hasattr(usage, "total_tokens") and usage.total_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)
+
+
+def _set_input_messages(span, messages):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Set input messages data on a span."""
+    if not _should_send_prompts():
+        return
+
+    if not messages:
+        return
+
+    try:
+        formatted_messages = []
+        system_prompt = None
+
+        # Extract system prompt from any ModelRequest with instructions
+        for msg in messages:
+            if hasattr(msg, "instructions") and msg.instructions:
+                system_prompt = msg.instructions
+                break
+
+        # Add system prompt as first message if present
+        if system_prompt:
+            formatted_messages.append(
+                {"role": "system", "content": [{"type": "text", "text": system_prompt}]}
+            )
+
+        for msg in messages:
+            if hasattr(msg, "parts"):
+                for part in msg.parts:
+                    role = "user"
+                    # Use isinstance checks with proper base classes
+                    if SystemPromptPart and isinstance(part, SystemPromptPart):
+                        role = "system"
+                    elif (
+                        (TextPart and isinstance(part, TextPart))
+                        or (ThinkingPart and isinstance(part, ThinkingPart))
+                        or (BaseToolCallPart and isinstance(part, BaseToolCallPart))
+                    ):
+                        role = "assistant"
+                    elif BaseToolReturnPart and isinstance(part, BaseToolReturnPart):
+                        role = "tool"
+
+                    content = []  # type: List[Dict[str, Any] | str]
+                    tool_calls = None
+                    tool_call_id = None
+
+                    # Handle ToolCallPart (assistant requesting tool use)
+                    if BaseToolCallPart and isinstance(part, BaseToolCallPart):
+                        tool_call_data = {}
+                        if hasattr(part, "tool_name"):
+                            tool_call_data["name"] = part.tool_name
+                        if hasattr(part, "args"):
+                            tool_call_data["arguments"] = safe_serialize(part.args)
+                        if tool_call_data:
+                            tool_calls = [tool_call_data]
+                    # Handle ToolReturnPart (tool result)
+                    elif BaseToolReturnPart and isinstance(part, BaseToolReturnPart):
+                        if hasattr(part, "tool_name"):
+                            tool_call_id = part.tool_name
+                        if hasattr(part, "content"):
+                            content.append({"type": "text", "text": str(part.content)})
+                    # Handle regular content
+                    elif hasattr(part, "content"):
+                        if isinstance(part.content, str):
+                            content.append({"type": "text", "text": part.content})
+                        elif isinstance(part.content, list):
+                            for item in part.content:
+                                if isinstance(item, str):
+                                    content.append({"type": "text", "text": item})
+                                else:
+                                    content.append(safe_serialize(item))
+                        else:
+                            content.append({"type": "text", "text": str(part.content)})
+
+                    # Add message if we have content or tool calls
+                    if content or tool_calls:
+                        message = {"role": role}  # type: Dict[str, Any]
+                        if content:
+                            message["content"] = content
+                        if tool_calls:
+                            message["tool_calls"] = tool_calls
+                        if tool_call_id:
+                            message["tool_call_id"] = tool_call_id
+                        formatted_messages.append(message)
+
+        if formatted_messages:
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, formatted_messages, unpack=False
+            )
+    except Exception:
+        # If we fail to format messages, just skip it
+        pass
+
+
+def _set_output_data(span, response):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Set output data on a span."""
+    if not _should_send_prompts():
+        return
+
+    if not response:
+        return
+
+    set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name)
+    try:
+        # Extract text from ModelResponse
+        if hasattr(response, "parts"):
+            texts = []
+            tool_calls = []
+
+            for part in response.parts:
+                if TextPart and isinstance(part, TextPart) and hasattr(part, "content"):
+                    texts.append(part.content)
+                elif BaseToolCallPart and isinstance(part, BaseToolCallPart):
+                    tool_call_data = {
+                        "type": "function",
+                    }
+                    if hasattr(part, "tool_name"):
+                        tool_call_data["name"] = part.tool_name
+                    if hasattr(part, "args"):
+                        tool_call_data["arguments"] = safe_serialize(part.args)
+                    tool_calls.append(tool_call_data)
+
+            if texts:
+                set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts)
+
+            if tool_calls:
+                span.set_data(
+                    SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)
+                )
+
+    except Exception:
+        # If we fail to format output, just skip it
+        pass
+
+
+def ai_client_span(messages, agent, model, model_settings):
+    # type: (Any, Any, Any, Any) -> sentry_sdk.tracing.Span
+    """Create a span for an AI client call (model request).
+
+    Args:
+        messages: Full conversation history (list of messages)
+        agent: Agent object
+        model: Model object
+        model_settings: Model settings
+    """
+    # Determine model name for span name
+    model_obj = model
+    if agent and hasattr(agent, "model"):
+        model_obj = agent.model
+
+    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.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+
+    _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:
+        # Try to get from Sentry scope
+        agent_data = (
+            sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {}
+        )
+        agent_obj = agent_data.get("_agent")
+
+    _set_available_tools(span, agent_obj)
+
+    # Set input messages (full conversation history)
+    if messages:
+        _set_input_messages(span, messages)
+
+    return span
+
+
+def update_ai_client_span(span, model_response):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Update the AI client span with response data."""
+    if not span:
+        return
+
+    # Set usage data if available
+    if model_response and hasattr(model_response, "usage"):
+        _set_usage_data(span, model_response.usage)
+
+    # Set output data
+    _set_output_data(span, model_response)