fix(pydantic-ai): Set system instructions after instructions are set on the request object

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/patches/_agent_run.py

@@ -0,0 +1,278 @@
+import sys
+import itertools
+from collections.abc import Sequence
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.integrations.pydantic_ai.utils import _should_send_prompts
+from sentry_sdk.utils import capture_internal_exceptions, reraise
+from sentry_sdk.consts import SPANDATA
+from sentry_sdk.ai import set_data_normalized
+
+from ..spans import invoke_agent_span, update_invoke_agent_span
+from ..utils import _capture_exception, pop_agent, push_agent
+
+from typing import TYPE_CHECKING
+
+try:
+    from pydantic_ai.agent import Agent  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Optional, Union
+
+    from pydantic_ai.agent.abstract import Instructions
+    from pydantic_ai.tools import SystemPromptFunc
+
+    from sentry_sdk._types import TextPart
+
+
+def _get_current_system_instructions(
+    agent_creation_time_instructions: "list[Union[str, SystemPromptFunc]]",
+    system_instructions: "Optional[Instructions[Any]]",
+) -> "list[str]":
+    instruction_texts: "list[str]" = []
+
+    instruction_texts += [
+        instruction
+        for instruction in agent_creation_time_instructions
+        if isinstance(instruction, str)
+    ]
+
+    if isinstance(system_instructions, str):
+        instruction_texts.append(system_instructions)
+
+    elif isinstance(system_instructions, Sequence):
+        instruction_texts += [
+            instruction
+            for instruction in system_instructions
+            if isinstance(instruction, str)
+        ]
+
+    return instruction_texts
+
+
+def _transform_system_instructions(
+    permanent_instructions: "list[str]", current_instructions: "list[str]"
+) -> "list[TextPart]":
+    return [
+        {
+            "type": "text",
+            "content": instruction,
+        }
+        for instruction in itertools.chain(permanent_instructions, current_instructions)
+    ]
+
+
+class _StreamingContextManagerWrapper:
+    """Wrapper for streaming methods that return async context managers."""
+
+    def __init__(
+        self,
+        agent: "Any",
+        original_ctx_manager: "Any",
+        user_prompt: "Any",
+        model: "Any",
+        model_settings: "Any",
+        is_streaming: bool = True,
+    ) -> None:
+        self.agent = agent
+        self.original_ctx_manager = original_ctx_manager
+        self.user_prompt = user_prompt
+        self.model = model
+        self.model_settings = model_settings
+        self.is_streaming = is_streaming
+        self._isolation_scope: "Any" = None
+        self._span: "Optional[sentry_sdk.tracing.Span]" = None
+        self._result: "Any" = None
+
+    async def __aenter__(self) -> "Any":
+        # Set up isolation scope and invoke_agent span
+        self._isolation_scope = sentry_sdk.isolation_scope()
+        self._isolation_scope.__enter__()
+
+        # Create invoke_agent span (will be closed in __aexit__)
+        self._span = invoke_agent_span(
+            self.user_prompt,
+            self.agent,
+            self.model,
+            self.model_settings,
+            self.is_streaming,
+        )
+        self._span.__enter__()
+
+        # Push agent to contextvar stack after span is successfully created and entered
+        # This ensures proper pairing with pop_agent() in __aexit__ even if exceptions occur
+        push_agent(self.agent, self.is_streaming)
+
+        # Enter the original context manager
+        result = await self.original_ctx_manager.__aenter__()
+        self._result = result
+        return result
+
+    async def __aexit__(self, exc_type: "Any", exc_val: "Any", exc_tb: "Any") -> None:
+        try:
+            # Exit the original context manager first
+            await self.original_ctx_manager.__aexit__(exc_type, exc_val, exc_tb)
+
+            # Update span with result if successful
+            if exc_type is None and self._result and self._span is not None:
+                update_invoke_agent_span(self._span, self._result)
+        finally:
+            # Pop agent from contextvar stack
+            pop_agent()
+
+            # Clean up invoke span
+            if self._span:
+                self._span.__exit__(exc_type, exc_val, exc_tb)
+
+            # Clean up isolation scope
+            if self._isolation_scope:
+                self._isolation_scope.__exit__(exc_type, exc_val, exc_tb)
+
+
+def _create_run_wrapper(
+    original_func: "Callable[..., Any]", is_streaming: bool = False
+) -> "Callable[..., Any]":
+    """
+    Wraps the Agent.run method to create an invoke_agent span.
+
+    Args:
+        original_func: The original run method
+        is_streaming: Whether this is a streaming method (for future use)
+    """
+
+    @wraps(original_func)
+    async def wrapper(self: "Agent", *args: "Any", **kwargs: "Any") -> "Any":
+        # Isolate each workflow so that when agents are run in asyncio tasks they
+        # don't touch each other's scopes
+        with sentry_sdk.isolation_scope():
+            # Extract parameters for the span
+            user_prompt = kwargs.get("user_prompt") or (args[0] if args else None)
+            model = kwargs.get("model")
+            model_settings = kwargs.get("model_settings")
+
+            permanent_instructions = list(self._system_prompts)
+
+            instructions: "Optional[Instructions[Any]]" = kwargs.get("instructions")
+            current_instructions = _get_current_system_instructions(
+                self._instructions, instructions
+            )
+
+            # Create invoke_agent span
+            with invoke_agent_span(
+                user_prompt, self, model, model_settings, is_streaming
+            ) as span:
+                if _should_send_prompts() and (
+                    len(permanent_instructions) > 0 or len(current_instructions) > 0
+                ):
+                    set_data_normalized(
+                        span,
+                        SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
+                        _transform_system_instructions(
+                            permanent_instructions, current_instructions
+                        ),
+                        unpack=False,
+                    )
+
+                # Push agent to contextvar stack after span is successfully created and entered
+                # This ensures proper pairing with pop_agent() in finally even if exceptions occur
+                push_agent(self, is_streaming)
+
+                try:
+                    result = await original_func(self, *args, **kwargs)
+
+                    # Update span with result
+                    update_invoke_agent_span(span, result)
+
+                    return result
+                except Exception as exc:
+                    exc_info = sys.exc_info()
+                    with capture_internal_exceptions():
+                        _capture_exception(exc)
+                    reraise(*exc_info)
+                finally:
+                    # Pop agent from contextvar stack
+                    pop_agent()
+
+    return wrapper
+
+
+def _create_streaming_wrapper(
+    original_func: "Callable[..., Any]",
+) -> "Callable[..., Any]":
+    """
+    Wraps run_stream method that returns an async context manager.
+    """
+
+    @wraps(original_func)
+    def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
+        # Extract parameters for the span
+        user_prompt = kwargs.get("user_prompt") or (args[0] if args else None)
+        model = kwargs.get("model")
+        model_settings = kwargs.get("model_settings")
+
+        # Call original function to get the context manager
+        original_ctx_manager = original_func(self, *args, **kwargs)
+
+        # Wrap it with our instrumentation
+        return _StreamingContextManagerWrapper(
+            agent=self,
+            original_ctx_manager=original_ctx_manager,
+            user_prompt=user_prompt,
+            model=model,
+            model_settings=model_settings,
+            is_streaming=True,
+        )
+
+    return wrapper
+
+
+def _create_streaming_events_wrapper(
+    original_func: "Callable[..., Any]",
+) -> "Callable[..., Any]":
+    """
+    Wraps run_stream_events method - no span needed as it delegates to run().
+
+    Note: run_stream_events internally calls self.run() with an event_stream_handler,
+    so the invoke_agent span will be created by the run() wrapper.
+    """
+
+    @wraps(original_func)
+    async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
+        # Just call the original generator - it will call run() which has the instrumentation
+        try:
+            async for event in original_func(self, *args, **kwargs):
+                yield event
+        except Exception as exc:
+            exc_info = sys.exc_info()
+            with capture_internal_exceptions():
+                _capture_exception(exc)
+            reraise(*exc_info)
+
+    return wrapper
+
+
+def _patch_agent_run() -> None:
+    """
+    Patches the Agent run methods to create spans for agent execution.
+
+    This patches both non-streaming (run, run_sync) and streaming
+    (run_stream, run_stream_events) methods.
+    """
+
+    # Store original methods
+    original_run = Agent.run
+    original_run_stream = Agent.run_stream
+    original_run_stream_events = Agent.run_stream_events
+
+    # Wrap and apply patches for non-streaming methods
+    Agent.run = _create_run_wrapper(original_run, is_streaming=False)
+
+    # Wrap and apply patches for streaming methods
+    Agent.run_stream = _create_streaming_wrapper(original_run_stream)
+    Agent.run_stream_events = _create_streaming_events_wrapper(
+        original_run_stream_events
+    )

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,12 @@ 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)
 
+            if messages:
+                _set_input_messages(span, messages)
+
             # Extract response from result if available
             model_response = None
             if hasattr(result, "model_response"):
@@ -86,7 +90,10 @@ 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:
+                if messages:
+                    _set_input_messages(span, messages)
+
                 # Call the original stream method
                 async with original_stream_method(self, ctx) as stream:
                     yield stream

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