set metadata in agent.__init__

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -101,6 +101,11 @@ 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 not metadata:
+            kwargs["metadata"] = {}  # Used as shared reference between hooks
+
         return original_init(self, *args, **kwargs)
 
     Agent.__init__ = patched_init
@@ -118,13 +123,12 @@ class PydanticAIIntegration(Integration):
     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.
+    The metadata on the RunContext instance is initialized with an empty dictionary in `Agent.__init__()`. The 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
@@ -158,7 +162,6 @@ def setup_once() -> None:
             from pydantic_ai.capabilities import Hooks
         except ImportError:
             Hooks = None
-            PydanticAIIntegration.are_request_hooks_available = False
 
         if Hooks is None:
             _patch_graph_nodes()

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -96,9 +96,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":
@@ -110,11 +107,6 @@ 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
@@ -148,9 +140,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":
@@ -159,11 +148,6 @@ 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)