fix(pydantic-ai): Use first-class hooks when available

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -1,8 +1,11 @@
-from sentry_sdk.integrations import DidNotEnable, Integration
+import functools
 
+from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.utils import capture_internal_exceptions
 
 try:
     import pydantic_ai  # type: ignore # noqa: F401
+    from pydantic_ai import Agent
 except ImportError:
     raise DidNotEnable("pydantic-ai not installed")
 
@@ -14,10 +17,119 @@
     _patch_tool_execution,
 )
 
+from .spans.ai_client import ai_client_span, update_ai_client_span
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+    from pydantic_ai import ModelRequestContext, RunContext
+    from pydantic_ai.messages import ModelResponse  # type: ignore
+    from pydantic_ai.capabilities import Hooks  # type: ignore
+
+
+def register_hooks(hooks: "Hooks") -> None:
+    """
+    Creates hooks for chat model calls and register the hooks by adding the hooks to the `capabilities` argument passed to `Agent.__init__()`.
+    """
+
+    @hooks.on.before_model_request  # type: ignore
+    async def on_request(
+        ctx: "RunContext[None]", request_context: "ModelRequestContext"
+    ) -> "ModelRequestContext":
+        run_context_metadata = ctx.metadata
+        if not isinstance(run_context_metadata, dict):
+            return request_context
+
+        span = ai_client_span(
+            messages=request_context.messages,
+            agent=None,
+            model=request_context.model,
+            model_settings=request_context.model_settings,
+        )
+
+        run_context_metadata["_sentry_span"] = span
+        span.__enter__()
+
+        return request_context
+
+    @hooks.on.after_model_request  # type: ignore
+    async def on_response(
+        ctx: "RunContext[None]",
+        *,
+        request_context: "ModelRequestContext",
+        response: "ModelResponse",
+    ) -> "ModelResponse":
+        run_context_metadata = ctx.metadata
+        if not isinstance(run_context_metadata, dict):
+            return response
+
+        span = run_context_metadata.pop("_sentry_span", None)
+        if span is None:
+            return response
+
+        update_ai_client_span(span, response)
+        span.__exit__(None, None, None)
+
+        return response
+
+    @hooks.on.model_request_error  # type: ignore
+    async def on_error(
+        ctx: "RunContext[None]",
+        *,
+        request_context: "ModelRequestContext",
+        error: "Exception",
+    ) -> "ModelResponse":
+        run_context_metadata = ctx.metadata
+
+        if not isinstance(run_context_metadata, dict):
+            raise error
+
+        span = run_context_metadata.pop("_sentry_span", None)
+        if span is None:
+            raise error
+
+        with capture_internal_exceptions():
+            span.__exit__(type(error), error, error.__traceback__)
+
+        raise error
+
+    original_init = Agent.__init__
+
+    @functools.wraps(original_init)
+    def patched_init(self: "Agent[Any, Any]", *args: "Any", **kwargs: "Any") -> None:
+        caps = list(kwargs.get("capabilities") or [])
+        caps.append(hooks)
+        kwargs["capabilities"] = caps
+
+        metadata = kwargs.get("metadata")
+        if metadata is None:
+            kwargs["metadata"] = {}  # Used as shared reference between hooks
+
+        return original_init(self, *args, **kwargs)
+
+    Agent.__init__ = patched_init
+
 
 class PydanticAIIntegration(Integration):
+    """
+    Typical interaction with the library:
+    1. The user creates an Agent instance with configuration, including system instructions sent to every model call.
+    2. The user calls `Agent.run()` or `Agent.run_stream()` to start an agent run. The latter can be used to incrementally receive progress.
+    - Each run invocation has `RunContext` objects that are passed to the library hooks.
+    3. In a loop, the agent repeatedly calls the model, maintaining a conversation history that includes previous messages and tool results, which is passed to each call.
+
+    Internally, Pydantic AI maintains an execution graph in which ModelRequestNode are responsible for model calls, including retries.
+    Hooks using the decorators provided by `pydantic_ai.capabilities` create and manage spans for model calls when these hooks are available (newer library versions).
+    The span is created in `on_request` and stored in the metadata of the `RunContext` object shared with `on_response` and `on_error`.
+
+    The metadata dictionary on the RunContext instance is initialized with `{"_sentry_span": None}` in the `_create_run_wrapper()` and `_create_streaming_wrapper()` wrappers that
+    instrument `Agent.run()` and `Agent.run_stream()`, respectively. A non-empty dictionary is required for the metadata object to be a shared reference between hooks.
+    """
+
     identifier = "pydantic_ai"
     origin = f"auto.ai.{identifier}"
+    are_request_hooks_available = True
 
     def __init__(
         self, include_prompts: bool = True, handled_tool_call_exceptions: bool = True
@@ -45,6 +157,18 @@ def setup_once() -> None:
         - Tool executions
         """
         _patch_agent_run()
-        _patch_graph_nodes()
-        _patch_model_request()
         _patch_tool_execution()
+
+        try:
+            from pydantic_ai.capabilities import Hooks
+        except ImportError:
+            Hooks = None
+            PydanticAIIntegration.are_request_hooks_available = False
+
+        if Hooks is None:
+            _patch_graph_nodes()
+            _patch_model_request()
+            return
+
+        hooks = Hooks()
+        register_hooks(hooks)

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -96,6 +96,9 @@ def _create_run_wrapper(
         original_func: The original run method
         is_streaming: Whether this is a streaming method (for future use)
     """
+    from sentry_sdk.integrations.pydantic_ai import (
+        PydanticAIIntegration,
+    )  # Required to avoid circular import
 
     @wraps(original_func)
     async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
@@ -107,6 +110,11 @@ async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
             model = kwargs.get("model")
             model_settings = kwargs.get("model_settings")
 
+            if PydanticAIIntegration.are_request_hooks_available:
+                metadata = kwargs.get("metadata")
+                if not metadata:
+                    kwargs["metadata"] = {"_sentry_span": None}
+
             # Create invoke_agent span
             with invoke_agent_span(
                 user_prompt, self, model, model_settings, is_streaming
@@ -140,6 +148,9 @@ def _create_streaming_wrapper(
     """
     Wraps run_stream method that returns an async context manager.
     """
+    from sentry_sdk.integrations.pydantic_ai import (
+        PydanticAIIntegration,
+    )  # Required to avoid circular import
 
     @wraps(original_func)
     def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
@@ -148,6 +159,11 @@ def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
         model = kwargs.get("model")
         model_settings = kwargs.get("model_settings")
 
+        if PydanticAIIntegration.are_request_hooks_available:
+            metadata = kwargs.get("metadata")
+            if not metadata:
+                kwargs["metadata"] = {"_sentry_span": None}
+
         # Call original function to get the context manager
         original_ctx_manager = original_func(self, *args, **kwargs)