fix(pydantic-ai): Use first-class hooks when available

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -7,6 +7,7 @@
     ai_client_span,
     update_ai_client_span,
 )
+from ..utils import _set_input_messages
 
 try:
     from pydantic_ai._agent_graph import ModelRequestNode  # type: ignore
@@ -59,9 +60,15 @@ 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(messages, None, model, model_settings) as span:
+        with ai_client_span(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"):
@@ -86,9 +93,16 @@ 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(messages, None, model, model_settings) as span:
+            with ai_client_span(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,6 +3,8 @@
 
 from sentry_sdk.integrations import DidNotEnable
 
+from ..utils import _set_input_messages
+
 try:
     from pydantic_ai import models  # type: ignore
 except ImportError:
@@ -32,7 +34,10 @@ async def wrapped_request(
             self: "Any", messages: "Any", *args: "Any", **kwargs: "Any"
         ) -> "Any":
             # Pass all messages (full conversation history)
-            with ai_client_span(messages, None, self, None) as span:
+            with ai_client_span(None, self, None) as span:
+                if messages:
+                    _set_input_messages(span, messages)
+
                 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,10 +1,6 @@
-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
@@ -20,178 +16,23 @@
     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, List, Dict
-    from pydantic_ai.messages import ModelMessage, SystemPromptPart  # type: ignore
-    from sentry_sdk._types import TextPart as SentryTextPart
+    from typing import Any
 
 try:
-    from pydantic_ai.messages import (
+    from pydantic_ai.messages import (  # type: ignore
         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:
@@ -236,7 +77,7 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None:
 
 
 def ai_client_span(
-    messages: "Any", agent: "Any", model: "Any", model_settings: "Any"
+    agent: "Any", model: "Any", model_settings: "Any"
 ) -> "sentry_sdk.tracing.Span":
     """Create a span for an AI client call (model request).
 
@@ -271,10 +112,6 @@ 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,
-)
-from .utils import (
     _serialize_binary_content_item,
     _serialize_image_url_item,
+)
+from .utils import (
     _set_usage_data,
 )
 

sentry_sdk/integrations/pydantic_ai/spans/utils.py

@@ -1,50 +1,15 @@
 """Utility functions for PydanticAI span instrumentation."""
 
 import sentry_sdk
-from sentry_sdk._types import BLOB_DATA_SUBSTITUTE
-from sentry_sdk.ai.utils import get_modality_from_mime_type
 from sentry_sdk.consts import SPANDATA
 
-from ..consts import DATA_URL_BASE64_REGEX
-
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Union, Dict, Any
+    from typing import Union
     from pydantic_ai.usage import RequestUsage, RunUsage  # type: ignore
 
 
-def _serialize_image_url_item(item: "Any") -> "Dict[str, Any]":
-    """Serialize an ImageUrl content item for span data.
-
-    For data URLs containing base64-encoded images, the content is redacted.
-    For regular HTTP URLs, the URL string is preserved.
-    """
-    url = str(item.url)
-    data_url_match = DATA_URL_BASE64_REGEX.match(url)
-
-    if data_url_match:
-        return {
-            "type": "image",
-            "content": BLOB_DATA_SUBSTITUTE,
-        }
-
-    return {
-        "type": "image",
-        "content": url,
-    }
-
-
-def _serialize_binary_content_item(item: "Any") -> "Dict[str, Any]":
-    """Serialize a BinaryContent item for span data, redacting the blob data."""
-    return {
-        "type": "blob",
-        "modality": get_modality_from_mime_type(item.media_type),
-        "mime_type": item.media_type,
-        "content": BLOB_DATA_SUBSTITUTE,
-    }
-
-
 def _set_usage_data(
     span: "sentry_sdk.tracing.Span", usage: "Union[RequestUsage, RunUsage]"
 ) -> None: