add error hook

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -62,7 +62,6 @@ def setup_once() -> None:
             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:
@@ -72,6 +71,9 @@ def setup_once() -> None:
 
         _patch_tool_execution()
 
+        # Assumptions:
+        # - Model requests within a run are sequential.
+        # - ctx.metadata is a shared dict instance between hooks.
         hooks = Hooks()
 
         @hooks.on.before_model_request  # type: ignore
@@ -113,6 +115,20 @@ async def on_response(
 
             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 isinstance(run_context_metadata, dict):
+                span = run_context_metadata.pop("_sentry_span", None)
+                if span is not None:
+                    span.__exit__(type(error), error, error.__traceback__)
+            raise error
+
         original_init = Agent.__init__
 
         @functools.wraps(original_init)

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -96,7 +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
+    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":
@@ -146,7 +148,9 @@ def _create_streaming_wrapper(
     """
     Wraps run_stream method that returns an async context manager.
     """
-    from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration
+    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":