merge

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -1,103 +1,16 @@
-from functools import wraps
-from contextlib import asynccontextmanager
-from contextvars import ContextVar
 from sentry_sdk.integrations import DidNotEnable, Integration
-from sentry_sdk.consts import SPANDATA
 
 try:
     import pydantic_ai  # type: ignore # noqa: F401
-    from pydantic_ai.capabilities.combined import CombinedCapability  # type: ignore
-    from pydantic_ai._agent_graph import ModelRequestNode
 except ImportError:
     raise DidNotEnable("pydantic-ai not installed")
 
 
 from .patches import (
     _patch_agent_run,
+    _patch_graph_nodes,
     _patch_tool_execution,
 )
-from .spans import (
-    ai_client_span,
-    update_ai_client_span,
-)
-
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import Any, Awaitable, Callable
-
-    from pydantic_ai._run_context import RunContext
-    from pydantic_ai.models import ModelRequestContext
-    from pydantic_ai.messages import ModelResponse
-
-
-_is_streaming: ContextVar[bool] = ContextVar(
-    "sentry_pydantic_ai_is_streaming", default=False
-)
-
-
-def _patch_wrap_model_request():
-    original_wrap_model_request = CombinedCapability.wrap_model_request
-
-    @wraps(original_wrap_model_request)
-    async def wrapped_wrap_model_request(
-        self,
-        ctx: "RunContext[Any]",
-        *,
-        request_context: "ModelRequestContext",
-        handler: "Callable[[ModelRequestContext], Awaitable[ModelResponse]]",
-    ) -> "Any":
-        with ai_client_span(
-            request_context.messages,
-            None,
-            request_context.model,
-            request_context.model_settings,
-        ) as span:
-            span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, _is_streaming.get())
-
-            result = await original_wrap_model_request(
-                self, ctx, request_context=request_context, handler=handler
-            )
-
-            update_ai_client_span(span, result)
-            return result
-
-    CombinedCapability.wrap_model_request = wrapped_wrap_model_request
-
-
-def _patch_model_request_node_run():
-    original_model_request_run = ModelRequestNode.run
-
-    @wraps(original_model_request_run)
-    async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any":
-        token = _is_streaming.set(False)
-        try:
-            return await original_model_request_run(self, ctx)
-        finally:
-            _is_streaming.reset(token)
-
-    ModelRequestNode.run = wrapped_model_request_run
-
-
-def _patch_model_request_node_stream():
-    original_model_request_stream = ModelRequestNode.stream
-
-    def create_wrapped_stream(
-        original_stream_method: "Callable[..., Any]",
-    ) -> "Callable[..., Any]":
-        @asynccontextmanager
-        @wraps(original_stream_method)
-        async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any":
-            token = _is_streaming.set(True)
-            try:
-                async with original_stream_method(self, ctx) as stream:
-                    yield stream
-            finally:
-                _is_streaming.reset(token)
-
-        return wrapped_model_request_stream
-
-    ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream)
 
 
 class PydanticAIIntegration(Integration):
@@ -130,8 +43,5 @@ def setup_once() -> None:
         - Tool executions
         """
         _patch_agent_run()
-        _patch_wrap_model_request()
+        _patch_graph_nodes()
         _patch_tool_execution()
-
-        _patch_model_request_node_run()
-        _patch_model_request_node_stream()

sentry_sdk/integrations/pydantic_ai/patches/__init__.py

@@ -1,2 +1,3 @@
 from .agent_run import _patch_agent_run  # noqa: F401
+from .graph_nodes import _patch_graph_nodes  # noqa: F401
 from .tools import _patch_tool_execution  # noqa: F401

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -0,0 +1,121 @@
+from contextlib import asynccontextmanager
+from functools import wraps
+
+from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.consts import SPANDATA
+
+from ..spans import (
+    ai_client_span,
+    update_ai_client_span,
+)
+
+try:
+    from pydantic_ai._agent_graph import ModelRequestNode  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable
+
+
+def _extract_span_data(node: "Any", ctx: "Any") -> "tuple[list[Any], Any, Any]":
+    """Extract common data needed for creating chat spans.
+
+    Returns:
+        Tuple of (messages, model, model_settings)
+    """
+    # Extract model and settings from context
+    model = None
+    model_settings = None
+    if hasattr(ctx, "deps"):
+        model = getattr(ctx.deps, "model", None)
+        model_settings = getattr(ctx.deps, "model_settings", None)
+
+    # Build full message list: history + current request
+    messages = []
+    if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"):
+        messages.extend(ctx.state.message_history)
+
+    current_request = getattr(node, "request", None)
+    if current_request:
+        messages.append(current_request)
+
+    return messages, model, model_settings
+
+
+def _patch_graph_nodes() -> None:
+    """
+    Patches the graph node execution to create appropriate spans.
+
+    ModelRequestNode -> Creates ai_client span for model requests
+    CallToolsNode -> Handles tool calls (spans created in tool patching)
+    """
+
+    # Patch ModelRequestNode to create ai_client spans
+    original_model_request_run = ModelRequestNode.run
+
+    @wraps(original_model_request_run)
+    async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any":
+        did_stream = getattr(self, "_did_stream", None)
+        cached_result = getattr(self, "_result", None)
+        if did_stream or cached_result is not None:
+            return await original_model_request_run(self, ctx)
+
+        messages, model, model_settings = _extract_span_data(self, ctx)
+
+        with ai_client_span(messages, None, model, model_settings) as span:
+            span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+
+            result = await original_model_request_run(self, ctx)
+
+            # Extract response from result if available
+            model_response = None
+            if hasattr(result, "model_response"):
+                model_response = result.model_response
+
+            update_ai_client_span(span, model_response)
+            return result
+
+    ModelRequestNode.run = wrapped_model_request_run
+
+    # Patch ModelRequestNode.stream for streaming requests
+    original_model_request_stream = ModelRequestNode.stream
+
+    def create_wrapped_stream(
+        original_stream_method: "Callable[..., Any]",
+    ) -> "Callable[..., Any]":
+        """Create a wrapper for ModelRequestNode.stream that creates chat spans."""
+
+        @asynccontextmanager
+        @wraps(original_stream_method)
+        async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any":
+            did_stream = getattr(self, "_did_stream", None)
+            if did_stream:
+                async with original_stream_method(self, ctx) as stream:
+                    yield stream
+
+            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:
+                span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+
+                # Call the original stream method
+                async with original_stream_method(self, ctx) as stream:
+                    yield stream
+
+                # After streaming completes, update span with response data
+                # The ModelRequestNode stores the final response in _result
+                model_response = None
+                if hasattr(self, "_result") and self._result is not None:
+                    # _result is a NextNode containing the model_response
+                    if hasattr(self._result, "model_response"):
+                        model_response = self._result.model_response
+
+                update_ai_client_span(span, model_response)
+
+        return wrapped_model_request_stream
+
+    ModelRequestNode.stream = create_wrapped_stream(original_model_request_stream)