pydantic-ai-slim 0.4.3__py3-none-any.whl → 0.4.4__py3-none-any.whl

pydantic_ai/tools.py

@@ -1,20 +1,15 @@
 from __future__ import annotations as _annotations
 
-import dataclasses
-import json
 from collections.abc import Awaitable, Sequence
 from dataclasses import dataclass, field
 from typing import Any, Callable, Generic, Literal, Union
 
-from opentelemetry.trace import Tracer
-from pydantic import ValidationError
 from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
 from pydantic_core import SchemaValidator, core_schema
 from typing_extensions import Concatenate, ParamSpec, Self, TypeAlias, TypeVar
 
-from . import _function_schema, _utils, messages as _messages
+from . import _function_schema, _utils
 from ._run_context import AgentDepsT, RunContext
-from .exceptions import ModelRetry, UnexpectedModelBehavior
 
 __all__ = (
     'AgentDepsT',
@@ -32,7 +27,6 @@ __all__ = (
     'ToolDefinition',
 )
 
-from .messages import ToolReturnPart
 
 ToolParams = ParamSpec('ToolParams', default=...)
 """Retrieval function param spec."""
@@ -173,12 +167,6 @@ class Tool(Generic[AgentDepsT]):
     This schema may be modified by the `prepare` function or by the Model class prior to including it in an API request.
     """
 
-    # TODO: Consider moving this current_retry state to live on something other than the tool.
-    #   We've worked around this for now by copying instances of the tool when creating new runs,
-    #   but this is a bit fragile. Moving the tool retry counts to live on the agent run state would likely clean things
-    #   up, though is also likely a larger effort to refactor.
-    current_retry: int = field(default=0, init=False)
-
     def __init__(
         self,
         function: ToolFuncEither[AgentDepsT],
@@ -303,6 +291,15 @@ class Tool(Generic[AgentDepsT]):
             function_schema=function_schema,
         )
 
+    @property
+    def tool_def(self):
+        return ToolDefinition(
+            name=self.name,
+            description=self.description,
+            parameters_json_schema=self.function_schema.json_schema,
+            strict=self.strict,
+        )
+
     async def prepare_tool_def(self, ctx: RunContext[AgentDepsT]) -> ToolDefinition | None:
         """Get the tool definition.
 
@@ -312,113 +309,11 @@ class Tool(Generic[AgentDepsT]):
         Returns:
             return a `ToolDefinition` or `None` if the tools should not be registered for this run.
         """
-        tool_def = ToolDefinition(
-            name=self.name,
-            description=self.description,
-            parameters_json_schema=self.function_schema.json_schema,
-            strict=self.strict,
-        )
+        base_tool_def = self.tool_def
         if self.prepare is not None:
-            return await self.prepare(ctx, tool_def)
+            return await self.prepare(ctx, base_tool_def)
         else:
-            return tool_def
-
-    async def run(
-        self,
-        message: _messages.ToolCallPart,
-        run_context: RunContext[AgentDepsT],
-        tracer: Tracer,
-        include_content: bool = False,
-    ) -> _messages.ToolReturnPart | _messages.RetryPromptPart:
-        """Run the tool function asynchronously.
-
-        This method wraps `_run` in an OpenTelemetry span.
-
-        See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>.
-        """
-        span_attributes = {
-            'gen_ai.tool.name': self.name,
-            # NOTE: this means `gen_ai.tool.call.id` will be included even if it was generated by pydantic-ai
-            'gen_ai.tool.call.id': message.tool_call_id,
-            **({'tool_arguments': message.args_as_json_str()} if include_content else {}),
-            'logfire.msg': f'running tool: {self.name}',
-            # add the JSON schema so these attributes are formatted nicely in Logfire
-            'logfire.json_schema': json.dumps(
-                {
-                    'type': 'object',
-                    'properties': {
-                        **(
-                            {
-                                'tool_arguments': {'type': 'object'},
-                                'tool_response': {'type': 'object'},
-                            }
-                            if include_content
-                            else {}
-                        ),
-                        'gen_ai.tool.name': {},
-                        'gen_ai.tool.call.id': {},
-                    },
-                }
-            ),
-        }
-        with tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
-            response = await self._run(message, run_context)
-            if include_content and span.is_recording():
-                span.set_attribute(
-                    'tool_response',
-                    response.model_response_str()
-                    if isinstance(response, ToolReturnPart)
-                    else response.model_response(),
-                )
-
-            return response
-
-    async def _run(
-        self, message: _messages.ToolCallPart, run_context: RunContext[AgentDepsT]
-    ) -> _messages.ToolReturnPart | _messages.RetryPromptPart:
-        try:
-            validator = self.function_schema.validator
-            if isinstance(message.args, str):
-                args_dict = validator.validate_json(message.args or '{}')
-            else:
-                args_dict = validator.validate_python(message.args or {})
-        except ValidationError as e:
-            return self._on_error(e, message)
-
-        ctx = dataclasses.replace(
-            run_context,
-            retry=self.current_retry,
-            tool_name=message.tool_name,
-            tool_call_id=message.tool_call_id,
-        )
-        try:
-            response_content = await self.function_schema.call(args_dict, ctx)
-        except ModelRetry as e:
-            return self._on_error(e, message)
-
-        self.current_retry = 0
-        return _messages.ToolReturnPart(
-            tool_name=message.tool_name,
-            content=response_content,
-            tool_call_id=message.tool_call_id,
-        )
-
-    def _on_error(
-        self, exc: ValidationError | ModelRetry, call_message: _messages.ToolCallPart
-    ) -> _messages.RetryPromptPart:
-        self.current_retry += 1
-        if self.max_retries is None or self.current_retry > self.max_retries:
-            raise UnexpectedModelBehavior(f'Tool exceeded max retries count of {self.max_retries}') from exc
-        else:
-            if isinstance(exc, ValidationError):
-                content = exc.errors(include_url=False, include_context=False)
-            else:
-                content = exc.message
-            return _messages.RetryPromptPart(
-                tool_name=call_message.tool_name,
-                content=content,
-                tool_call_id=call_message.tool_call_id,
-            )
+            return base_tool_def
 
 
 ObjectJsonSchema: TypeAlias = dict[str, Any]
@@ -429,6 +324,9 @@ This type is used to define tools parameters (aka arguments) in [ToolDefinition]
 With PEP-728 this should be a TypedDict with `type: Literal['object']`, and `extra_parts=Any`
 """
 
+ToolKind: TypeAlias = Literal['function', 'output', 'deferred']
+"""Kind of tool."""
+
 
 @dataclass(repr=False)
 class ToolDefinition:
@@ -440,7 +338,7 @@ class ToolDefinition:
     name: str
     """The name of the tool."""
 
-    parameters_json_schema: ObjectJsonSchema
+    parameters_json_schema: ObjectJsonSchema = field(default_factory=lambda: {'type': 'object', 'properties': {}})
     """The JSON schema for the tool's parameters."""
 
     description: str | None = None
@@ -464,4 +362,13 @@ class ToolDefinition:
     Note: this is currently only supported by OpenAI models.
     """
 
+    kind: ToolKind = field(default='function')
+    """The kind of tool:
+
+    - `'function'`: a tool that will be executed by Pydantic AI during an agent run and has its result returned to the model
+    - `'output'`: a tool that passes through an output value that ends the run
+    - `'deferred'`: a tool whose result will be produced outside of the Pydantic AI agent run in which it was called, because it depends on an upstream service (or user) or could take longer to generate than it's reasonable to keep the agent process running.
+        When the model calls a deferred tool, the agent run ends with a `DeferredToolCalls` object and a new run is expected to be started at a later point with the message history and new `ToolReturnPart`s corresponding to each deferred call.
+    """
+
     __repr__ = _utils.dataclasses_no_defaults_repr

pydantic_ai/toolsets/__init__.py

@@ -0,0 +1,22 @@
+from .abstract import AbstractToolset, ToolsetTool
+from .combined import CombinedToolset
+from .deferred import DeferredToolset
+from .filtered import FilteredToolset
+from .function import FunctionToolset
+from .prefixed import PrefixedToolset
+from .prepared import PreparedToolset
+from .renamed import RenamedToolset
+from .wrapper import WrapperToolset
+
+__all__ = (
+    'AbstractToolset',
+    'ToolsetTool',
+    'CombinedToolset',
+    'DeferredToolset',
+    'FilteredToolset',
+    'FunctionToolset',
+    'PrefixedToolset',
+    'RenamedToolset',
+    'PreparedToolset',
+    'WrapperToolset',
+)

pydantic_ai/toolsets/abstract.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol
+
+from pydantic_core import SchemaValidator
+from typing_extensions import Self
+
+from .._run_context import AgentDepsT, RunContext
+from ..tools import ToolDefinition, ToolsPrepareFunc
+
+if TYPE_CHECKING:
+    from .filtered import FilteredToolset
+    from .prefixed import PrefixedToolset
+    from .prepared import PreparedToolset
+    from .renamed import RenamedToolset
+
+
+class SchemaValidatorProt(Protocol):
+    """Protocol for a Pydantic Core `SchemaValidator` or `PluggableSchemaValidator` (which is private but API-compatible)."""
+
+    def validate_json(
+        self,
+        input: str | bytes | bytearray,
+        *,
+        allow_partial: bool | Literal['off', 'on', 'trailing-strings'] = False,
+        **kwargs: Any,
+    ) -> Any: ...
+
+    def validate_python(
+        self, input: Any, *, allow_partial: bool | Literal['off', 'on', 'trailing-strings'] = False, **kwargs: Any
+    ) -> Any: ...
+
+
+@dataclass
+class ToolsetTool(Generic[AgentDepsT]):
+    """Definition of a tool available on a toolset.
+
+    This is a wrapper around a plain tool definition that includes information about:
+
+    - the toolset that provided it, for use in error messages
+    - the maximum number of retries to attempt if the tool call fails
+    - the validator for the tool's arguments
+    """
+
+    toolset: AbstractToolset[AgentDepsT]
+    """The toolset that provided this tool, for use in error messages."""
+    tool_def: ToolDefinition
+    """The tool definition for this tool, including the name, description, and parameters."""
+    max_retries: int
+    """The maximum number of retries to attempt if the tool call fails."""
+    args_validator: SchemaValidator | SchemaValidatorProt
+    """The Pydantic Core validator for the tool's arguments.
+
+    For example, a [`pydantic.TypeAdapter(...).validator`](https://docs.pydantic.dev/latest/concepts/type_adapter/) or [`pydantic_core.SchemaValidator`](https://docs.pydantic.dev/latest/api/pydantic_core/#pydantic_core.SchemaValidator).
+    """
+
+
+class AbstractToolset(ABC, Generic[AgentDepsT]):
+    """A toolset is a collection of tools that can be used by an agent.
+
+    It is responsible for:
+
+    - Listing the tools it contains
+    - Validating the arguments of the tools
+    - Calling the tools
+
+    See [toolset docs](../toolsets.md) for more information.
+    """
+
+    @property
+    def name(self) -> str:
+        """The name of the toolset for use in error messages."""
+        return self.__class__.__name__.replace('Toolset', ' toolset')
+
+    @property
+    def tool_name_conflict_hint(self) -> str:
+        """A hint for how to avoid name conflicts with other toolsets for use in error messages."""
+        return 'Rename the tool or wrap the toolset in a `PrefixedToolset` to avoid name conflicts.'
+
+    async def __aenter__(self) -> Self:
+        """Enter the toolset context.
+
+        This is where you can set up network connections in a concrete implementation.
+        """
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        """Exit the toolset context.
+
+        This is where you can tear down network connections in a concrete implementation.
+        """
+        return None
+
+    @abstractmethod
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        """The tools that are available in this toolset."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        """Call a tool with the given arguments.
+
+        Args:
+            name: The name of the tool to call.
+            tool_args: The arguments to pass to the tool.
+            ctx: The run context.
+            tool: The tool definition returned by [`get_tools`][pydantic_ai.toolsets.AbstractToolset.get_tools] that was called.
+        """
+        raise NotImplementedError()
+
+    def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None:
+        """Run a visitor function on all concrete toolsets that are not wrappers (i.e. they implement their own tool listing and calling)."""
+        visitor(self)
+
+    def filtered(
+        self, filter_func: Callable[[RunContext[AgentDepsT], ToolDefinition], bool]
+    ) -> FilteredToolset[AgentDepsT]:
+        """Returns a new toolset that filters this toolset's tools using a filter function that takes the agent context and the tool definition.
+
+        See [toolset docs](../toolsets.md#filtering-tools) for more information.
+        """
+        from .filtered import FilteredToolset
+
+        return FilteredToolset(self, filter_func)
+
+    def prefixed(self, prefix: str) -> PrefixedToolset[AgentDepsT]:
+        """Returns a new toolset that prefixes the names of this toolset's tools.
+
+        See [toolset docs](../toolsets.md#prefixing-tool-names) for more information.
+        """
+        from .prefixed import PrefixedToolset
+
+        return PrefixedToolset(self, prefix)
+
+    def prepared(self, prepare_func: ToolsPrepareFunc[AgentDepsT]) -> PreparedToolset[AgentDepsT]:
+        """Returns a new toolset that prepares this toolset's tools using a prepare function that takes the agent context and the original tool definitions.
+
+        See [toolset docs](../toolsets.md#preparing-tool-definitions) for more information.
+        """
+        from .prepared import PreparedToolset
+
+        return PreparedToolset(self, prepare_func)
+
+    def renamed(self, name_map: dict[str, str]) -> RenamedToolset[AgentDepsT]:
+        """Returns a new toolset that renames this toolset's tools using a dictionary mapping new names to original names.
+
+        See [toolset docs](../toolsets.md#renaming-tools) for more information.
+        """
+        from .renamed import RenamedToolset
+
+        return RenamedToolset(self, name_map)

pydantic_ai/toolsets/combined.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Sequence
+from contextlib import AsyncExitStack
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+from typing_extensions import Self
+
+from .._run_context import AgentDepsT, RunContext
+from .._utils import get_async_lock
+from ..exceptions import UserError
+from .abstract import AbstractToolset, ToolsetTool
+
+
+@dataclass
+class _CombinedToolsetTool(ToolsetTool[AgentDepsT]):
+    """A tool definition for a combined toolset tools that keeps track of the source toolset and tool."""
+
+    source_toolset: AbstractToolset[AgentDepsT]
+    source_tool: ToolsetTool[AgentDepsT]
+
+
+@dataclass
+class CombinedToolset(AbstractToolset[AgentDepsT]):
+    """A toolset that combines multiple toolsets.
+
+    See [toolset docs](../toolsets.md#combining-toolsets) for more information.
+    """
+
+    toolsets: Sequence[AbstractToolset[AgentDepsT]]
+
+    _enter_lock: asyncio.Lock = field(compare=False, init=False)
+    _entered_count: int = field(init=False)
+    _exit_stack: AsyncExitStack | None = field(init=False)
+
+    def __post_init__(self):
+        self._enter_lock = get_async_lock()
+        self._entered_count = 0
+        self._exit_stack = None
+
+    async def __aenter__(self) -> Self:
+        async with self._enter_lock:
+            if self._entered_count == 0:
+                self._exit_stack = AsyncExitStack()
+                for toolset in self.toolsets:
+                    await self._exit_stack.enter_async_context(toolset)
+            self._entered_count += 1
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        async with self._enter_lock:
+            self._entered_count -= 1
+            if self._entered_count == 0 and self._exit_stack is not None:
+                await self._exit_stack.aclose()
+                self._exit_stack = None
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        toolsets_tools = await asyncio.gather(*(toolset.get_tools(ctx) for toolset in self.toolsets))
+        all_tools: dict[str, ToolsetTool[AgentDepsT]] = {}
+
+        for toolset, tools in zip(self.toolsets, toolsets_tools):
+            for name, tool in tools.items():
+                if existing_tools := all_tools.get(name):
+                    raise UserError(
+                        f'{toolset.name} defines a tool whose name conflicts with existing tool from {existing_tools.toolset.name}: {name!r}. {toolset.tool_name_conflict_hint}'
+                    )
+
+                all_tools[name] = _CombinedToolsetTool(
+                    toolset=tool.toolset,
+                    tool_def=tool.tool_def,
+                    max_retries=tool.max_retries,
+                    args_validator=tool.args_validator,
+                    source_toolset=toolset,
+                    source_tool=tool,
+                )
+        return all_tools
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        assert isinstance(tool, _CombinedToolsetTool)
+        return await tool.source_toolset.call_tool(name, tool_args, ctx, tool.source_tool)
+
+    def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None:
+        for toolset in self.toolsets:
+            toolset.apply(visitor)

pydantic_ai/toolsets/deferred.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Any
+
+from pydantic_core import SchemaValidator, core_schema
+
+from .._run_context import AgentDepsT, RunContext
+from ..tools import ToolDefinition
+from .abstract import AbstractToolset, ToolsetTool
+
+TOOL_SCHEMA_VALIDATOR = SchemaValidator(schema=core_schema.any_schema())
+
+
+@dataclass
+class DeferredToolset(AbstractToolset[AgentDepsT]):
+    """A toolset that holds deferred tools whose results will be produced outside of the Pydantic AI agent run in which they were called.
+
+    See [toolset docs](../toolsets.md#deferred-toolset), [`ToolDefinition.kind`][pydantic_ai.tools.ToolDefinition.kind], and [`DeferredToolCalls`][pydantic_ai.output.DeferredToolCalls] for more information.
+    """
+
+    tool_defs: list[ToolDefinition]
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        return {
+            tool_def.name: ToolsetTool(
+                toolset=self,
+                tool_def=replace(tool_def, kind='deferred'),
+                max_retries=0,
+                args_validator=TOOL_SCHEMA_VALIDATOR,
+            )
+            for tool_def in self.tool_defs
+        }
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        raise NotImplementedError('Deferred tools cannot be called')

pydantic_ai/toolsets/filtered.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+from .._run_context import AgentDepsT, RunContext
+from ..tools import ToolDefinition
+from .abstract import ToolsetTool
+from .wrapper import WrapperToolset
+
+
+@dataclass
+class FilteredToolset(WrapperToolset[AgentDepsT]):
+    """A toolset that filters the tools it contains using a filter function that takes the agent context and the tool definition.
+
+    See [toolset docs](../toolsets.md#filtering-tools) for more information.
+    """
+
+    filter_func: Callable[[RunContext[AgentDepsT], ToolDefinition], bool]
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        return {
+            name: tool for name, tool in (await super().get_tools(ctx)).items() if self.filter_func(ctx, tool.tool_def)
+        }