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

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)
+        }

pydantic_ai/toolsets/function.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Sequence
+from dataclasses import dataclass, field, replace
+from typing import Any, Callable, overload
+
+from pydantic.json_schema import GenerateJsonSchema
+
+from .._run_context import AgentDepsT, RunContext
+from ..exceptions import UserError
+from ..tools import (
+    DocstringFormat,
+    GenerateToolJsonSchema,
+    Tool,
+    ToolFuncEither,
+    ToolParams,
+    ToolPrepareFunc,
+)
+from .abstract import AbstractToolset, ToolsetTool
+
+
+@dataclass
+class _FunctionToolsetTool(ToolsetTool[AgentDepsT]):
+    """A tool definition for a function toolset tool that keeps track of the function to call."""
+
+    call_func: Callable[[dict[str, Any], RunContext[AgentDepsT]], Awaitable[Any]]
+
+
+@dataclass(init=False)
+class FunctionToolset(AbstractToolset[AgentDepsT]):
+    """A toolset that lets Python functions be used as tools.
+
+    See [toolset docs](../toolsets.md#function-toolset) for more information.
+    """
+
+    max_retries: int = field(default=1)
+    tools: dict[str, Tool[Any]] = field(default_factory=dict)
+
+    def __init__(self, tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [], max_retries: int = 1):
+        """Build a new function toolset.
+
+        Args:
+            tools: The tools to add to the toolset.
+            max_retries: The maximum number of retries for each tool during a run.
+        """
+        self.max_retries = max_retries
+        self.tools = {}
+        for tool in tools:
+            if isinstance(tool, Tool):
+                self.add_tool(tool)
+            else:
+                self.add_function(tool)
+
+    @overload
+    def tool(self, func: ToolFuncEither[AgentDepsT, ToolParams], /) -> ToolFuncEither[AgentDepsT, ToolParams]: ...
+
+    @overload
+    def tool(
+        self,
+        /,
+        *,
+        name: str | None = None,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        strict: bool | None = None,
+    ) -> Callable[[ToolFuncEither[AgentDepsT, ToolParams]], ToolFuncEither[AgentDepsT, ToolParams]]: ...
+
+    def tool(
+        self,
+        func: ToolFuncEither[AgentDepsT, ToolParams] | None = None,
+        /,
+        *,
+        name: str | None = None,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        strict: bool | None = None,
+    ) -> Any:
+        """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
+
+        Can decorate a sync or async functions.
+
+        The docstring is inspected to extract both the tool description and description of each parameter,
+        [learn more](../tools.md#function-tools-and-schema).
+
+        We can't add overloads for every possible signature of tool, since the return type is a recursive union
+        so the signature of functions decorated with `@toolset.tool` is obscured.
+
+        Example:
+        ```python
+        from pydantic_ai import Agent, RunContext
+        from pydantic_ai.toolsets.function import FunctionToolset
+
+        toolset = FunctionToolset()
+
+        @toolset.tool
+        def foobar(ctx: RunContext[int], x: int) -> int:
+            return ctx.deps + x
+
+        @toolset.tool(retries=2)
+        async def spam(ctx: RunContext[str], y: float) -> float:
+            return ctx.deps + y
+
+        agent = Agent('test', toolsets=[toolset], deps_type=int)
+        result = agent.run_sync('foobar', deps=1)
+        print(result.output)
+        #> {"foobar":1,"spam":1.0}
+        ```
+
+        Args:
+            func: The tool function to register.
+            name: The name of the tool, defaults to the function name.
+            retries: The number of retries to allow for this tool, defaults to the agent's default retries,
+                which defaults to 1.
+            prepare: custom method to prepare the tool definition for each step, return `None` to omit this
+                tool from a given step. This is useful if you want to customise a tool at call time,
+                or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
+            docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
+                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
+            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
+            strict: Whether to enforce JSON schema compliance (only affects OpenAI).
+                See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+        """
+
+        def tool_decorator(
+            func_: ToolFuncEither[AgentDepsT, ToolParams],
+        ) -> ToolFuncEither[AgentDepsT, ToolParams]:
+            # noinspection PyTypeChecker
+            self.add_function(
+                func_,
+                None,
+                name,
+                retries,
+                prepare,
+                docstring_format,
+                require_parameter_descriptions,
+                schema_generator,
+                strict,
+            )
+            return func_
+
+        return tool_decorator if func is None else tool_decorator(func)
+
+    def add_function(
+        self,
+        func: ToolFuncEither[AgentDepsT, ToolParams],
+        takes_ctx: bool | None = None,
+        name: str | None = None,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDepsT] | None = None,
+        docstring_format: DocstringFormat = 'auto',
+        require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
+        strict: bool | None = None,
+    ) -> None:
+        """Add a function as a tool to the toolset.
+
+        Can take a sync or async function.
+
+        The docstring is inspected to extract both the tool description and description of each parameter,
+        [learn more](../tools.md#function-tools-and-schema).
+
+        Args:
+            func: The tool function to register.
+            takes_ctx: Whether the function takes a [`RunContext`][pydantic_ai.tools.RunContext] as its first argument. If `None`, this is inferred from the function signature.
+            name: The name of the tool, defaults to the function name.
+            retries: The number of retries to allow for this tool, defaults to the agent's default retries,
+                which defaults to 1.
+            prepare: custom method to prepare the tool definition for each step, return `None` to omit this
+                tool from a given step. This is useful if you want to customise a tool at call time,
+                or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
+            docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
+                Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
+            require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
+            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
+            strict: Whether to enforce JSON schema compliance (only affects OpenAI).
+                See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+        """
+        tool = Tool[AgentDepsT](
+            func,
+            takes_ctx=takes_ctx,
+            name=name,
+            max_retries=retries,
+            prepare=prepare,
+            docstring_format=docstring_format,
+            require_parameter_descriptions=require_parameter_descriptions,
+            schema_generator=schema_generator,
+            strict=strict,
+        )
+        self.add_tool(tool)
+
+    def add_tool(self, tool: Tool[AgentDepsT]) -> None:
+        """Add a tool to the toolset.
+
+        Args:
+            tool: The tool to add.
+        """
+        if tool.name in self.tools:
+            raise UserError(f'Tool name conflicts with existing tool: {tool.name!r}')
+        if tool.max_retries is None:
+            tool.max_retries = self.max_retries
+        self.tools[tool.name] = tool
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        tools: dict[str, ToolsetTool[AgentDepsT]] = {}
+        for original_name, tool in self.tools.items():
+            run_context = replace(ctx, tool_name=original_name, retry=ctx.retries.get(original_name, 0))
+            tool_def = await tool.prepare_tool_def(run_context)
+            if not tool_def:
+                continue
+
+            new_name = tool_def.name
+            if new_name in tools:
+                if new_name != original_name:
+                    raise UserError(f'Renaming tool {original_name!r} to {new_name!r} conflicts with existing tool.')
+                else:
+                    raise UserError(f'Tool name conflicts with previously renamed tool: {new_name!r}.')
+
+            tools[new_name] = _FunctionToolsetTool(
+                toolset=self,
+                tool_def=tool_def,
+                max_retries=tool.max_retries if tool.max_retries is not None else self.max_retries,
+                args_validator=tool.function_schema.validator,
+                call_func=tool.function_schema.call,
+            )
+        return tools
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        assert isinstance(tool, _FunctionToolsetTool)
+        return await tool.call_func(tool_args, ctx)

pydantic_ai/toolsets/prefixed.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Any
+
+from .._run_context import AgentDepsT, RunContext
+from .abstract import ToolsetTool
+from .wrapper import WrapperToolset
+
+
+@dataclass
+class PrefixedToolset(WrapperToolset[AgentDepsT]):
+    """A toolset that prefixes the names of the tools it contains.
+
+    See [toolset docs](../toolsets.md#prefixing-tool-names) for more information.
+    """
+
+    prefix: str
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        return {
+            new_name: replace(
+                tool,
+                toolset=self,
+                tool_def=replace(tool.tool_def, name=new_name),
+            )
+            for name, tool in (await super().get_tools(ctx)).items()
+            if (new_name := f'{self.prefix}_{name}')
+        }
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        original_name = name.removeprefix(self.prefix + '_')
+        ctx = replace(ctx, tool_name=original_name)
+        tool = replace(tool, tool_def=replace(tool.tool_def, name=original_name))
+        return await super().call_tool(original_name, tool_args, ctx, tool)

pydantic_ai/toolsets/prepared.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+
+from .._run_context import AgentDepsT, RunContext
+from ..exceptions import UserError
+from ..tools import ToolsPrepareFunc
+from .abstract import ToolsetTool
+from .wrapper import WrapperToolset
+
+
+@dataclass
+class PreparedToolset(WrapperToolset[AgentDepsT]):
+    """A toolset that prepares the tools it contains 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.
+    """
+
+    prepare_func: ToolsPrepareFunc[AgentDepsT]
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        original_tools = await super().get_tools(ctx)
+        original_tool_defs = [tool.tool_def for tool in original_tools.values()]
+        prepared_tool_defs_by_name = {
+            tool_def.name: tool_def for tool_def in (await self.prepare_func(ctx, original_tool_defs) or [])
+        }
+
+        if len(prepared_tool_defs_by_name.keys() - original_tools.keys()) > 0:
+            raise UserError(
+                'Prepare function cannot add or rename tools. Use `FunctionToolset.add_function()` or `RenamedToolset` instead.'
+            )
+
+        return {
+            name: replace(original_tools[name], tool_def=tool_def)
+            for name, tool_def in prepared_tool_defs_by_name.items()
+        }

pydantic_ai/toolsets/renamed.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Any
+
+from .._run_context import AgentDepsT, RunContext
+from .abstract import ToolsetTool
+from .wrapper import WrapperToolset
+
+
+@dataclass
+class RenamedToolset(WrapperToolset[AgentDepsT]):
+    """A toolset that renames the tools it contains using a dictionary mapping new names to original names.
+
+    See [toolset docs](../toolsets.md#renaming-tools) for more information.
+    """
+
+    name_map: dict[str, str]
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        original_to_new_name_map = {v: k for k, v in self.name_map.items()}
+        original_tools = await super().get_tools(ctx)
+        tools: dict[str, ToolsetTool[AgentDepsT]] = {}
+        for original_name, tool in original_tools.items():
+            new_name = original_to_new_name_map.get(original_name, None)
+            if new_name:
+                tools[new_name] = replace(
+                    tool,
+                    toolset=self,
+                    tool_def=replace(tool.tool_def, name=new_name),
+                )
+            else:
+                tools[original_name] = tool
+        return tools
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        original_name = self.name_map.get(name, name)
+        ctx = replace(ctx, tool_name=original_name)
+        tool = replace(tool, tool_def=replace(tool.tool_def, name=original_name))
+        return await super().call_tool(original_name, tool_args, ctx, tool)