ref(pydantic-ai): Use tool hooks

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

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

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.
 

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"]