ref(pydantic-ai): Use hooks.on.run()

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -19,10 +19,8 @@
 from importlib.metadata import PackageNotFoundError, version
 from typing import TYPE_CHECKING
 
-from .patches import (
-    _patch_agent_run,
-)
 from .spans.ai_client import ai_client_span, update_ai_client_span
+from .spans.invoke_agent import invoke_agent_span, update_invoke_agent_span
 
 if TYPE_CHECKING:
     from typing import Any
@@ -35,11 +33,13 @@
         ToolCallPart,
         ToolDefinition,
     )
+    from pydantic_ai.agent import AgentRunResult
     from pydantic_ai.capabilities import (
         Hooks,
         RawToolArgs,
         ValidatedToolArgs,
         WrapModelRequestHandler,
+        WrapRunHandler,
         WrapToolExecuteHandler,
     )
     from pydantic_ai.messages import ModelResponse
@@ -59,7 +59,7 @@ async def on_model_request(
     ) -> "ModelResponse":
         with ai_client_span(
             messages=request_context.messages,
-            agent=None,
+            agent=ctx.agent,
             model=request_context.model,
             model_settings=request_context.model_settings,
         ) as span:
@@ -118,6 +118,28 @@ async def sentry_wrap_tool_execute(
                         _capture_exception(exc, handled=True)
                 reraise(*exc_info)
 
+    @hooks.on.run
+    async def sentry_wrap_run(
+        ctx: "RunContext[Any]",
+        *,
+        handler: "WrapRunHandler",
+    ) -> "AgentRunResult[Any]":
+        with sentry_sdk.isolation_scope(), invoke_agent_span(
+            user_prompt=ctx.prompt,
+            agent=ctx.agent,
+            model=ctx.model,
+            model_settings=ctx.model_settings,
+        ) as span:
+            try:
+                result = await handler()
+                update_invoke_agent_span(span, result)
+                return result
+            except Exception as exc:
+                exc_info = sys.exc_info()
+                with capture_internal_exceptions():
+                    _capture_exception(exc, handled=False)
+                reraise(*exc_info)
+
     original_init = Agent.__init__
 
     @functools.wraps(original_init)
@@ -185,8 +207,6 @@ def setup_once() -> None:
         if PYDANTIC_AI_VERSION is None:
             return
 
-        _patch_agent_run()
-
         try:
             from pydantic_ai.capabilities import Hooks
         except ImportError:

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -18,8 +18,6 @@
     _set_model_data,
     _should_send_inputs,
     _should_send_outputs,
-    get_current_agent,
-    get_is_streaming,
 )
 from .utils import (
     _serialize_binary_content_item,
@@ -30,6 +28,7 @@
 if TYPE_CHECKING:
     from typing import Any, Dict, List, Optional, Union
 
+    from pydantic_ai import Agent
     from pydantic_ai.messages import ModelMessage, ModelResponse, SystemPromptPart
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
@@ -271,7 +270,7 @@ def _set_output_data(
 
 def ai_client_span(
     messages: "list[ModelMessage]",
-    agent: "Any",
+    agent: "Optional[Agent[Any, Any]]",
     model: "Model",
     model_settings: "Optional[ModelSettings]",
 ) -> "StreamedSpan":
@@ -283,29 +282,20 @@ def ai_client_span(
         model: Model object
         model_settings: Model settings
     """
-    # Determine model name for span name
-    model_obj = model
-    if agent and hasattr(agent, "model"):
-        model_obj = agent.model
-
-    model_name = _get_model_name(model_obj) or "unknown"
+    model_name = _get_model_name(model) or "unknown"
 
     span = sentry_sdk.traces.start_span(
         name=f"chat {model_name}",
         attributes={
             "sentry.op": OP.GEN_AI_CHAT,
             "sentry.origin": SPAN_ORIGIN,
             SPANDATA.GEN_AI_OPERATION_NAME: "chat",
-            SPANDATA.GEN_AI_RESPONSE_STREAMING: get_is_streaming(),
         },
     )
 
     _set_agent_data(span, agent)
-    _set_model_data(span, model, model_settings)
-
-    # Add available tools if agent is available
-    agent_obj = agent or get_current_agent()
-    _set_available_tools(span, agent_obj)
+    _set_model_data(span, agent, model, model_settings)
+    _set_available_tools(span, agent)
 
     # Set input messages (full conversation history)
     if messages:

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py

@@ -22,7 +22,12 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Optional, Sequence, Union
+
+    from pydantic_ai import Agent, UserContent
+    from pydantic_ai.models import AbstractModel
+    from pydantic_ai.realtime.settings import RealtimeModelSettings
+    from pydantic_ai.settings import ModelSettings
 
 try:
     from pydantic_ai.messages import BinaryContent, ImageUrl
@@ -32,11 +37,10 @@
 
 
 def invoke_agent_span(
-    user_prompt: "Any",
-    agent: "Any",
-    model: "Any",
-    model_settings: "Any",
-    is_streaming: bool = False,
+    user_prompt: "Optional[Union[str, Sequence[UserContent]]]",
+    agent: "Optional[Agent]",
+    model: "AbstractModel",
+    model_settings: "Optional[Union[ModelSettings, RealtimeModelSettings]]",
 ) -> "StreamedSpan":
     """Create a span for invoking the agent."""
     # Determine agent name for span
@@ -54,7 +58,7 @@ def invoke_agent_span(
     )
 
     _set_agent_data(span, agent)
-    _set_model_data(span, model, model_settings)
+    _set_model_data(span, agent, model, model_settings)
     _set_available_tools(span, agent)
 
     # Add user prompt and system prompts if available and prompts are enabled