diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index 7bd15a04df..2988301791 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -1,11 +1,17 @@ import functools +import sys +import sentry_sdk from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version -from sentry_sdk.utils import parse_version +from sentry_sdk.utils import capture_internal_exceptions, parse_version, reraise + +from .spans import execute_tool_span, update_execute_tool_span +from .utils import _capture_exception try: import pydantic_ai # noqa: F401 from pydantic_ai import Agent + from pydantic_ai.exceptions import ToolRetryError except ImportError: raise DidNotEnable("pydantic-ai not installed") @@ -15,15 +21,27 @@ from .patches import ( _patch_agent_run, - _patch_tool_execution, ) from .spans.ai_client import ai_client_span, update_ai_client_span if TYPE_CHECKING: from typing import Any - from pydantic_ai import ModelRequestContext, RunContext - from pydantic_ai.capabilities import Hooks, WrapModelRequestHandler + from pydantic import ValidationError + from pydantic_ai import ( + ModelRequestContext, + ModelRetry, + RunContext, + ToolCallPart, + ToolDefinition, + ) + from pydantic_ai.capabilities import ( + Hooks, + RawToolArgs, + ValidatedToolArgs, + WrapModelRequestHandler, + WrapToolExecuteHandler, + ) from pydantic_ai.messages import ModelResponse @@ -50,6 +68,56 @@ async def on_model_request( update_ai_client_span(span, response) return response + @hooks.on.tool_validate_error + async def sentry_on_tool_validate_error( + ctx: "RunContext[Any]", + *, + call: "ToolCallPart", + tool_def: "ToolDefinition", + args: "RawToolArgs", + error: "ValidationError | ModelRetry", + ) -> "ValidatedToolArgs": + with capture_internal_exceptions(): + integration = sentry_sdk.get_client().get_integration( + PydanticAIIntegration, + ) + if integration is not None and integration.handled_tool_call_exceptions: + _capture_exception(error, handled=True) + + raise error + + @hooks.on.tool_execute + async def sentry_wrap_tool_execute( + ctx: "RunContext[Any]", + *, + call: "ToolCallPart", + tool_def: "ToolDefinition", + args: "ValidatedToolArgs", + handler: "WrapToolExecuteHandler", + ) -> "Any": + with execute_tool_span( + tool_name=call.tool_name, + tool_args=args, + agent=ctx.agent, + tool_definition=tool_def, + ) as span: + try: + result = await handler(args) + update_execute_tool_span(span, result) + return result + except ToolRetryError as exc: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + integration = sentry_sdk.get_client().get_integration( + PydanticAIIntegration, + ) + if ( + integration is not None + and integration.handled_tool_call_exceptions + ): + _capture_exception(exc, handled=True) + reraise(*exc_info) + original_init = Agent.__init__ @functools.wraps(original_init) @@ -118,7 +186,6 @@ def setup_once() -> None: return _patch_agent_run() - _patch_tool_execution() try: from pydantic_ai.capabilities import Hooks diff --git a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py index ad6c22216b..0130dc1c4c 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/__init__.py @@ -1,2 +1 @@ from .agent_run import _patch_agent_run # noqa: F401 -from .tools import _patch_tool_execution # noqa: F401 diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py deleted file mode 100644 index 21d1135300..0000000000 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ /dev/null @@ -1,100 +0,0 @@ -import sys -from contextlib import nullcontext -from functools import wraps -from typing import TYPE_CHECKING - -import sentry_sdk -from sentry_sdk.integrations import DidNotEnable -from sentry_sdk.utils import capture_internal_exceptions, reraise - -from ..spans import execute_tool_span, update_execute_tool_span -from ..utils import _capture_exception, get_current_agent - -if TYPE_CHECKING: - from typing import Any - -try: - try: - from pydantic_ai.tool_manager import ToolManager - except ImportError: - from pydantic_ai._tool_manager import ToolManager # type: ignore - - from pydantic_ai.exceptions import ToolRetryError -except ImportError: - raise DidNotEnable("pydantic-ai not installed") - - -def _patch_tool_execution() -> None: - if hasattr(ToolManager, "execute_tool_call"): - _patch_execute_tool_call() - - -def _patch_execute_tool_call() -> None: - original_execute_tool_call = ToolManager.execute_tool_call - - @wraps(original_execute_tool_call) - async def wrapped_execute_tool_call( - self: "Any", validated: "Any", *args: "Any", **kwargs: "Any" - ) -> "Any": - if not validated or not hasattr(validated, "call"): - return await original_execute_tool_call(self, validated, *args, **kwargs) - - # Extract tool info before calling original - call = validated.call - name = call.tool_name - tool = self.tools.get(name) if self.tools else None - selected_tool_definition = getattr(tool, "tool_def", None) - - # 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_definition=selected_tool_definition, - ) - if validated.args_valid - else nullcontext() - ) as span: - try: - result = await original_execute_tool_call( - self, - validated, - *args, - **kwargs, - ) - if span is not None: - update_execute_tool_span(span, result) - return result - except ToolRetryError as exc: - exc_info = sys.exc_info() - with capture_internal_exceptions(): - # Avoid circular import due to multi-file integration structure - from sentry_sdk.integrations.pydantic_ai import ( - PydanticAIIntegration, - ) - - integration = sentry_sdk.get_client().get_integration( - PydanticAIIntegration - ) - if ( - integration is not None - and integration.handled_tool_call_exceptions - ): - _capture_exception(exc, handled=True) - reraise(*exc_info) - - return await original_execute_tool_call(self, validated, *args, **kwargs) - - ToolManager.execute_tool_call = wrapped_execute_tool_call # type: ignore[method-assign] diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py index 7a8974f072..045367ae88 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py @@ -11,14 +11,15 @@ if TYPE_CHECKING: from typing import Any, Optional + from pydantic_ai import Agent from pydantic_ai._tool_manager import ToolDefinition # type: ignore def execute_tool_span( tool_name: str, - tool_args: "Any", - agent: "Any", - tool_definition: "Optional[ToolDefinition]" = None, + tool_args: "dict[str, Any]", + agent: "Optional[Agent[Any, Any]]", + tool_definition: "ToolDefinition", ) -> "StreamedSpan": """Create a span for tool execution. diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 96e324c374..15280cf4d9 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -549,10 +549,8 @@ def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int: assert result is None if handled_tool_call_exceptions: - ( - error, - model_behaviour_error, - ) = (item.payload for item in items if item.type == "event") + events = [item.payload for item in items if item.type == "event"] + error = events[0] assert error["level"] == "error" assert error["exception"]["values"][0]["mechanism"]["handled"]