use hooks

Author: alexander-alderman-webb

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,6 +16,14 @@
     _patch_tool_execution,
 )
 
+from .spans.ai_client import ai_client_span, update_ai_client_span
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pydantic_ai import ModelRequestContext, RunContext
+    from pydantic_ai.messages import ModelResponse
+
 
 class PydanticAIIntegration(Integration):
     identifier = "pydantic_ai"
@@ -45,6 +55,57 @@ def setup_once() -> None:
         - Tool executions
         """
         _patch_agent_run()
-        _patch_graph_nodes()
-        _patch_model_request()
+
+        try:
+            from pydantic_ai.capabilities import Hooks
+
+            hooks = Hooks()
+
+            @hooks.on.before_model_request
+            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,
+                )
+                ctx.metadata["_sentry_span"] = span
+                span.__enter__()
+
+                return request_context
+
+            @hooks.on.after_model_request
+            async def on_response(
+                ctx: "RunContext[None]",
+                *,
+                request_context: "ModelRequestContext",
+                response: "ModelResponse",
+            ) -> "ModelResponse":
+                span = ctx.metadata["_sentry_span"]
+                if span is None:
+                    return response
+
+                update_ai_client_span(span, response)
+                span.__exit__(None, None, None)
+                del ctx.metadata["_sentry_span"]
+
+                return response
+
+            original_init = Agent.__init__
+
+            @functools.wraps(original_init)
+            def patched_init(self, *args, **kwargs):
+                caps = list(kwargs.get("capabilities") or [])
+                caps.append(hooks)
+                kwargs["capabilities"] = caps
+                original_init(self, *args, **kwargs)
+
+            Agent.__init__ = patched_init
+
+        except ImportError:
+            _patch_graph_nodes()
+            _patch_model_request()
+
         _patch_tool_execution()

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -107,6 +107,10 @@ async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
             model = kwargs.get("model")
             model_settings = kwargs.get("model_settings")
 
+            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,6 +152,10 @@ def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
         model = kwargs.get("model")
         model_settings = kwargs.get("model_settings")
 
+        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)
 

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -7,7 +7,6 @@
     ai_client_span,
     update_ai_client_span,
 )
-from ..utils import _set_input_messages
 
 try:
     from pydantic_ai._agent_graph import ModelRequestNode  # type: ignore
@@ -60,15 +59,9 @@ def _patch_graph_nodes() -> None:
     async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any":
         messages, model, model_settings = _extract_span_data(self, ctx)
 
-        with ai_client_span(None, model, model_settings) as span:
+        with ai_client_span(messages, None, model, model_settings) as span:
             result = await original_model_request_run(self, ctx)
 
-            # The instructions are added in `_prepare_request` that runs as part of `ModelRequestNode.run`, so the input
-            # must be recorded after the call. See _get_instructions() added with
-            # https://github.com/pydantic/pydantic-ai/commit/f5271434a56c7a3bb5a3c93f2d1236d8b18afe3e
-            if messages:
-                _set_input_messages(span, messages)
-
             # Extract response from result if available
             model_response = None
             if hasattr(result, "model_response"):
@@ -93,16 +86,9 @@ async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any":
             messages, model, model_settings = _extract_span_data(self, ctx)
 
             # Create chat span for streaming request
-            with ai_client_span(None, model, model_settings) as span:
+            with ai_client_span(messages, None, model, model_settings) as span:
                 # Call the original stream method
                 async with original_stream_method(self, ctx) as stream:
-                    # The instructions are added in `_prepare_request` that runs as part of __aenter__ on the
-                    # context manager returned by `ModelRequestNode.stream()`, so the input must be recorded after the
-                    # call. See _get_instructions() added with
-                    # https://github.com/pydantic/pydantic-ai/commit/f5271434a56c7a3bb5a3c93f2d1236d8b18afe3e
-                    if messages:
-                        _set_input_messages(span, messages)
-
                     yield stream
 
                 # After streaming completes, update span with response data

sentry_sdk/integrations/pydantic_ai/patches/model_request.py

@@ -3,8 +3,6 @@
 
 from sentry_sdk.integrations import DidNotEnable
 
-from ..utils import _set_input_messages
-
 try:
     from pydantic_ai import models  # type: ignore
 except ImportError:
@@ -34,10 +32,7 @@ async def wrapped_request(
             self: "Any", messages: "Any", *args: "Any", **kwargs: "Any"
         ) -> "Any":
             # Pass all messages (full conversation history)
-            with ai_client_span(None, self, None) as span:
-                if messages:
-                    _set_input_messages(span, messages)
-
+            with ai_client_span(messages, None, self, None) as span:
                 result = await original_request(self, messages, *args, **kwargs)
                 update_ai_client_span(span, result)
                 return result

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -1,6 +1,10 @@
+import json
+
 import sentry_sdk
 from sentry_sdk.ai.utils import (
+    normalize_message_roles,
     set_data_normalized,
+    truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.utils import safe_serialize
@@ -16,23 +20,178 @@
     get_is_streaming,
 )
 from .utils import (
+    _serialize_binary_content_item,
+    _serialize_image_url_item,
     _set_usage_data,
 )
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, List, Dict
+    from pydantic_ai.messages import ModelMessage, SystemPromptPart  # type: ignore
+    from sentry_sdk._types import TextPart as SentryTextPart
 
 try:
-    from pydantic_ai.messages import (  # type: ignore
+    from pydantic_ai.messages import (
         BaseToolCallPart,
+        BaseToolReturnPart,
+        SystemPromptPart,
+        UserPromptPart,
         TextPart,
+        ThinkingPart,
+        BinaryContent,
+        ImageUrl,
     )
 except ImportError:
     # Fallback if these classes are not available
     BaseToolCallPart = None
+    BaseToolReturnPart = None
+    SystemPromptPart = None
+    UserPromptPart = None
     TextPart = None
+    ThinkingPart = None
+    BinaryContent = None
+    ImageUrl = None
+
+
+def _transform_system_instructions(
+    permanent_instructions: "list[SystemPromptPart]",
+    current_instructions: "list[str]",
+) -> "list[SentryTextPart]":
+    text_parts: "list[SentryTextPart]" = [
+        {
+            "type": "text",
+            "content": instruction.content,
+        }
+        for instruction in permanent_instructions
+    ]
+
+    text_parts.extend(
+        {
+            "type": "text",
+            "content": instruction,
+        }
+        for instruction in current_instructions
+    )
+
+    return text_parts
+
+
+def _get_system_instructions(
+    messages: "list[ModelMessage]",
+) -> "tuple[list[SystemPromptPart], list[str]]":
+    permanent_instructions = []
+    current_instructions = []
+
+    for msg in messages:
+        if hasattr(msg, "parts"):
+            for part in msg.parts:
+                if SystemPromptPart and isinstance(part, SystemPromptPart):
+                    permanent_instructions.append(part)
+
+        if hasattr(msg, "instructions") and msg.instructions is not None:
+            current_instructions.append(msg.instructions)
+
+    return permanent_instructions, current_instructions
+
+
+def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> None:
+    """Set input messages data on a span."""
+    if not _should_send_prompts():
+        return
+
+    if not messages:
+        return
+
+    permanent_instructions, current_instructions = _get_system_instructions(messages)
+    if len(permanent_instructions) > 0 or len(current_instructions) > 0:
+        span.set_data(
+            SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
+            json.dumps(
+                _transform_system_instructions(
+                    permanent_instructions, current_instructions
+                )
+            ),
+        )
+
+    try:
+        formatted_messages = []
+
+        for msg in messages:
+            if hasattr(msg, "parts"):
+                for part in msg.parts:
+                    role = "user"
+                    # Use isinstance checks with proper base classes
+                    if SystemPromptPart and isinstance(part, SystemPromptPart):
+                        continue
+                    elif (
+                        (TextPart and isinstance(part, TextPart))
+                        or (ThinkingPart and isinstance(part, ThinkingPart))
+                        or (BaseToolCallPart and isinstance(part, BaseToolCallPart))
+                    ):
+                        role = "assistant"
+                    elif BaseToolReturnPart and isinstance(part, BaseToolReturnPart):
+                        role = "tool"
+
+                    content: "List[Dict[str, Any] | str]" = []
+                    tool_calls = None
+                    tool_call_id = None
+
+                    # Handle ToolCallPart (assistant requesting tool use)
+                    if BaseToolCallPart and isinstance(part, BaseToolCallPart):
+                        tool_call_data = {}
+                        if hasattr(part, "tool_name"):
+                            tool_call_data["name"] = part.tool_name
+                        if hasattr(part, "args"):
+                            tool_call_data["arguments"] = safe_serialize(part.args)
+                        if tool_call_data:
+                            tool_calls = [tool_call_data]
+                    # Handle ToolReturnPart (tool result)
+                    elif BaseToolReturnPart and isinstance(part, BaseToolReturnPart):
+                        if hasattr(part, "tool_name"):
+                            tool_call_id = part.tool_name
+                        if hasattr(part, "content"):
+                            content.append({"type": "text", "text": str(part.content)})
+                    # Handle regular content
+                    elif hasattr(part, "content"):
+                        if isinstance(part.content, str):
+                            content.append({"type": "text", "text": part.content})
+                        elif isinstance(part.content, list):
+                            for item in part.content:
+                                if isinstance(item, str):
+                                    content.append({"type": "text", "text": item})
+                                elif ImageUrl and isinstance(item, ImageUrl):
+                                    content.append(_serialize_image_url_item(item))
+                                elif BinaryContent and isinstance(item, BinaryContent):
+                                    content.append(_serialize_binary_content_item(item))
+                                else:
+                                    content.append(safe_serialize(item))
+                        else:
+                            content.append({"type": "text", "text": str(part.content)})
+                    # Add message if we have content or tool calls
+                    if content or tool_calls:
+                        message: "Dict[str, Any]" = {"role": role}
+                        if content:
+                            message["content"] = content
+                        if tool_calls:
+                            message["tool_calls"] = tool_calls
+                        if tool_call_id:
+                            message["tool_call_id"] = tool_call_id
+                        formatted_messages.append(message)
+
+        if formatted_messages:
+            normalized_messages = normalize_message_roles(formatted_messages)
+            scope = sentry_sdk.get_current_scope()
+            messages_data = truncate_and_annotate_messages(
+                normalized_messages, span, scope
+            )
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
+            )
+    except Exception:
+        # If we fail to format messages, just skip it
+        pass
 
 
 def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
@@ -77,7 +236,7 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
 
 
 def ai_client_span(
-    agent: "Any", model: "Any", model_settings: "Any"
+    messages: "Any", agent: "Any", model: "Any", model_settings: "Any"
 ) -> "sentry_sdk.tracing.Span":
     """Create a span for an AI client call (model request).
 
@@ -112,6 +271,10 @@ def ai_client_span(
     agent_obj = agent or get_current_agent()
     _set_available_tools(span, agent_obj)
 
+    # Set input messages (full conversation history)
+    if messages:
+        _set_input_messages(span, messages)
+
     return span
 
 

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py

@@ -13,10 +13,10 @@
     _set_available_tools,
     _set_model_data,
     _should_send_prompts,
-    _serialize_binary_content_item,
-    _serialize_image_url_item,
 )
 from .utils import (
+    _serialize_binary_content_item,
+    _serialize_image_url_item,
     _set_usage_data,
 )