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

pydantic_ai/_function_schema.py

@@ -5,10 +5,10 @@ This module has to use numerous internal Pydantic APIs and is therefore brittle
 
 from __future__ import annotations as _annotations
 
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, Callable, Union, cast
+from typing import TYPE_CHECKING, Any, Concatenate, cast, get_origin
 
 from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -17,7 +17,7 @@ from pydantic.fields import FieldInfo
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
-from typing_extensions import Concatenate, ParamSpec, TypeIs, TypeVar, get_origin
+from typing_extensions import ParamSpec, TypeIs, TypeVar
 
 from ._griffe import doc_descriptions
 from ._run_context import RunContext
@@ -30,7 +30,7 @@ if TYPE_CHECKING:
 __all__ = ('function_schema',)
 
 
-@dataclass
+@dataclass(kw_only=True)
 class FunctionSchema:
     """Internal information about a function schema."""
 
@@ -231,7 +231,7 @@ R = TypeVar('R')
 
 WithCtx = Callable[Concatenate[RunContext[Any], P], R]
 WithoutCtx = Callable[P, R]
-TargetFunc = Union[WithCtx[P, R], WithoutCtx[P, R]]
+TargetFunc = WithCtx[P, R] | WithoutCtx[P, R]
 
 
 def _takes_ctx(function: TargetFunc[P, R]) -> TypeIs[WithCtx[P, R]]:

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,
@@ -154,6 +154,7 @@ class ModelResponsePartsManager:
         *,
         vendor_part_id: Hashable | None,
         content: str | None = None,
+        id: str | None = None,
         signature: str | None = None,
     ) -> ModelResponseStreamEvent:
         """Handle incoming thinking content, creating or updating a ThinkingPart in the manager as appropriate.
@@ -167,6 +168,7 @@ class ModelResponsePartsManager:
                 of thinking. If None, a new part will be created unless the latest part is already
                 a ThinkingPart.
             content: The thinking content to append to the appropriate ThinkingPart.
+            id: An optional id for the thinking part.
             signature: An optional signature for the thinking content.
 
         Returns:
@@ -197,7 +199,7 @@ class ModelResponsePartsManager:
             if content is not None:
                 # There is no existing thinking part that should be updated, so create a new one
                 new_part_index = len(self._parts)
-                part = ThinkingPart(content=content, signature=signature)
+                part = ThinkingPart(content=content, id=id, signature=signature)
                 if vendor_part_id is not None:  # pragma: no branch
                     self._vendor_id_to_part_index[vendor_part_id] = new_part_index
                 self._parts.append(part)
@@ -262,14 +264,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,9 +12,9 @@ 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
+from .usage import UsageLimits
 
 
 @dataclass
@@ -68,7 +67,11 @@ class ToolManager(Generic[AgentDepsT]):
             return None
 
     async def handle_call(
-        self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+        self,
+        call: ToolCallPart,
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+        usage_limits: UsageLimits | None = None,
     ) -> Any:
         """Handle a tool call by validating the arguments, calling the tool, and handling retries.
 
@@ -76,13 +79,14 @@ class ToolManager(Generic[AgentDepsT]):
             call: The tool call part to handle.
             allow_partial: Whether to allow partial validation of the tool arguments.
             wrap_validation_errors: Whether to wrap validation errors in a retry prompt part.
+            usage_limits: Optional usage limits to check before executing tools.
         """
         if self.tools is None or self.ctx is None:
             raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
 
         if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
-            # Output tool calls are not traced
-            return await self._call_tool(call, allow_partial, wrap_validation_errors)
+            # Output tool calls are not traced and not counted
+            return await self._call_tool(call, allow_partial, wrap_validation_errors, count_tool_usage=False)
         else:
             return await self._call_tool_traced(
                 call,
@@ -90,9 +94,17 @@ class ToolManager(Generic[AgentDepsT]):
                 wrap_validation_errors,
                 self.ctx.tracer,
                 self.ctx.trace_include_content,
+                usage_limits,
             )
 
-    async def _call_tool(self, call: ToolCallPart, allow_partial: bool, wrap_validation_errors: bool) -> Any:
+    async def _call_tool(
+        self,
+        call: ToolCallPart,
+        allow_partial: bool,
+        wrap_validation_errors: bool,
+        usage_limits: UsageLimits | None = None,
+        count_tool_usage: bool = True,
+    ) -> Any:
         if self.tools is None or self.ctx is None:
             raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
 
@@ -106,6 +118,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,
@@ -120,7 +135,15 @@ class ToolManager(Generic[AgentDepsT]):
             else:
                 args_dict = validator.validate_python(call.args or {}, allow_partial=pyd_allow_partial)
 
-            return await self.toolset.call_tool(name, args_dict, ctx, tool)
+            if usage_limits is not None and count_tool_usage:
+                usage_limits.check_before_tool_call(self.ctx.usage)
+
+            result = await self.toolset.call_tool(name, args_dict, ctx, tool)
+
+            if count_tool_usage:
+                self.ctx.usage.tool_calls += 1
+
+            return result
         except (ValidationError, ModelRetry) as e:
             max_retries = tool.max_retries if tool is not None else 1
             current_retry = self.ctx.retries.get(name, 0)
@@ -159,6 +182,7 @@ class ToolManager(Generic[AgentDepsT]):
         wrap_validation_errors: bool,
         tracer: Tracer,
         include_content: bool = False,
+        usage_limits: UsageLimits | None = None,
     ) -> Any:
         """See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
         span_attributes = {
@@ -188,7 +212,7 @@ class ToolManager(Generic[AgentDepsT]):
         }
         with tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
             try:
-                tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
+                tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits)
             except ToolRetryError as e:
                 part = e.tool_retry
                 if include_content and span.is_recording():
@@ -204,23 +228,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()