Add prepare_tools param to Agent class (#1474)

Co-authored-by: Douwe Maan <[email protected]>
Co-authored-by: Marcelo Trylesinski <[email protected]>

Author: 3 people

docs/tools.md

@@ -553,6 +553,100 @@ print(test_model.last_model_request_parameters.function_tools)
 
 _(This example is complete, it can be run "as is")_
 
+## Agent-wide Dynamic Tool Preparation {#prepare-tools}
+
+In addition to per-tool `prepare` methods, you can also define an agent-wide `prepare_tools` function. This function is called at each step of a run and allows you to filter or modify the list of all tool definitions available to the agent for that step. This is especially useful if you want to enable or disable multiple tools at once, or apply global logic based on the current context.
+
+The `prepare_tools` function should be of type [`ToolsPrepareFunc`][pydantic_ai.tools.ToolsPrepareFunc], which takes the [`RunContext`][pydantic_ai.tools.RunContext] and a list of [`ToolDefinition`][pydantic_ai.tools.ToolDefinition], and returns a new list of tool definitions (or `None` to disable all tools for that step).
+
+!!! note
+    The list of tool definitions passed to `prepare_tools` includes both regular tools and tools from any MCP servers attached to the agent.
+
+Here's an example that makes all tools strict if the model is an OpenAI model:
+
+```python {title="agent_prepare_tools_customize.py" noqa="I001"}
+from dataclasses import replace
+from typing import Union
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.models.test import TestModel
+
+
+async def turn_on_strict_if_openai(
+    ctx: RunContext[None], tool_defs: list[ToolDefinition]
+) -> Union[list[ToolDefinition], None]:
+    if ctx.model.system == 'openai':
+        return [replace(tool_def, strict=True) for tool_def in tool_defs]
+    return tool_defs
+
+
+test_model = TestModel()
+agent = Agent(test_model, prepare_tools=turn_on_strict_if_openai)
+
+
+@agent.tool_plain
+def echo(message: str) -> str:
+    return message
+
+
+agent.run_sync('testing...')
+assert test_model.last_model_request_parameters.function_tools[0].strict is None
+
+# Set the system attribute of the test_model to 'openai'
+test_model._system = 'openai'
+
+agent.run_sync('testing with openai...')
+assert test_model.last_model_request_parameters.function_tools[0].strict
+```
+
+_(This example is complete, it can be run "as is")_
+
+Here's another example that conditionally filters out the tools by name if the dependency (`ctx.deps`) is `True`:
+
+```python {title="agent_prepare_tools_filter_out.py" noqa="I001"}
+from typing import Union
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.tools import Tool, ToolDefinition
+
+
+def launch_potato(target: str) -> str:
+    return f'Potato launched at {target}!'
+
+
+async def filter_out_tools_by_name(
+    ctx: RunContext[bool], tool_defs: list[ToolDefinition]
+) -> Union[list[ToolDefinition], None]:
+    if ctx.deps:
+        return [tool_def for tool_def in tool_defs if tool_def.name != 'launch_potato']
+    return tool_defs
+
+
+agent = Agent(
+    'test',
+    tools=[Tool(launch_potato)],
+    prepare_tools=filter_out_tools_by_name,
+    deps_type=bool,
+)
+
+result = agent.run_sync('testing...', deps=False)
+print(result.output)
+#> {"launch_potato":"Potato launched at a!"}
+result = agent.run_sync('testing...', deps=True)
+print(result.output)
+#> success (no tool calls)
+```
+
+_(This example is complete, it can be run "as is")_
+
+You can use `prepare_tools` to:
+
+- Dynamically enable or disable tools based on the current model, dependencies, or other context
+- Modify tool definitions globally (e.g., set all tools to strict mode, change descriptions, etc.)
+
+If both per-tool `prepare` and agent-wide `prepare_tools` are used, the per-tool `prepare` is applied first to each tool, and then `prepare_tools` is called with the resulting list of tool definitions.
+
 
 ## Tool Execution and Retries {#tool-retries}
 

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -26,7 +26,7 @@
 )
 from .result import OutputDataT, ToolOutput
 from .settings import ModelSettings, merge_model_settings
-from .tools import RunContext, Tool, ToolDefinition
+from .tools import RunContext, Tool, ToolDefinition, ToolsPrepareFunc
 
 if TYPE_CHECKING:
     from .mcp import MCPServer
@@ -97,6 +97,8 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
 
     tracer: Tracer
 
+    prepare_tools: ToolsPrepareFunc[DepsT] | None = None
+
 
 class AgentNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]]):
     """The base class for all agent nodes.
@@ -241,6 +243,11 @@ async def add_mcp_server_tools(server: MCPServer) -> None:
         *map(add_mcp_server_tools, ctx.deps.mcp_servers),
     )
 
+    if ctx.deps.prepare_tools:
+        # Prepare the tools using the provided function
+        # This also acts over tool definitions pulled from MCP servers
+        function_tool_defs = await ctx.deps.prepare_tools(run_context, function_tool_defs) or []
+
     output_schema = ctx.deps.output_schema
     return models.ModelRequestParameters(
         function_tools=function_tool_defs,

pydantic_ai_slim/pydantic_ai/agent.py

@@ -42,6 +42,7 @@
     ToolFuncPlain,
     ToolParams,
     ToolPrepareFunc,
+    ToolsPrepareFunc,
 )
 
 # Re-exporting like this improves auto-import behavior in PyCharm
@@ -131,6 +132,11 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     The type of data output by agent runs, used to validate the data returned by the model, defaults to `str`.
     """
 
+    prepare_tools: ToolsPrepareFunc[AgentDepsT] | None
+    """
+    Function invoked on each step, allowing the tools to be modified and filtered out as needed.
+    """
+
     instrument: InstrumentationSettings | bool | None
     """Options to automatically instrument with OpenTelemetry."""
 
@@ -172,6 +178,7 @@ def __init__(
         retries: int = 1,
         output_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -200,6 +207,7 @@ def __init__(
         result_tool_description: str | None = None,
         result_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -223,6 +231,7 @@ def __init__(
         retries: int = 1,
         output_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -251,6 +260,9 @@ def __init__(
             output_retries: The maximum number of retries to allow for result validation, defaults to `retries`.
             tools: Tools to register with the agent, you can also register tools via the decorators
                 [`@agent.tool`][pydantic_ai.Agent.tool] and [`@agent.tool_plain`][pydantic_ai.Agent.tool_plain].
+            prepare_tools: custom method to prepare the tool definition of all tools for each step.
+                This is useful if you want to customize the definition of multiple tools or you want to register
+                a subset of tools for a given step. See [`ToolsPrepareFunc`][pydantic_ai.tools.ToolsPrepareFunc]
             mcp_servers: MCP servers to register with the agent. You should register a [`MCPServer`][pydantic_ai.mcp.MCPServer]
                 for each server you want the agent to connect to.
             defer_model_check: by default, if you provide a [named][pydantic_ai.models.KnownModelName] model,
@@ -334,6 +346,7 @@ def __init__(
         self._default_retries = retries
         self._max_result_retries = output_retries if output_retries is not None else retries
         self._mcp_servers = mcp_servers
+        self._prepare_tools = prepare_tools
         for tool in tools:
             if isinstance(tool, Tool):
                 self._register_tool(tool)
@@ -694,6 +707,7 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
             mcp_servers=self._mcp_servers,
             default_retries=self._default_retries,
             tracer=tracer,
+            prepare_tools=self._prepare_tools,
             get_instructions=get_instructions,
         )
         start_node = _agent_graph.UserPromptNode[AgentDepsT](

pydantic_ai_slim/pydantic_ai/tools.py

@@ -29,6 +29,7 @@
     'ToolFuncEither',
     'ToolParams',
     'ToolPrepareFunc',
+    'ToolsPrepareFunc',
     'Tool',
     'ObjectJsonSchema',
     'ToolDefinition',
@@ -133,6 +134,37 @@ def hitchhiker(ctx: RunContext[int], answer: str) -> str:
 Usage `ToolPrepareFunc[AgentDepsT]`.
 """
 
+ToolsPrepareFunc: TypeAlias = (
+    'Callable[[RunContext[AgentDepsT], list[ToolDefinition]], Awaitable[list[ToolDefinition] | None]]'
+)
+"""Definition of a function that can prepare the tool definition of all tools for each step.
+This is useful if you want to customize the definition of multiple tools or you want to register
+a subset of tools for a given step.
+
+Example — here `turn_on_strict_if_openai` is valid as a `ToolsPrepareFunc`:
+
+```python {noqa="I001"}
+from dataclasses import replace
+from typing import Union
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.tools import ToolDefinition
+
+
+async def turn_on_strict_if_openai(
+    ctx: RunContext[None], tool_defs: list[ToolDefinition]
+) -> Union[list[ToolDefinition], None]:
+    if ctx.model.system == 'openai':
+        return [replace(tool_def, strict=True) for tool_def in tool_defs]
+    return tool_defs
+
+agent = Agent('openai:gpt-4o', prepare_tools=turn_on_strict_if_openai)
+```
+
+Usage `ToolsPrepareFunc[AgentDepsT]`.
+"""
+
+
 DocstringFormat = Literal['google', 'numpy', 'sphinx', 'auto']
 """Supported docstring formats.
 

tests/test_tools.py

@@ -1,5 +1,5 @@
 import json
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Annotated, Any, Callable, Literal, Union
 
 import pydantic_core
@@ -954,3 +954,37 @@ def get_score(data: Data) -> int: ...  # pragma: no branch
             'strict': None,
         }
     )
+
+
+def test_dynamic_tools_agent_wide():
+    async def prepare_tool_defs(
+        ctx: RunContext[int], tool_defs: list[ToolDefinition]
+    ) -> Union[list[ToolDefinition], None]:
+        if ctx.deps == 42:
+            return []
+        elif ctx.deps == 43:
+            return None
+        elif ctx.deps == 21:
+            return [replace(tool_def, strict=True) for tool_def in tool_defs]
+        return tool_defs
+
+    agent = Agent('test', deps_type=int, prepare_tools=prepare_tool_defs)
+
+    @agent.tool
+    def foobar(ctx: RunContext[int], x: int, y: str) -> str:
+        return f'{ctx.deps} {x} {y}'
+
+    result = agent.run_sync('', deps=42)
+    assert result.output == snapshot('success (no tool calls)')
+
+    result = agent.run_sync('', deps=43)
+    assert result.output == snapshot('success (no tool calls)')
+
+    with agent.override(model=FunctionModel(get_json_schema)):
+        result = agent.run_sync('', deps=21)
+        json_schema = json.loads(result.output)
+        assert agent._function_tools['foobar'].strict is None
+        assert json_schema['strict'] is True
+
+    result = agent.run_sync('', deps=1)
+    assert result.output == snapshot('{"foobar":"1 0 a"}')