ref(pydantic-ai): Replace individual hooks with hooks.on.model_request (#7327)

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -1,7 +1,7 @@
 import functools
 
 from sentry_sdk.integrations import DidNotEnable, Integration
-from sentry_sdk.utils import capture_internal_exceptions, parse_version
+from sentry_sdk.utils import parse_version
 
 try:
     import pydantic_ai  # noqa: F401
@@ -24,7 +24,7 @@
     from typing import Any
 
     from pydantic_ai import ModelRequestContext, RunContext
-    from pydantic_ai.capabilities import Hooks
+    from pydantic_ai.capabilities import Hooks, WrapModelRequestHandler
     from pydantic_ai.messages import ModelResponse
 
 
@@ -33,67 +33,24 @@ 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
-    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
-    async def on_response(
+    @hooks.on.model_request
+    async def on_model_request(
         ctx: "RunContext[None]",
         *,
         request_context: "ModelRequestContext",
-        response: "ModelResponse",
+        handler: "WrapModelRequestHandler",
     ) -> "ModelResponse":
-        run_context_metadata = ctx.metadata
-        if not isinstance(run_context_metadata, dict):
-            return response
+        with ai_client_span(
+            messages=request_context.messages,
+            agent=None,
+            model=request_context.model,
+            model_settings=request_context.model_settings,
+        ) as span:
+            response = await handler(request_context)
 
-        span = run_context_metadata.pop("_sentry_span", None)
-        if span is None:
+            update_ai_client_span(span, response)
             return response
 
-        update_ai_client_span(span, response)
-        span.__exit__(None, None, None)
-
-        return response
-
-    @hooks.on.model_request_error
-    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)
@@ -102,10 +59,6 @@ def patched_init(self: "Agent[Any, Any]", *args: "Any", **kwargs: "Any") -> None
         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  # type: ignore[method-assign]
@@ -129,7 +82,6 @@ class PydanticAIIntegration(Integration):
 
     identifier = "pydantic_ai"
     origin = f"auto.ai.{identifier}"
-    using_request_hooks = False
 
     def __init__(
         self, include_prompts: bool = True, handled_tool_call_exceptions: bool = True
@@ -159,7 +111,6 @@ def setup_once() -> None:
         _patch_agent_run()
         _patch_tool_execution()
 
-        PydanticAIIntegration.using_request_hooks = False
         try:
             PYDANTIC_AI_VERSION = version("pydantic-ai-slim")
         except PackageNotFoundError:
@@ -182,6 +133,5 @@ def setup_once() -> None:
         except ImportError:
             return
 
-        PydanticAIIntegration.using_request_hooks = True
         hooks = Hooks()
         register_hooks(hooks)

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -95,9 +95,6 @@ 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":
@@ -109,11 +106,6 @@ async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
             model = kwargs.get("model")
             model_settings = kwargs.get("model_settings")
 
-            if PydanticAIIntegration.using_request_hooks:
-                metadata = kwargs.get("metadata")
-                if metadata is None:
-                    kwargs["metadata"] = {"_sentry_span": None}
-
             # Create invoke_agent span
             with invoke_agent_span(
                 user_prompt, self, model, model_settings, is_streaming
@@ -147,9 +139,6 @@ 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":
@@ -158,11 +147,6 @@ def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
         model = kwargs.get("model")
         model_settings = kwargs.get("model_settings")
 
-        if PydanticAIIntegration.using_request_hooks:
-            metadata = kwargs.get("metadata")
-            if metadata is None:
-                kwargs["metadata"] = {"_sentry_span": None}
-
         # Call original function to get the context manager
         original_ctx_manager = original_func(self, *args, **kwargs)