Skip to content

ref(pydantic-ai): Resolve version divergence in _compat and restructure modules #7196

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ericapisani wants to merge 1 commit into from
Draft

ref(pydantic-ai): Resolve version divergence in _compat and restructure modules #7196

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
170 changes: 17 additions & 153 deletions sentry_sdk/integrations/pydantic_ai/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,148 +1,30 @@
import functools

from sentry_sdk.integrations import DidNotEnable, Integration
from sentry_sdk.utils import capture_internal_exceptions, parse_version

try:
import pydantic_ai # noqa: F401
from pydantic_ai import Agent
except ImportError:
raise DidNotEnable("pydantic-ai not installed")


from importlib.metadata import PackageNotFoundError, version
from typing import TYPE_CHECKING

from .patches import (
_patch_agent_run,
_patch_graph_nodes,
_patch_tool_execution,
)
from .spans.ai_client import ai_client_span, update_ai_client_span

if TYPE_CHECKING:
from typing import Any

from pydantic_ai import ModelRequestContext, RunContext
from pydantic_ai.capabilities import Hooks
from pydantic_ai.messages import ModelResponse


def register_hooks(hooks: "Hooks") -> None:
"""
Creates hooks for chat model calls and register the hooks by adding the hooks to the `capabilities` argument passed to `Agent.__init__()`.

The chat span opened in on_request is stored in the run's `RunContext.metadata`
dict, which pydantic-ai shares by reference between the hooks of one run. This
keeps span pairing correct per run (even for overlapping runs in one task) and
covers every entry point that fires request hooks (including `Agent.iter()`,
which the Agent.run/run_stream wrappers never see). It requires seeding a
metadata dict in `patched_init` below when the user did not provide one.
"""

@hooks.on.before_model_request
async def on_request(
ctx: "RunContext[None]", request_context: "ModelRequestContext"
) -> "ModelRequestContext":
run_context_metadata = ctx.metadata
if not isinstance(run_context_metadata, dict):
return request_context

span = None
with capture_internal_exceptions():
span = ai_client_span(
messages=request_context.messages,
agent=None,
model=request_context.model,
model_settings=request_context.model_settings,
)

if span is None:
return request_context

run_context_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":
run_context_metadata = ctx.metadata
if not isinstance(run_context_metadata, dict):
return response

span = run_context_metadata.pop("_sentry_span", None)
if span is None:
return response

with capture_internal_exceptions():
update_ai_client_span(span, response)
span.__exit__(None, None, None)

return response

@hooks.on.model_request_error
async def on_error(
ctx: "RunContext[None]",
*,
request_context: "ModelRequestContext",
error: "Exception",
) -> "ModelResponse":
run_context_metadata = ctx.metadata

if not isinstance(run_context_metadata, dict):
raise error

span = run_context_metadata.pop("_sentry_span", None)
if span is None:
raise error

with capture_internal_exceptions():
span.__exit__(type(error), error, error.__traceback__)

raise error

original_init = Agent.__init__

@functools.wraps(original_init)
def patched_init(self: "Agent[Any, Any]", *args: "Any", **kwargs: "Any") -> None:
caps = list(kwargs.get("capabilities") or [])
caps.append(hooks)
kwargs["capabilities"] = caps

metadata = kwargs.get("metadata")
if metadata is None:
kwargs["metadata"] = {} # Used as shared reference between hooks

return original_init(self, *args, **kwargs)

Agent.__init__ = patched_init # type: ignore[method-assign]


class PydanticAIIntegration(Integration):
"""
Typical interaction with the library:
1. The user creates an Agent instance with configuration, including system instructions sent to every model call.
2. The user calls `Agent.run()` or `Agent.run_stream()` to start an agent run. The latter can be used to incrementally receive progress.
3. In a loop, the agent repeatedly calls the model, maintaining a conversation history that includes previous messages and tool results, which is passed to each call.

Internally, Pydantic AI maintains an execution graph in which ModelRequestNode are responsible for model calls, including retries.
Hooks using the decorators provided by `pydantic_ai.capabilities` create and manage spans for model calls when these hooks are available (newer library versions);
older versions are instrumented by patching the graph nodes directly (see patches/graph_nodes.py).

The wrappers around `Agent.run()` and `Agent.run_stream()` track each in-flight run on a contextvar stack (see _run_context.py); the tool patches and span
helpers read the current agent from there. The request hooks pair each chat span with its model request through the run's `RunContext.metadata` dict
(see register_hooks), which stays correct per run and also covers entry points the wrappers don't instrument, such as `Agent.iter()`.
How the integration is put together:
- _compat.py resolves the installed pydantic-ai version and every version-dependent decision, once, at import time.
- _extract.py is the only module that reads pydantic-ai object internals; it returns plain data structures.
- _spans.py creates spans and writes extracted data onto them.
- _run_context.py tracks each in-flight run on a contextvar stack; the tool wrapper and span helpers read the current agent from there.
- _wrap_agent.py instruments Agent.run / Agent.run_stream (invoke_agent spans, isolation scopes, run tracking).
- _wrap_model.py emits chat spans for model requests via one of two backends chosen in _compat: request hooks (>= 1.73), paired per run through RunContext.metadata, or graph-node patching (older versions).
- _wrap_tools.py instruments the single ToolManager method all tool calls flow through (execute_tool spans).
"""

identifier = "pydantic_ai"
origin = f"auto.ai.{identifier}"
using_request_hooks = False

def __init__(
self, include_prompts: bool = True, handled_tool_call_exceptions: bool = True
Expand All @@ -169,32 +51,14 @@ def setup_once() -> None:
- Model requests (AI client calls)
- Tool executions
"""
# Deferred imports keep `import sentry_sdk.integrations.pydantic_ai`
# cheap when the integration is never enabled; they are the only
# intra-package imports in this module, keeping the import graph
# acyclic.
from ._wrap_agent import _patch_agent_run
from ._wrap_model import install_model_backend
from ._wrap_tools import _patch_tool_execution

_patch_agent_run()
_patch_tool_execution()

PydanticAIIntegration.using_request_hooks = False
try:
PYDANTIC_AI_VERSION = version("pydantic-ai-slim")
except PackageNotFoundError:
return

PYDANTIC_AI_VERSION = parse_version(PYDANTIC_AI_VERSION)
if PYDANTIC_AI_VERSION is None:
return

# ModelRequestContext.model added in https://github.com/pydantic/pydantic-ai/commit/f1260dfe09907f17688eee1646daf898fc428d4c
if PYDANTIC_AI_VERSION < (
1,
73,
):
_patch_graph_nodes()
return

try:
from pydantic_ai.capabilities import Hooks
except ImportError:
return

PydanticAIIntegration.using_request_hooks = True
hooks = Hooks()
register_hooks(hooks)
install_model_backend()
70 changes: 70 additions & 0 deletions sentry_sdk/integrations/pydantic_ai/_compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
"""Version detection and version-dependent imports for pydantic-ai.

Everything here is resolved once at import time. The rest of the integration
consumes the resulting constants instead of probing versions or attributes at
call time, so "what runs on version X" is answered entirely by this module.
"""

from typing import TYPE_CHECKING

from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.utils import package_version

try:
from pydantic_ai import messages as _messages
from pydantic_ai.agent import Agent # noqa: F401
from pydantic_ai.exceptions import ToolRetryError # noqa: F401

try:
from pydantic_ai.tool_manager import ToolManager
except ImportError:
# older versions
from pydantic_ai._tool_manager import ToolManager # type: ignore
except ImportError:
raise DidNotEnable("pydantic-ai not installed")

if TYPE_CHECKING:
from typing import Optional

# Message part classes are resolved individually so that a single upstream
# rename degrades only the extraction paths that need that class, instead of
# silently disabling all of them at once.
BaseToolCallPart = getattr(_messages, "BaseToolCallPart", None)
BaseToolReturnPart = getattr(_messages, "BaseToolReturnPart", None)
BinaryContent = getattr(_messages, "BinaryContent", None)
ImageUrl = getattr(_messages, "ImageUrl", None)
SystemPromptPart = getattr(_messages, "SystemPromptPart", None)
TextPart = getattr(_messages, "TextPart", None)
ThinkingPart = getattr(_messages, "ThinkingPart", None)

PYDANTIC_AI_VERSION = package_version("pydantic-ai-slim")

# The ToolManager method through which all tool calls flow; renamed from
# _call_tool to execute_tool_call in newer versions. None means the method
# could not be found and tool instrumentation is skipped.
TOOL_CALL_METHOD: "Optional[str]" = None
if hasattr(ToolManager, "execute_tool_call"):
TOOL_CALL_METHOD = "execute_tool_call"
elif hasattr(ToolManager, "_call_tool"):
TOOL_CALL_METHOD = "_call_tool"

# Request hooks (pydantic_ai.capabilities) are usable from 1.73 on, when
# ModelRequestContext.model was added:
# https://github.com/pydantic/pydantic-ai/commit/f1260dfe09907f17688eee1646daf898fc428d4c
USES_REQUEST_HOOKS = False
if PYDANTIC_AI_VERSION is not None and PYDANTIC_AI_VERSION >= (1, 73):
try:
from pydantic_ai.capabilities import Hooks # noqa: F401

USES_REQUEST_HOOKS = True
except ImportError:
USES_REQUEST_HOOKS = False

# Which mechanism emits chat spans for model requests: request hooks on new
# versions, graph-node patching on old ones. None (unknown version, or hooks
# unavailable on a new version) means only agent and tool spans are emitted.
MODEL_BACKEND: "Optional[str]" = None
if USES_REQUEST_HOOKS:
MODEL_BACKEND = "hooks"
elif PYDANTIC_AI_VERSION is not None and PYDANTIC_AI_VERSION < (1, 73):
MODEL_BACKEND = "graph_nodes"
32 changes: 11 additions & 21 deletions sentry_sdk/integrations/pydantic_ai/_extract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@
private attributes and version-dependent shapes) so that upstream library
changes are absorbed here rather than throughout the integration. The one
exception is control-flow state read at the patch points themselves (e.g.
ModelRequestNode._did_stream in patches/graph_nodes.py and Tool.tool_def in
patches/tools.py); everything else consumes the plain data structures
ModelRequestNode._did_stream in _wrap_model.py and Tool.tool_def in
_wrap_tools.py); everything else consumes the plain data structures
returned here.
"""

Expand All @@ -18,25 +18,15 @@
from sentry_sdk.consts import SPANDATA
from sentry_sdk.utils import safe_serialize

try:
from pydantic_ai.messages import (
BaseToolCallPart,
BaseToolReturnPart,
BinaryContent,
ImageUrl,
SystemPromptPart,
TextPart,
ThinkingPart,
)
except ImportError:
# Fallback if these classes are not available
BaseToolCallPart = None # type: ignore[misc,assignment]
BaseToolReturnPart = None # type: ignore[misc,assignment]
BinaryContent = None # type: ignore[misc,assignment]
ImageUrl = None # type: ignore[misc,assignment]
SystemPromptPart = None # type: ignore[misc,assignment]
TextPart = None # type: ignore[misc,assignment]
ThinkingPart = None # type: ignore[misc,assignment]
from ._compat import (
BaseToolCallPart,
BaseToolReturnPart,
BinaryContent,
ImageUrl,
SystemPromptPart,
TextPart,
ThinkingPart,
)

if TYPE_CHECKING:
from typing import Any, Dict, List, Optional
Expand Down
22 changes: 11 additions & 11 deletions sentry_sdk/integrations/pydantic_ai/_run_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,18 +46,18 @@ def get_is_streaming() -> bool:
@contextmanager
def agent_run_scope(agent: "Any", is_streaming: bool = False) -> "Iterator[AgentRun]":
"""Track an agent run on the contextvar stack for the duration of the
with block."""
with block.

On exit, exactly this run is removed from the stack (by identity, not a
token reset), so streaming runs that exit out of LIFO order or in a
different asyncio task never erase other still-active runs.
"""
run = AgentRun(agent=agent, is_streaming=is_streaming)
token = _agent_run_stack.set(_agent_run_stack.get() + (run,))
_agent_run_stack.set(_agent_run_stack.get() + (run,))
try:
yield run
finally:
try:
_agent_run_stack.reset(token)
except (LookupError, ValueError):
# A streaming run's context manager can be exited in a different
# asyncio task (and therefore a different Context) than it was
# entered in, in which case the token cannot be reset. The stack
# entry only lives in the entering task's context copy, so there
# is nothing to clean up.
pass
stack = _agent_run_stack.get()
new_stack = tuple(r for r in stack if r is not run)
if len(new_stack) != len(stack):
_agent_run_stack.set(new_stack)
Loading
Loading