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

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -1,8 +1,10 @@
-from sentry_sdk.integrations import DidNotEnable, Integration
+import functools
 
+from sentry_sdk.integrations import DidNotEnable, Integration
 
 try:
     import pydantic_ai  # type: ignore # noqa: F401
+    from pydantic_ai import Agent
 except ImportError:
     raise DidNotEnable("pydantic-ai not installed")
 
@@ -14,10 +16,20 @@
     _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
+
 
 class PydanticAIIntegration(Integration):
     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 +57,71 @@
         - Tool executions
         """
         _patch_agent_run()
-        _patch_graph_nodes()
-        _patch_model_request()
+
+        try:
+            from pydantic_ai.capabilities import Hooks  # type: ignore
+        except ImportError:
+            Hooks = None
+            # Save to populate ctx.metadata
+            PydanticAIIntegration.are_request_hooks_available = True
+
+        if Hooks is None:
+            _patch_graph_nodes()
+            _patch_model_request()
+            return
+
         _patch_tool_execution()
+
+        hooks = Hooks()
+
+        @hooks.on.before_model_request  # type: ignore
+        async def on_request(
+            ctx: "RunContext[None]", request_context: "ModelRequestContext"
+        ) -> "ModelRequestContext":
+            span = ai_client_span(
+                messages=request_context.messages,
+                agent=None,
+                model=request_context.model,
+                model_settings=request_context.model_settings,
+            )
+            run_context_metadata = ctx.metadata
+            if isinstance(run_context_metadata, dict):
+                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["_sentry_span"]
+            if span is None:
+                return response
+
+            update_ai_client_span(span, response)
+            span.__exit__(None, None, None)
+            del run_context_metadata["_sentry_span"]
+
+            return response
+
+        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
+            return original_init(self, *args, **kwargs)
+
+        Agent.__init__ = patched_init

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -96,6 +96,7 @@ 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
 
     @wraps(original_func)
     async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
@@ -107,6 +108,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 +146,7 @@ def _create_streaming_wrapper(
     """
     Wraps run_stream method that returns an async context manager.
     """
+    from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration
 
     @wraps(original_func)
     def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
@@ -148,6 +155,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)