pydantic-ai-slim 0.8.1__py3-none-any.whl → 1.0.0b1__py3-none-any.whl

pydantic_ai/_griffe.py

@@ -2,9 +2,10 @@ from __future__ import annotations as _annotations
 
 import logging
 import re
+from collections.abc import Callable
 from contextlib import contextmanager
 from inspect import Signature
-from typing import TYPE_CHECKING, Any, Callable, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from griffe import Docstring, DocstringSectionKind, Object as GriffeObject
 

pydantic_ai/_otel_messages.py

@@ -5,10 +5,10 @@ Based on https://github.com/lmolkova/semantic-conventions/blob/eccd1f806e426a32c
 
 from __future__ import annotations
 
-from typing import Literal
+from typing import Literal, TypeAlias
 
 from pydantic import JsonValue
-from typing_extensions import NotRequired, TypeAlias, TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 
 class TextPart(TypedDict):

pydantic_ai/_output.py

@@ -3,9 +3,9 @@ from __future__ import annotations as _annotations
 import inspect
 import json
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable, Sequence
+from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload
 
 from pydantic import TypeAdapter, ValidationError
 from pydantic_core import SchemaValidator, to_json
@@ -15,7 +15,7 @@ from . import _function_schema, _utils, messages as _messages
 from ._run_context import AgentDepsT, RunContext
 from .exceptions import ModelRetry, ToolRetryError, UserError
 from .output import (
-    DeferredToolCalls,
+    DeferredToolRequests,
     NativeOutput,
     OutputDataT,
     OutputMode,
@@ -49,12 +49,12 @@ At some point, it may make sense to change the input to OutputValidatorFunc to b
 resolve these potential variance issues.
 """
 
-OutputValidatorFunc = Union[
-    Callable[[RunContext[AgentDepsT], OutputDataT_inv], OutputDataT_inv],
-    Callable[[RunContext[AgentDepsT], OutputDataT_inv], Awaitable[OutputDataT_inv]],
-    Callable[[OutputDataT_inv], OutputDataT_inv],
-    Callable[[OutputDataT_inv], Awaitable[OutputDataT_inv]],
-]
+OutputValidatorFunc = (
+    Callable[[RunContext[AgentDepsT], OutputDataT_inv], OutputDataT_inv]
+    | Callable[[RunContext[AgentDepsT], OutputDataT_inv], Awaitable[OutputDataT_inv]]
+    | Callable[[OutputDataT_inv], OutputDataT_inv]
+    | Callable[[OutputDataT_inv], Awaitable[OutputDataT_inv]]
+)
 """
 A function that always takes and returns the same type of data (which is the result type of an agent run), and:
 
@@ -196,7 +196,7 @@ class OutputValidator(Generic[AgentDepsT, OutputDataT_inv]):
 
 @dataclass
 class BaseOutputSchema(ABC, Generic[OutputDataT]):
-    allows_deferred_tool_calls: bool
+    allows_deferred_tools: bool
 
     @abstractmethod
     def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
@@ -249,10 +249,10 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
         """Build an OutputSchema dataclass from an output type."""
         raw_outputs = _flatten_output_spec(output_spec)
 
-        outputs = [output for output in raw_outputs if output is not DeferredToolCalls]
-        allows_deferred_tool_calls = len(outputs) < len(raw_outputs)
-        if len(outputs) == 0 and allows_deferred_tool_calls:
-            raise UserError('At least one output type must be provided other than `DeferredToolCalls`.')
+        outputs = [output for output in raw_outputs if output is not DeferredToolRequests]
+        allows_deferred_tools = len(outputs) < len(raw_outputs)
+        if len(outputs) == 0 and allows_deferred_tools:
+            raise UserError('At least one output type must be provided other than `DeferredToolRequests`.')
 
         if output := next((output for output in outputs if isinstance(output, NativeOutput)), None):
             if len(outputs) > 1:
@@ -265,7 +265,7 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
                     description=output.description,
                     strict=output.strict,
                 ),
-                allows_deferred_tool_calls=allows_deferred_tool_calls,
+                allows_deferred_tools=allows_deferred_tools,
             )
         elif output := next((output for output in outputs if isinstance(output, PromptedOutput)), None):
             if len(outputs) > 1:
@@ -278,7 +278,7 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
                     description=output.description,
                 ),
                 template=output.template,
-                allows_deferred_tool_calls=allows_deferred_tool_calls,
+                allows_deferred_tools=allows_deferred_tools,
             )
 
         text_outputs: Sequence[type[str] | TextOutput[OutputDataT]] = []
@@ -313,21 +313,21 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
 
             if toolset:
                 return ToolOrTextOutputSchema(
-                    processor=text_output_schema, toolset=toolset, allows_deferred_tool_calls=allows_deferred_tool_calls
+                    processor=text_output_schema,
+                    toolset=toolset,
+                    allows_deferred_tools=allows_deferred_tools,
                 )
             else:
-                return PlainTextOutputSchema(
-                    processor=text_output_schema, allows_deferred_tool_calls=allows_deferred_tool_calls
-                )
+                return PlainTextOutputSchema(processor=text_output_schema, allows_deferred_tools=allows_deferred_tools)
 
         if len(tool_outputs) > 0:
-            return ToolOutputSchema(toolset=toolset, allows_deferred_tool_calls=allows_deferred_tool_calls)
+            return ToolOutputSchema(toolset=toolset, allows_deferred_tools=allows_deferred_tools)
 
         if len(other_outputs) > 0:
             schema = OutputSchemaWithoutMode(
                 processor=cls._build_processor(other_outputs, name=name, description=description, strict=strict),
                 toolset=toolset,
-                allows_deferred_tool_calls=allows_deferred_tool_calls,
+                allows_deferred_tools=allows_deferred_tools,
             )
             if default_mode:
                 schema = schema.with_default_mode(default_mode)
@@ -371,23 +371,19 @@ class OutputSchemaWithoutMode(BaseOutputSchema[OutputDataT]):
         self,
         processor: ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT],
         toolset: OutputToolset[Any] | None,
-        allows_deferred_tool_calls: bool,
+        allows_deferred_tools: bool,
     ):
-        super().__init__(allows_deferred_tool_calls)
+        super().__init__(allows_deferred_tools)
         self.processor = processor
         self._toolset = toolset
 
     def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
         if mode == 'native':
-            return NativeOutputSchema(
-                processor=self.processor, allows_deferred_tool_calls=self.allows_deferred_tool_calls
-            )
+            return NativeOutputSchema(processor=self.processor, allows_deferred_tools=self.allows_deferred_tools)
         elif mode == 'prompted':
-            return PromptedOutputSchema(
-                processor=self.processor, allows_deferred_tool_calls=self.allows_deferred_tool_calls
-            )
+            return PromptedOutputSchema(processor=self.processor, allows_deferred_tools=self.allows_deferred_tools)
         elif mode == 'tool':
-            return ToolOutputSchema(toolset=self.toolset, allows_deferred_tool_calls=self.allows_deferred_tool_calls)
+            return ToolOutputSchema(toolset=self.toolset, allows_deferred_tools=self.allows_deferred_tools)
         else:
             assert_never(mode)
 
@@ -550,8 +546,8 @@ class PromptedOutputSchema(StructuredTextOutputSchema[OutputDataT]):
 class ToolOutputSchema(OutputSchema[OutputDataT]):
     _toolset: OutputToolset[Any] | None
 
-    def __init__(self, toolset: OutputToolset[Any] | None, allows_deferred_tool_calls: bool):
-        super().__init__(allows_deferred_tool_calls)
+    def __init__(self, toolset: OutputToolset[Any] | None, allows_deferred_tools: bool):
+        super().__init__(allows_deferred_tools)
         self._toolset = toolset
 
     @property
@@ -575,9 +571,9 @@ class ToolOrTextOutputSchema(ToolOutputSchema[OutputDataT], PlainTextOutputSchem
         self,
         processor: PlainTextOutputProcessor[OutputDataT] | None,
         toolset: OutputToolset[Any] | None,
-        allows_deferred_tool_calls: bool,
+        allows_deferred_tools: bool,
     ):
-        super().__init__(toolset=toolset, allows_deferred_tool_calls=allows_deferred_tool_calls)
+        super().__init__(toolset=toolset, allows_deferred_tools=allows_deferred_tools)
         self.processor = processor
 
     @property

pydantic_ai/_parts_manager.py

@@ -15,7 +15,7 @@ from __future__ import annotations as _annotations
 
 from collections.abc import Hashable
 from dataclasses import dataclass, field, replace
-from typing import Any, Union
+from typing import Any
 
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 from pydantic_ai.messages import (
@@ -38,7 +38,7 @@ VendorId = Hashable
 Type alias for a vendor identifier, which can be any hashable type (e.g., a string, UUID, etc.)
 """
 
-ManagedPart = Union[ModelResponsePart, ToolCallPartDelta]
+ManagedPart = ModelResponsePart | ToolCallPartDelta
 """
 A union of types that are managed by the ModelResponsePartsManager.
 Because many vendors have streaming APIs that may produce not-fully-formed tool calls,
@@ -262,14 +262,14 @@ class ModelResponsePartsManager:
             if tool_name is None and self._parts:
                 part_index = len(self._parts) - 1
                 latest_part = self._parts[part_index]
-                if isinstance(latest_part, (ToolCallPart, ToolCallPartDelta)):  # pragma: no branch
+                if isinstance(latest_part, ToolCallPart | ToolCallPartDelta):  # pragma: no branch
                     existing_matching_part_and_index = latest_part, part_index
         else:
             # vendor_part_id is provided, so look up the corresponding part or delta
             part_index = self._vendor_id_to_part_index.get(vendor_part_id)
             if part_index is not None:
                 existing_part = self._parts[part_index]
-                if not isinstance(existing_part, (ToolCallPartDelta, ToolCallPart)):
+                if not isinstance(existing_part, ToolCallPartDelta | ToolCallPart):
                     raise UnexpectedModelBehavior(f'Cannot apply a tool call delta to {existing_part=}')
                 existing_matching_part_and_index = existing_part, part_index
 

pydantic_ai/_run_context.py

@@ -18,7 +18,7 @@ AgentDepsT = TypeVar('AgentDepsT', default=None, contravariant=True)
 """Type variable for agent dependencies."""
 
 
-@dataclasses.dataclass(repr=False)
+@dataclasses.dataclass(repr=False, kw_only=True)
 class RunContext(Generic[AgentDepsT]):
     """Information about the current call."""
 
@@ -46,5 +46,7 @@ class RunContext(Generic[AgentDepsT]):
     """Number of retries so far."""
     run_step: int = 0
     """The current step in the run."""
+    tool_call_approved: bool = False
+    """Whether a tool call that required approval has now been approved."""
 
     __repr__ = _utils.dataclasses_no_defaults_repr

pydantic_ai/_system_prompt.py

@@ -1,9 +1,9 @@
 from __future__ import annotations as _annotations
 
 import inspect
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
-from typing import Any, Callable, Generic, cast
+from typing import Any, Generic, cast
 
 from . import _utils
 from ._run_context import AgentDepsT, RunContext

pydantic_ai/_tool_manager.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import json
-from collections.abc import Iterable
 from dataclasses import dataclass, field, replace
 from typing import Any, Generic
 
@@ -13,7 +12,6 @@ from . import messages as _messages
 from ._run_context import AgentDepsT, RunContext
 from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior
 from .messages import ToolCallPart
-from .output import DeferredToolCalls
 from .tools import ToolDefinition
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 
@@ -106,6 +104,9 @@ class ToolManager(Generic[AgentDepsT]):
                     msg = 'No tools available.'
                 raise ModelRetry(f'Unknown tool name: {name!r}. {msg}')
 
+            if tool.tool_def.defer:
+                raise RuntimeError('Deferred tools cannot be called')
+
             ctx = replace(
                 self.ctx,
                 tool_name=name,
@@ -204,23 +205,3 @@ class ToolManager(Generic[AgentDepsT]):
                 )
 
         return tool_result
-
-    def get_deferred_tool_calls(self, parts: Iterable[_messages.ModelResponsePart]) -> DeferredToolCalls | None:
-        """Get the deferred tool calls from the model response parts."""
-        deferred_calls_and_defs = [
-            (part, tool_def)
-            for part in parts
-            if isinstance(part, _messages.ToolCallPart)
-            and (tool_def := self.get_tool_def(part.tool_name))
-            and tool_def.kind == 'deferred'
-        ]
-        if not deferred_calls_and_defs:
-            return None
-
-        deferred_calls: list[_messages.ToolCallPart] = []
-        deferred_tool_defs: dict[str, ToolDefinition] = {}
-        for part, tool_def in deferred_calls_and_defs:
-            deferred_calls.append(part)
-            deferred_tool_defs[part.tool_name] = tool_def
-
-        return DeferredToolCalls(deferred_calls, deferred_tool_defs)

pydantic_ai/_utils.py

@@ -4,34 +4,28 @@ import asyncio
 import functools
 import inspect
 import re
-import sys
 import time
 import uuid
-import warnings
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterator
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, fields, is_dataclass
 from datetime import datetime, timezone
 from functools import partial
 from types import GenericAlias
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, overload
+from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeGuard, TypeVar, get_args, get_origin, overload
 
 from anyio.to_thread import run_sync
 from pydantic import BaseModel, TypeAdapter
 from pydantic.json_schema import JsonSchemaValue
 from typing_extensions import (
     ParamSpec,
-    TypeAlias,
-    TypeGuard,
     TypeIs,
-    get_args,
-    get_origin,
     is_typeddict,
 )
 from typing_inspection import typing_objects
 from typing_inspection.introspection import is_union_origin
 
-from pydantic_graph._utils import AbstractSpan, get_event_loop
+from pydantic_graph._utils import AbstractSpan
 
 from . import exceptions
 
@@ -96,7 +90,7 @@ class Some(Generic[T]):
     value: T
 
 
-Option: TypeAlias = Union[Some[T], None]
+Option: TypeAlias = Some[T] | None
 """Analogous to Rust's `Option` type, usage: `Option[Thing]` is equivalent to `Some[Thing] | None`."""
 
 
@@ -459,28 +453,22 @@ def strip_markdown_fences(text: str) -> str:
     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 get_args(tp)
+        return tuple(_unwrap_annotated(arg) for arg in get_args(tp))
     else:
         return ()
-
-
-# The `asyncio.Lock` `loop` argument was deprecated in 3.8 and removed in 3.10,
-# but 3.9 still needs it to have the intended behavior.
-
-if sys.version_info < (3, 10):
-
-    def get_async_lock() -> asyncio.Lock:  # pragma: lax no cover
-        with warnings.catch_warnings():
-            warnings.simplefilter('ignore', DeprecationWarning)
-            return asyncio.Lock(loop=get_event_loop())
-else:
-
-    def get_async_lock() -> asyncio.Lock:  # pragma: lax no cover
-        return asyncio.Lock()

pydantic_ai/ag_ui.py

@@ -8,12 +8,11 @@ from __future__ import annotations
 
 import json
 import uuid
-from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Callable, Iterable, Mapping, Sequence
 from dataclasses import Field, dataclass, replace
 from http import HTTPStatus
 from typing import (
     Any,
-    Callable,
     ClassVar,
     Final,
     Generic,
@@ -46,11 +45,11 @@ from .messages import (
     UserPromptPart,
 )
 from .models import KnownModelName, Model
-from .output import DeferredToolCalls, OutputDataT, OutputSpec
+from .output import OutputDataT, OutputSpec
 from .settings import ModelSettings
-from .tools import AgentDepsT, ToolDefinition
+from .tools import AgentDepsT, DeferredToolRequests, ToolDefinition
 from .toolsets import AbstractToolset
-from .toolsets.deferred import DeferredToolset
+from .toolsets.external import ExternalToolset
 from .usage import RunUsage, UsageLimits
 
 try:
@@ -343,7 +342,7 @@ async def run_ag_ui(
 
         async with agent.iter(
             user_prompt=None,
-            output_type=[output_type or agent.output_type, DeferredToolCalls],
+            output_type=[output_type or agent.output_type, DeferredToolRequests],
             message_history=messages,
             model=model,
             deps=deps,
@@ -515,7 +514,7 @@ async def _handle_tool_result_event(
     content = result.content
     if isinstance(content, BaseEvent):
         yield content
-    elif isinstance(content, (str, bytes)):  # pragma: no branch
+    elif isinstance(content, str | bytes):  # pragma: no branch
         # Avoid iterable check for strings and bytes.
         pass
     elif isinstance(content, Iterable):  # pragma: no branch
@@ -681,7 +680,7 @@ class _ToolCallNotFoundError(_RunError, ValueError):
         )
 
 
-class _AGUIFrontendToolset(DeferredToolset[AgentDepsT]):
+class _AGUIFrontendToolset(ExternalToolset[AgentDepsT]):
     def __init__(self, tools: list[AGUITool]):
         super().__init__(
             [