from __future__ import annotations as _annotations import ast import asyncio import copy import functools import inspect import re import sys import textwrap import time import uuid from collections.abc import ( AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Callable, Generator, Iterable, Iterator, ) from concurrent.futures import Executor from contextlib import asynccontextmanager, contextmanager, suppress from contextvars import ContextVar, copy_context from dataclasses import MISSING, dataclass, fields, is_dataclass from datetime import datetime, timezone from enum import Enum from types import GenericAlias from typing import ( TYPE_CHECKING, Any, Generic, TypeAlias, TypeGuard, get_args, get_origin, overload, ) import anyio from anyio.to_thread import run_sync from pydantic import BaseModel, TypeAdapter from pydantic._internal import _decorators, _typing_extra from pydantic.json_schema import JsonSchemaValue from typing_extensions import ParamSpec, TypeIs, TypeVar, is_typeddict from typing_inspection import typing_objects from typing_inspection.introspection import is_union_origin from pydantic_graph._utils import ( AbstractSpan, run_until_complete as _graph_run_until_complete, ) from pydantic_graph.exceptions import UnsupportedEventLoopError from pydantic_graph.util import get_callable_name from .exceptions import UserError if sys.version_info < (3, 11): from exceptiongroup import BaseExceptionGroup as BaseExceptionGroup # pragma: lax no cover else: BaseExceptionGroup = BaseExceptionGroup # pragma: lax no cover AbstractSpan = AbstractSpan if TYPE_CHECKING: from pydantic_ai.agent import AgentRun, AgentRunResult from pydantic_graph import GraphRun from . import messages as _messages from .tools import ObjectJsonSchema _P = ParamSpec('_P') _R = TypeVar('_R') _disable_threads: ContextVar[bool] = ContextVar('_disable_threads', default=sys.platform == 'emscripten') _thread_executor: ContextVar[Executor | None] = ContextVar('_thread_executor', default=None) _in_sync_callback: ContextVar[bool] = ContextVar('_in_sync_callback', default=False) # Any cancellation delivered while awaiting a worker thread abandons that thread, not just the one a # deadline schedules: `anyio` cannot tell them apart. That is acceptable only because this dial is set # tightly around calls that are already armed with a deadline, whose owner asked for a timeout. _abandon_on_cancel: ContextVar[bool] = ContextVar('_abandon_on_cancel', default=False) def check_no_nested_sync_run() -> None: """Reject sync agent entry points inside sync callbacks dispatched by Pydantic AI. Sync tools, output functions, and similar callbacks are dispatched through [`run_in_executor`][pydantic_ai._utils.run_in_executor], which flags the callback's context — whether the callback runs on a worker thread or inline under [`disable_threads`][pydantic_ai._utils.disable_threads]. On a worker thread, a nested sync run starts a second event loop that can deadlock against an async resource bound to the parent run's loop; inline, it would drive the already-running loop and fail anyway. Either way we fail fast with guidance instead. """ if _in_sync_callback.get(): raise UserError( '`Agent.run_sync()` and `Agent.run_stream_sync()` cannot be used inside a synchronous tool, ' 'output function, or other function called during an agent run, as they can deadlock the run. ' 'Make the function `async def` and use `await agent.run(...)` or `async with agent.run_stream(...)` instead.' ) def run_until_complete(coro: Awaitable[_R]) -> _R: """Run `coro` to completion on the event loop, for use by the sync wrappers. Wraps `pydantic_graph`'s `run_until_complete()` to report an event loop that can't be driven by the caller -- like Temporal's workflow event loop -- as a `UserError`, which is how the rest of the library reports usage mistakes, and which durable execution integrations know to treat as a deterministic failure rather than an infrastructure one to retry. """ try: return _graph_run_until_complete(coro) except UnsupportedEventLoopError as e: raise UserError(e.message) from e @contextmanager def disable_threads() -> Generator[None]: """Context manager to disable thread-based execution for sync functions. Inside this context, sync functions will execute inline rather than being sent to a thread pool via [`anyio.to_thread.run_sync`][anyio.to_thread.run_sync]. This is useful in environments where threading is restricted, such as Temporal workflows which use a sandboxed event loop. On emscripten, sync callbacks already run inline by default because Python threads are unavailable there. Yields: None """ token = _disable_threads.set(True) try: yield finally: _disable_threads.reset(token) @contextmanager def using_thread_executor(executor: Executor) -> Generator[None]: """Context manager to use a custom executor for running sync functions in threads. Inside this context, sync functions will be executed using the provided executor via [`asyncio.get_running_loop().run_in_executor()`][asyncio.loop.run_in_executor] instead of the default [`anyio.to_thread.run_sync`][anyio.to_thread.run_sync]. This is useful in long-running servers (e.g. FastAPI) where thread accumulation from ephemeral anyio worker threads can be a problem, and you want to use a bounded `ThreadPoolExecutor` instead. Args: executor: The executor to use for running sync functions. Yields: None """ token = _thread_executor.set(executor) try: yield finally: _thread_executor.reset(token) @contextmanager def abandon_threads_on_cancel() -> Generator[None]: """Context manager to abandon worker threads running sync functions when they're cancelled. Inside this context, a cancellation delivered while awaiting a worker thread abandons that thread -- it runs to completion in the background and its result is discarded -- instead of waiting for it to finish. Outside it, [`anyio.to_thread.run_sync`][anyio.to_thread.run_sync] shields the await, so the cancellation is only delivered once the thread returns. This is used around calls that carry a deadline, so that [`anyio.fail_after`][anyio.fail_after] can actually raise `TimeoutError` when a sync function overruns it, rather than only after it returns. Yields: None """ token = _abandon_on_cancel.set(True) try: yield finally: _abandon_on_cancel.reset(token) async def run_in_executor(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R: def call_with_sync_agent_guard() -> _R: token = _in_sync_callback.set(True) try: return func(*args, **kwargs) finally: _in_sync_callback.reset(token) if _disable_threads.get(): return call_with_sync_agent_guard() executor = _thread_executor.get() if executor is not None: loop = asyncio.get_running_loop() ctx = copy_context() return await loop.run_in_executor(executor, ctx.run, call_with_sync_agent_guard) return await run_sync(call_with_sync_agent_guard, abandon_on_cancel=_abandon_on_cancel.get()) def is_async_generator_already_running(exc: RuntimeError) -> bool: return 'asynchronous generator is already running' in str(exc) def is_model_like(type_: Any) -> bool: """Check if something is a pydantic model, dataclass or typedict. These should all generate a JSON Schema with `{"type": "object"}` and therefore be usable directly as function parameters. """ return ( isinstance(type_, type) and not isinstance(type_, GenericAlias) and ( issubclass(type_, BaseModel) or is_dataclass(type_) # pyright: ignore[reportUnknownArgumentType] or is_typeddict(type_) # pyright: ignore[reportUnknownArgumentType] or getattr(type_, '__is_model_like__', False) # pyright: ignore[reportUnknownArgumentType] ) ) def check_object_json_schema(schema: JsonSchemaValue) -> ObjectJsonSchema: from .exceptions import UserError if schema.get('type') == 'object': return schema elif ref := schema.get('$ref'): prefix = '#/$defs/' # Return the referenced schema unless it contains additional nested references. if ( ref.startswith(prefix) and (resolved := schema.get('$defs', {}).get(ref[len(prefix) :])) and resolved.get('type') == 'object' and not _contains_ref(resolved) ): return resolved return schema else: raise UserError('Schema must be an object') def _contains_ref(obj: JsonSchemaValue | list[JsonSchemaValue]) -> bool: """Recursively check if an object contains any $ref keys.""" items: Iterable[JsonSchemaValue] if isinstance(obj, dict): if '$ref' in obj: return True items = obj.values() else: items = obj return any(isinstance(item, dict | list) and _contains_ref(item) for item in items) # pyright: ignore[reportUnknownArgumentType] T = TypeVar('T') def check_tools_prepare_func_result(result: Iterable[T] | None, prepare_func: Any) -> list[T]: """Validate and normalize a tool-prepare callback result.""" if result is None: raise UserError( f'Prepare function {get_callable_name(prepare_func)!r} returned `None`; ' 'return `[]` to expose no tools, or return `tool_defs` to pass them through unchanged.' ) return list(result) @dataclass class Some(Generic[T]): """Analogous to Rust's `Option::Some` type.""" value: T Option: TypeAlias = Some[T] | None """Analogous to Rust's `Option` type, usage: `Option[Thing]` is equivalent to `Some[Thing] | None`.""" async def gather(*coros: Awaitable[T]) -> list[T]: """Run awaitables concurrently via an `anyio` task group and return results in input order. Unlike `asyncio.gather`, a failure in one coroutine cancels the rest instead of leaving them as orphan background tasks. If exactly one task fails, its exception is re-raised directly to match `asyncio.gather`'s shape; multi-failure cases propagate as an `ExceptionGroup`. """ sentinel = Unset() results: list[T | Unset] = [sentinel] * len(coros) async def _run(index: int, coro: Awaitable[T]) -> None: results[index] = await coro try: async with anyio.create_task_group() as tg: for i, coro in enumerate(coros): tg.start_soon(_run, i, coro) except BaseExceptionGroup as eg: if len(eg.exceptions) == 1: exc = eg.exceptions[0] exc.__suppress_context__ = True raise exc raise final_results: list[T] = [] for result in results: assert not isinstance(result, Unset) final_results.append(result) return final_results async def cancel_and_drain(*tasks: asyncio.Task[Any], msg: object = None) -> None: """Cancel any tasks still running and wait for them to finish unwinding. Cleanup-only: results and exceptions from `tasks` are intentionally discarded so a cancelled child cannot replace an exception already propagating in the caller. Use after `asyncio.create_task` when an outer cancel/exception means the spawned tasks must be torn down before the caller exits. """ for task in tasks: if not task.done(): task.cancel(msg=msg) # Pydantic Graph runs nodes under AnyIO cancel scopes. Once the outer scope # is cancelled, AnyIO uses level cancellation and can keep re-cancelling at # each await. Shield the drain so child tasks get one explicit cancel above, # then can finish normal async `finally` cleanup before we re-raise. with anyio.CancelScope(shield=True): await asyncio.gather(*tasks, return_exceptions=True) def raise_if_cancelling() -> None: """Re-assert an external cancellation that a completed step absorbed (level-triggered backstop). A step the run awaits — a Temporal activity under `WAIT_CANCELLATION_COMPLETED`, an `event_stream_handler`, a capability hook — can catch the `CancelledError` injected by `task.cancel()` and return normally. asyncio even *delegates* a task's cancellation to the future it is currently awaiting, so an awaited child task that absorbs its cancel silently completes the awaiting task too. Either way the cancellation is an edge the framework never sees, and without a re-check the run would complete as if it was never cancelled. Well-behaved consumers of their own cancellation, like `asyncio.timeout()` and AnyIO cancel scopes, balance `Task.cancelling()` back down with `Task.uncancel()`, so a positive count at a step boundary is treated as a still-pending cancellation of the run and re-raised. This is deliberately a *policy*, not a proof of external intent: code awaited by the run that cancels its own task as an internal wake-up and suppresses the `CancelledError` without calling `Task.uncancel()` (a pre-3.11 idiom) will be read as a cancelled run. Call this only after the just-completed step's results have been recorded to message history, so cancellation never discards completed work. The re-raise is a fresh `CancelledError`: the originally-injected exception object (and any message it carried) was consumed by whatever absorbed it and cannot be recovered — the cancellation *state* is re-asserted, not the original exception. On Python 3.10 `Task.cancelling()` does not exist and this is a no-op: an absorbed external cancellation cannot be reliably detected there, so the cancellation guarantee is documented as best-effort on 3.10. """ if sys.version_info < (3, 11): # pragma: lax no cover return try: task = asyncio.current_task() except RuntimeError: # pragma: no cover # No running asyncio loop (e.g. a Trio-backed run). return if task is not None and task.cancelling() > 0: raise asyncio.CancelledError('pydantic-ai: re-asserting a cancellation absorbed by a completed step') class Unset: """A singleton to represent an unset value.""" pass UNSET = Unset() def is_set(t_or_unset: T | Unset) -> TypeGuard[T]: return t_or_unset is not UNSET def replace_no_init(obj: T, **changes: Any) -> T: """Return a shallow copy of a dataclass instance with `changes` applied to its fields. Use instead of `dataclasses.replace` on instances of subclassable dataclasses: `replace` reconstructs through `type(obj).__init__`, which crashes for subclasses whose custom `__init__` doesn't accept the dataclass field names (https://github.com/pydantic/pydantic-ai/issues/6674). Copying preserves the subclass and all of its state, and never re-runs `__init__`/`__post_init__` — the caller must refresh any state it derives from the changed fields. Not a drop-in for `replace`: fields declared `init=False` are carried over rather than reset, so call sites that rely on `replace` resetting derived state (e.g. per-run state isolation) must keep using `replace`. """ assert is_dataclass(obj) field_names = {f.name for f in fields(obj)} if unknown := changes.keys() - field_names: raise TypeError(f'Invalid field name(s) for {type(obj).__name__}: {", ".join(sorted(unknown))}') new_obj = copy.copy(obj) if new_obj is obj: # A `__copy__` that returns `self` (immutable-style classes) would make the loop below # mutate the original in place, silently leaking the changes to everyone holding it. raise TypeError(f'Cannot replace fields on {type(obj).__name__}: its `__copy__` does not return a new instance') for name, value in changes.items(): # `object.__setattr__` so frozen dataclasses work too: `new_obj` is a fresh copy no # caller has seen yet, the same way a frozen dataclass's own `__init__` assigns fields. object.__setattr__(new_obj, name, value) return new_obj async def _cleanup_temporal_group( task: asyncio.Task[Any] | None, aiterator: AsyncIterator[Any], ) -> None: """Clean up pending task and async iterator after group_by_temporal exits.""" if task: task.cancel('Cancelling group_by_temporal pending task') with suppress(asyncio.CancelledError, StopAsyncIteration): await task aclose = getattr(aiterator, 'aclose', None) if aclose is not None: # pragma: no branch await aclose() async def aclose_if_supported(stream: AsyncIterable[Any]) -> None: """Close an async iterable if it exposes an `aclose` method.""" aclose: Callable[[], Awaitable[None]] | None = getattr(stream, 'aclose', None) if aclose is not None: await aclose() async def aclose_all(streams: Iterable[AsyncIterable[Any]]) -> None: """Close every async iterable, then propagate any close failures.""" errors: list[BaseException] = [] for stream in streams: try: await aclose_if_supported(stream) except BaseException as error: errors.append(error) if len(errors) == 1: raise errors[0] if errors: raise BaseExceptionGroup('Errors closing async iterables', errors) @asynccontextmanager async def group_by_temporal( aiterable: AsyncIterable[T], soft_max_interval: float | None ) -> AsyncGenerator[AsyncIterable[list[T]]]: """Group items from an async iterable into lists based on time interval between them. Effectively, this debounces the iterator. This returns a context manager usable as an iterator so any pending tasks can be cancelled if an error occurs during iteration. Usage: ```python async with group_by_temporal(yield_groups(), 0.1) as groups_iter: async for groups in groups_iter: print(groups) ``` Args: aiterable: The async iterable to group. soft_max_interval: Maximum interval over which to group items, this should avoid a trickle of items causing a group to never be yielded. It's a soft max in the sense that once we're over this time, we yield items as soon as `anext(aiter)` returns. If `None`, no grouping/debouncing is performed Returns: A context manager usable as an async iterable of lists of items produced by the input async iterable. """ # we might wait for the next item more than once, so we store the task to await next time task: asyncio.Task[T] | None = None aiterator = aiter(aiterable) if soft_max_interval is None: async def async_iter_groups() -> AsyncIterator[list[T]]: async for item in aiterator: yield [item] else: async def async_iter_groups() -> AsyncIterator[list[T]]: nonlocal task assert soft_max_interval is not None and soft_max_interval >= 0, ( 'soft_max_interval must be a positive number' ) buffer: list[T] = [] group_start_time: float | None = None while True: if group_start_time is None: # group hasn't started, we just wait for the maximum interval wait_time = soft_max_interval else: # wait for the time remaining in the group wait_time = soft_max_interval - (time.monotonic() - group_start_time) # if there's no current task, we get the next one if task is None: # anext(aiter) returns an Awaitable[T], not a Coroutine which asyncio.create_task expects # so far, this doesn't seem to be a problem task = asyncio.create_task(anext(aiterator)) # pyright: ignore[reportArgumentType,reportUnknownVariableType] # we use asyncio.wait to avoid cancelling the coroutine if it's not done done, _ = await asyncio.wait((task,), timeout=wait_time) if done: # the one task we waited for completed try: item = done.pop().result() except StopAsyncIteration: # if the task raised StopAsyncIteration, we're done iterating if buffer: yield buffer task = None break else: # we got an item, add it to the buffer and set task to None to get the next item buffer.append(item) task = None # if this is the first item in the group, set the group start time if group_start_time is None: group_start_time = time.monotonic() elif buffer: # otherwise if the task timeout expired and we have items in the buffer, yield the buffer yield buffer # clear the buffer and reset the group start time ready for the next group buffer = [] group_start_time = None try: yield async_iter_groups() finally: await _cleanup_temporal_group(task, aiterator) def sync_anext(iterator: Iterator[T]) -> T: """Get the next item from a sync iterator, raising `StopAsyncIteration` if it's exhausted. Useful when iterating over a sync iterator in an async context. """ try: return next(iterator) except StopIteration as e: raise StopAsyncIteration() from e def now_utc() -> datetime: return datetime.now(tz=timezone.utc) def fill_run_metadata(message: _messages.ModelMessage, *, run_id: str | None, conversation_id: str | None) -> None: """Fill in framework-tracked metadata (`timestamp`, `run_id`, `conversation_id`) that's still unset. Producer-supplied values are preserved; only unset fields are filled in. Centralizing the field list here means a new framework-tracked field only needs to be handled in one place, rather than every site that materializes a message into the history. """ message.timestamp = message.timestamp or now_utc() message.run_id = message.run_id or run_id message.conversation_id = message.conversation_id or conversation_id def validate_uploaded_file_provider(item: _messages.UploadedFile, *, system: str, model_type_name: str) -> None: """Raise `UserError` if an `UploadedFile` references a different provider than the model it was passed to.""" if item.provider_name != system: raise UserError( f'UploadedFile with `provider_name={item.provider_name!r}` cannot be used with {model_type_name}. ' f'Expected `provider_name` to be `{system!r}`.' ) def guard_tool_call_id( t: _messages.ToolCallPart | _messages.ToolReturnPart | _messages.RetryPromptPart | _messages.NativeToolCallPart | _messages.NativeToolReturnPart, ) -> str: """Type guard that either returns the tool call id or generates a new one if it's None.""" return t.tool_call_id or generate_tool_call_id() TOOL_NAME_SANITIZER = re.compile(r'[^a-zA-Z0-9_-]') """Regex matching characters not allowed in tool names by most providers.""" def sanitize_tool_name(name: str) -> str: """Replace characters outside `[a-zA-Z0-9_-]` with `_`.""" return TOOL_NAME_SANITIZER.sub('_', name) TOOL_CALL_ID_PREFIX = 'pyd_ai_' def generate_tool_call_id() -> str: """Generate a tool call id. Ensure that the tool call id is unique. """ return f'{TOOL_CALL_ID_PREFIX}{uuid.uuid4().hex}' SourceT = TypeVar('SourceT', bound=AsyncIterable[Any], default=AsyncIterable[T]) class PeekableAsyncStream(Generic[T, SourceT]): """Wraps an async iterable of type T and allows peeking at the *next* item without consuming it. We only buffer one item at a time (the next item). Once that item is yielded, it is discarded. This is a single-pass stream. """ def __init__(self, source: SourceT): self.source = source self._source_iter: AsyncIterator[T] | None = None self._buffer: T | Unset = UNSET self._exhausted = False # Serialize access to the underlying source so cancelling an in-flight pull releases the lock before # `aclose()` closes it. A debounced consumer (`group_by_temporal`) prefetches the next item in a background # task, so the source generator can be mid-`anext` when the stream is abandoned; closing it concurrently # would raise `RuntimeError: aclose(): asynchronous generator is already running`. self._source_lock = anyio.Lock() self._pull_scopes: set[anyio.CancelScope] = set() async def peek(self) -> T | Unset: """Returns the next item that would be yielded without consuming it. Returns None if the stream is exhausted. """ if self._exhausted: return UNSET # If we already have a buffered item, just return it. if not isinstance(self._buffer, Unset): return self._buffer # Otherwise, we need to fetch the next item from the underlying iterator. if self._source_iter is None: self._source_iter = aiter(self.source) with anyio.CancelScope() as scope: self._pull_scopes.add(scope) try: async with self._source_lock: try: self._buffer = await anext(self._source_iter) except StopAsyncIteration: self._exhausted = True return UNSET return self._buffer finally: self._pull_scopes.discard(scope) # Only reached when `aclose()` cancelled the scope: the stream is closed, so iteration is over. self._exhausted = True return UNSET async def is_exhausted(self) -> bool: """Returns True if the stream is exhausted, False otherwise.""" return isinstance(await self.peek(), Unset) def __aiter__(self) -> AsyncIterator[T]: # For a single-pass iteration, we can return self as the iterator. return self async def __anext__(self) -> T: """Yields the buffered item if present, otherwise fetches the next item from the underlying source. Raises StopAsyncIteration if the stream is exhausted. """ if self._exhausted: raise StopAsyncIteration # If we have a buffered item, yield it. if not isinstance(self._buffer, Unset): item = self._buffer self._buffer = UNSET return item # Otherwise, fetch the next item from the source. if self._source_iter is None: self._source_iter = aiter(self.source) with anyio.CancelScope() as scope: self._pull_scopes.add(scope) try: async with self._source_lock: try: return await anext(self._source_iter) except StopAsyncIteration: self._exhausted = True raise finally: self._pull_scopes.discard(scope) # Only reached when `aclose()` cancelled the scope: the stream is closed, so iteration is over. self._exhausted = True raise StopAsyncIteration async def aclose(self) -> None: self._exhausted = True for scope in self._pull_scopes: scope.cancel() value = self._source_iter if self._source_iter is not None else self.source # Wait for the cancelled pull to release the source before closing it, so we don't close a # generator that's still running. async with self._source_lock: await aclose_if_supported(value) def get_traceparent(x: AgentRun | AgentRunResult | GraphRun[Any, Any, Any]) -> str: return x._traceparent(required=False) or '' # type: ignore[reportPrivateUsage] def dataclasses_no_defaults_repr(self: Any) -> str: """Exclude fields with values equal to the field default. A field is shown when its value differs from an explicit `default`. Fields that are required or that only have a `default_factory` have no plain default to compare against here, so they are always shown (the `default_factory` is deliberately not called: some factories are impure, e.g. `uuid7()` or `now_utc()`, and `repr()` must stay observational). The comparison is guarded because a value whose `__ne__`/`__bool__` does not return a plain `bool` (e.g. a numpy array or pandas `Series`/`DataFrame`) would otherwise make `repr()` raise `ValueError`, which breaks logging and traceback formatting of the message history. """ def include_field(f: Any) -> bool: if not f.repr: return False if f.default is MISSING: return True try: return bool(getattr(self, f.name) != f.default) except Exception: # `repr()` must never raise, regardless of how a field value implements `__ne__`/`__bool__` # (e.g. numpy/pandas return non-bool comparisons), so the broad catch here is intentional. return True kv_pairs = (f'{f.name}={getattr(self, f.name)!r}' for f in fields(self) if include_field(f)) return f'{self.__class__.__qualname__}({", ".join(kv_pairs)})' def own_annotations(cls: type) -> dict[str, Any]: """The annotations written on `cls` itself, without forcing lazy (PEP 649) evaluation. On Python 3.14+ annotations are evaluated lazily, so a field referencing a class defined later in the module must not be evaluated while the class is still being built — `@dataclass` defers it, and anything inspecting a class from `__init_subclass__` has to as well. `Format.FORWARDREF` never raises `NameError`, so an unresolvable annotation comes back as a `ForwardRef` instead. """ if sys.version_info >= (3, 14): import annotationlib return dict(annotationlib.get_annotations(cls, format=annotationlib.Format.FORWARDREF)) return dict(inspect.get_annotations(cls)) def declares_dataclass_fields(cls: type) -> bool: """Whether `@dataclass` would give `cls` fields of its own, i.e. it annotates a non-`ClassVar` name. Annotations are inspected without being evaluated (see [`own_annotations`][]), so a `ClassVar` that can't be resolved yet is still recognized from the way it was written. """ for annotation in own_annotations(cls).values(): if typing_objects.is_classvar(get_origin(annotation)) or typing_objects.is_classvar(annotation): continue # An unevaluated annotation arrives as a string or `ForwardRef`; match how it was written. text = annotation if isinstance(annotation, str) else getattr(annotation, '__forward_arg__', None) if text is not None and re.match(r'^(typing\.)?ClassVar\b', text.strip()): continue return True return False def copy_dataclass_fields(src: Any, dst_cls: type, **overrides: Any) -> Any: """Shared utility for typed-part narrowers — preserves base fields when promoting to a typed subclass. Construct a new dataclass instance from `src`'s fields, overriding selected ones. Lets typed-part narrowers stay maintainable when fields are added to the base class — base-class field changes flow through automatically instead of needing every narrower to be updated by hand. """ field_values: dict[str, Any] = {f.name: getattr(src, f.name) for f in fields(src)} field_values.update(overrides) return dst_cls(**field_values) _datetime_ta = TypeAdapter(datetime) def number_to_datetime(x: int | float) -> datetime: return _datetime_ta.validate_python(x) AwaitableCallable = Callable[..., Awaitable[T]] @overload def is_async_callable(obj: AwaitableCallable[T]) -> TypeIs[AwaitableCallable[T]]: ... @overload def is_async_callable(obj: Any) -> TypeIs[AwaitableCallable[Any]]: ... def is_async_callable(obj: Any) -> Any: """Correctly check if a callable is async. This function was copied from Starlette: https://github.com/encode/starlette/blob/78da9b9e218ab289117df7d62aee200ed4c59617/starlette/_utils.py#L36-L40 """ while isinstance(obj, functools.partial): obj = obj.func return inspect.iscoroutinefunction(obj) or (callable(obj) and inspect.iscoroutinefunction(obj.__call__)) async def await_maybe(value: T | Awaitable[T]) -> T: """Await `value` if it is awaitable, otherwise return it unchanged. Use this to resolve the result of calling a callback typed as `X | Awaitable[X]`, regardless of how the awaitable is produced: an `async def`, or a plain `def` / callable object that *returns* a coroutine. [`is_async_callable`][pydantic_ai._utils.is_async_callable] can't detect the latter because it inspects the callable rather than its result, so dispatching on it alone drops such a callback's coroutine un-awaited. """ if inspect.isawaitable(value): return await value return value def takes_run_context(callable_obj: Callable[..., Any]) -> bool: """Check if a callable takes a `RunContext` as its first argument. Args: callable_obj: The callable to check. Returns: `True` if the callable takes a `RunContext` as first argument, `False` otherwise. Raises: UserError: If the callable has annotations that cannot be resolved at runtime. """ from ._run_context import RunContext first_param_type = get_first_param_type(callable_obj) if first_param_type is None: return False return first_param_type is RunContext or get_origin(first_param_type) is RunContext def get_first_param_type(callable_obj: Callable[..., Any]) -> Any | None: """Get the type annotation of the first parameter of a callable. Handles regular functions, methods, and callable classes with __call__. Uses Pydantic internals to properly resolve type hints including forward references. Args: callable_obj: The callable to inspect. Returns: The type annotation of the first parameter, or None if the callable has no introspectable annotations. Raises: UserError: If the callable has annotations that cannot be resolved at runtime. """ try: sig = inspect.signature(callable_obj) except ValueError: return None try: first_param_name = next(iter(sig.parameters.keys())) except StopIteration: return None # See https://github.com/pydantic/pydantic/pull/11451 for a similar implementation in Pydantic callable_for_hints = callable_obj name = get_callable_name(callable_obj) if not isinstance(callable_obj, _decorators._function_like): # pyright: ignore[reportPrivateUsage] call_func = getattr(type(callable_obj), '__call__', None) if call_func is not None: callable_for_hints = call_func # Name the class rather than the (address-bearing) repr of the instance. name = type(callable_obj).__name__ else: return None # pragma: no cover try: type_hints = _typing_extra.get_function_type_hints(_decorators.unwrap_wrapped_function(callable_for_hints)) except NameError as e: # A `NameError` means the callable does have annotations, we just can't resolve them. Treating that # like "no annotations" would silently pick the wrong calling convention, so we surface it instead. raise UserError( f'Unable to resolve the type annotations of {name!r}: {e}. ' 'This typically happens when a type is imported inside an `if TYPE_CHECKING:` block ' 'in a module that uses `from __future__ import annotations`. ' 'Pydantic AI resolves these annotations at runtime to determine how to call the function, ' 'so every type used in its signature needs to be imported at runtime as well.' ) from e except (TypeError, AttributeError): return None return type_hints.get(first_param_name) def get_function_type_hints(func: Any) -> dict[str, Any]: """Resolve type hints for a function, including forward references. Wraps `pydantic._internal._typing_extra.get_function_type_hints` so callers don't need to import Pydantic internals directly. """ return _typing_extra.get_function_type_hints(func) def _update_mapped_json_schema_refs(s: dict[str, Any], name_mapping: dict[str, str]) -> None: """Update $refs in a schema to use the new names from name_mapping.""" if '$ref' in s: ref = s['$ref'] if ref.startswith('#/$defs/'): # pragma: no branch original_name = ref[8:] # Remove '#/$defs/' new_name = name_mapping.get(original_name, original_name) s['$ref'] = f'#/$defs/{new_name}' # Recursively update refs in properties if 'properties' in s: props: dict[str, dict[str, Any]] = s['properties'] for prop in props.values(): _update_mapped_json_schema_refs(prop, name_mapping) # Handle arrays if 'items' in s and isinstance(s['items'], dict): items: dict[str, Any] = s['items'] # pyright: ignore[reportUnknownVariableType] _update_mapped_json_schema_refs(items, name_mapping) if 'prefixItems' in s: prefix_items: list[dict[str, Any]] = s['prefixItems'] for item in prefix_items: _update_mapped_json_schema_refs(item, name_mapping) # Handle additionalProperties if 'additionalProperties' in s and isinstance(s['additionalProperties'], dict): additional_props: dict[str, Any] = s['additionalProperties'] # pyright: ignore[reportUnknownVariableType] _update_mapped_json_schema_refs(additional_props, name_mapping) # Handle unions and composition keywords for keyword in ['anyOf', 'oneOf', 'allOf']: if keyword in s: keyword_items: list[dict[str, Any]] = s[keyword] for item in keyword_items: _update_mapped_json_schema_refs(item, name_mapping) # Handle negation if 'not' in s and isinstance(s['not'], dict): not_schema: dict[str, Any] = s['not'] # pyright: ignore[reportUnknownVariableType] _update_mapped_json_schema_refs(not_schema, name_mapping) def _unique_def_name(name: str, schema: dict[str, Any], all_defs: dict[str, dict[str, Any]]) -> str: """Generate a unique definition name by appending the schema title and/or a numeric suffix.""" new_name = name if title := schema.get('title'): new_name = f'{title}_{name}' i = 1 original_new_name = new_name new_name = f'{new_name}_{i}' while new_name in all_defs: i += 1 new_name = f'{original_new_name}_{i}' return new_name def merge_json_schema_defs(schemas: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]]]: """Merges the `$defs` from different JSON schemas into a single deduplicated `$defs`, handling name collisions of `$defs` that are not the same, and rewrites `$ref`s to point to the new `$defs`. Returns a tuple of the rewritten schemas and a dictionary of the new `$defs`. """ all_defs: dict[str, dict[str, Any]] = {} rewritten_schemas: list[dict[str, Any]] = [] for schema in schemas: if '$defs' not in schema: rewritten_schemas.append(schema) continue schema = schema.copy() defs = schema.pop('$defs', None) schema_name_mapping: dict[str, str] = {} # Process definitions and build mapping for name, def_schema in defs.items(): if name not in all_defs: all_defs[name] = def_schema schema_name_mapping[name] = name elif def_schema != all_defs[name]: # Different def with same name — assign a unique name schema_name_mapping[name] = _unique_def_name(name, schema, all_defs) all_defs[schema_name_mapping[name]] = def_schema # else: structurally equal — handled below # Defs that are structurally equal (same dict) may still be semantically # different if they contain $refs that point to defs that were renamed in # this schema. E.g. both schemas have Wrapper={$ref Inner}, but their # Inner defs differ, so Schema B's Inner was renamed to Inner_1. The shared # Wrapper is not actually equal — Schema B needs its own copy with updated refs. # Loop until stable, since creating a copy can trigger further copies # in defs that reference it (transitive chains). changed = True while changed: changed = False for name, def_schema in defs.items(): if name not in schema_name_mapping: updated = copy.deepcopy(def_schema) _update_mapped_json_schema_refs(updated, schema_name_mapping) if updated != def_schema: schema_name_mapping[name] = _unique_def_name(name, schema, all_defs) all_defs[schema_name_mapping[name]] = updated changed = True else: schema_name_mapping[name] = name # Update refs inside definitions so internal cross-references # (e.g. Outer referencing Inner which was renamed to Inner_1) are corrected. for new_name in schema_name_mapping.values(): _update_mapped_json_schema_refs(all_defs[new_name], schema_name_mapping) _update_mapped_json_schema_refs(schema, schema_name_mapping) rewritten_schemas.append(schema) return rewritten_schemas, all_defs _MARKDOWN_FENCES_PATTERN = re.compile(r'```(?:\w+)?\r?\n(\{.*?\})\s*(?:\r?\n?```|\Z)', flags=re.DOTALL) def strip_markdown_fences(text: str) -> str: if text.startswith('{'): return text match = re.search(_MARKDOWN_FENCES_PATTERN, text) if match: return match.group(1) return text def _unwrap_annotated(tp: Any) -> Any: origin = get_origin(tp) while typing_objects.is_annotated(origin): tp = tp.__origin__ origin = get_origin(tp) return tp def get_union_args(tp: Any) -> tuple[Any, ...]: """Extract the arguments of a Union type if `tp` is a union, otherwise return an empty tuple.""" if typing_objects.is_typealiastype(tp): tp = tp.__value__ tp = _unwrap_annotated(tp) origin = get_origin(tp) if is_union_origin(origin): return tuple(_unwrap_annotated(arg) for arg in get_args(tp)) else: return () def get_event_loop() -> asyncio.AbstractEventLoop: try: event_loop = asyncio.get_event_loop() except RuntimeError: event_loop = None if event_loop is None or event_loop.is_closed(): event_loop = asyncio.new_event_loop() asyncio.set_event_loop(event_loop) return event_loop def is_str_dict(obj: Any) -> TypeGuard[dict[str, Any]]: """Check if obj is a dict, narrowing the type to `dict[str, Any]`.""" return isinstance(obj, dict) def is_text_like_media_type(media_type: str) -> bool: """Check if a media type represents text-like content. Returns True for `text/*`, JSON, XML, YAML, and their structured syntax suffixes. """ return ( media_type.startswith('text/') or media_type == 'application/json' or media_type.endswith('+json') or media_type == 'application/xml' or media_type.endswith('+xml') or media_type in ('application/x-yaml', 'application/yaml') ) def format_inlined_text_file(text: str, *, media_type: str, identifier: str) -> str: """Format text file content with delimiters for inlining into a text prompt.""" return '\n'.join( [ f'-----BEGIN FILE id="{identifier}" type="{media_type}"-----', text, f'-----END FILE id="{identifier}"-----', ] ) _TOKEN_SPLIT_PATTERN = re.compile(r'[\s",.:]+') def estimate_string_tokens(text: str) -> int: """Roughly estimate the number of tokens in a string by splitting on whitespace and punctuation. Shared by the test models, which report a plausible usage count without pulling in a tokenizer. Blank text counts as one token, so a caller that wants zero for it guards the call itself. """ return len(_TOKEN_SPLIT_PATTERN.split(text.strip())) def enum_member_docstrings(cls: type[Enum]) -> dict[str, str]: """The docstring under each member of an `Enum`, by member name. Pydantic reads a docstring under a model field with `use_attribute_docstrings`, but not one under an enum member; this does the same for enums, so each option can be described where it is declared. Empty when the source is not available, such as for a class defined in the REPL. """ try: source = inspect.getsource(cls) except (OSError, TypeError): return {} class_def = ast.parse(textwrap.dedent(source)).body[0] if not isinstance(class_def, ast.ClassDef): # pragma: no cover return {} docstrings: dict[str, str] = {} for previous, node in zip(class_def.body, class_def.body[1:]): if not ( isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant) and isinstance(node.value.value, str) ): continue # A member is a plain or an annotated assignment; a string after anything else describes no option, # and neither does one after a name that is not a member, such as `_ignore_`. if isinstance(previous, ast.Assign): targets = previous.targets elif isinstance(previous, ast.AnnAssign): targets = [previous.target] else: continue for name in [target.id for target in targets if isinstance(target, ast.Name) and target.id in cls.__members__]: docstrings[name] = inspect.cleandoc(node.value.value) return docstrings