Let toolsets be built dynamically based on run context (#2366)

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

Author: strawgate

docs/toolsets.md

@@ -5,10 +5,11 @@ A toolset represents a collection of [tools](tools.md) that can be registered wi
 
 Toolsets are used (among many other things) to define [MCP servers](mcp/client.md) available to an agent. Pydantic AI includes many kinds of toolsets which are described below, and you can define a [custom toolset](#building-a-custom-toolset) by inheriting from the [`AbstractToolset`][pydantic_ai.toolsets.AbstractToolset] class.
 
-The toolsets that will be available during an agent run can be specified in three different ways:
+The toolsets that will be available during an agent run can be specified in four different ways:
 
-* at agent construction time, via the [`toolsets`][pydantic_ai.Agent.__init__] keyword argument to `Agent`
-* at agent run time, via the `toolsets` keyword argument to [`agent.run()`][pydantic_ai.agent.AbstractAgent.run], [`agent.run_sync()`][pydantic_ai.agent.AbstractAgent.run_sync], [`agent.run_stream()`][pydantic_ai.agent.AbstractAgent.run_stream], or [`agent.iter()`][pydantic_ai.Agent.iter]. These toolsets will be additional to those provided to the `Agent` constructor
+* at agent construction time, via the [`toolsets`][pydantic_ai.Agent.__init__] keyword argument to `Agent`, which takes toolset instances as well as functions that generate toolsets [dynamically](#dynamically-building-a-toolset) based on the agent [run context][pydantic_ai.tools.RunContext]
+* at agent run time, via the `toolsets` keyword argument to [`agent.run()`][pydantic_ai.agent.AbstractAgent.run], [`agent.run_sync()`][pydantic_ai.agent.AbstractAgent.run_sync], [`agent.run_stream()`][pydantic_ai.agent.AbstractAgent.run_stream], or [`agent.iter()`][pydantic_ai.Agent.iter]. These toolsets will be additional to those registered on the `Agent`
+* [dynamically](#dynamically-building-a-toolset), via the [`@agent.toolset`][pydantic_ai.Agent.toolset] decorator which lets you build a toolset based on the agent [run context][pydantic_ai.tools.RunContext]
 * as a contextual override, via the `toolsets` keyword argument to the [`agent.override()`][pydantic_ai.Agent.iter] context manager. These toolsets will replace those provided at agent construction or run time during the life of the context manager
 
 ```python {title="toolsets.py"}
@@ -330,15 +331,11 @@ print(test_model.last_model_request_parameters.function_tools)
 
 1. We're using [`TestModel`][pydantic_ai.models.test.TestModel] here because it makes it easy to see which tools were available on each run.
 
-### Wrapping a Toolset
+### Changing Tool Execution
 
 [`WrapperToolset`][pydantic_ai.toolsets.WrapperToolset] wraps another toolset and delegates all responsibility to it.
 
-It is is a no-op by default, but enables some useful abilities:
-
-#### Changing Tool Execution
-
-You can subclass `WrapperToolset` to change the wrapped toolset's tool execution behavior by overriding the [`call_tool()`][pydantic_ai.toolsets.AbstractToolset.call_tool] method.
+It is is a no-op by default, but you can subclass `WrapperToolset` to change the wrapped toolset's tool execution behavior by overriding the [`call_tool()`][pydantic_ai.toolsets.AbstractToolset.call_tool] method.
 
 ```python {title="logging_toolset.py" requires="function_toolset.py,combined_toolset.py,renamed_toolset.py,prepared_toolset.py"}
 import asyncio
@@ -392,47 +389,68 @@ print(LOG)
 
 _(This example is complete, it can be run "as is")_
 
-#### Modifying Toolsets During a Run
+## Dynamically Building a Toolset
+
+Toolsets can be built dynamically ahead of each agent run or run step using a function that takes the agent [run context][pydantic_ai.tools.RunContext] and returns a toolset or `None`. This is useful when a toolset (like an MCP server) depends on information specific to an agent run, like its [dependencies](./dependencies.md).
+
+To register a dynamic toolset, you can pass a function that takes [`RunContext`][pydantic_ai.tools.RunContext] to the `toolsets` argument of the `Agent` constructor, or you can wrap a compliant function in the [`@agent.toolset`][pydantic_ai.Agent.toolset] decorator.
 
-You can change the `WrapperToolset`'s `wrapped` property during an agent run to swap out one toolset for another starting at the next run step.
+By default, the function will be called again ahead of each agent run step. If you are using the decorator, you can optionally provide a `per_run_step=False` argument to indicate that the toolset only needs to be built once for the entire run.
 
-To add or remove available toolsets, you can wrap a [`CombinedToolset`](#combining-toolsets) and replace it during the run with one that can include fewer, more, or entirely different toolsets.
+```python {title="dynamic_toolset.py", requires="function_toolset.py"}
+from dataclasses import dataclass
+from typing import Literal
 
-```python {title="wrapper_toolset.py" requires="function_toolset.py"}
 from function_toolset import weather_toolset, datetime_toolset
 
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.models.test import TestModel
-from pydantic_ai.toolsets import WrapperToolset
 
-togglable_toolset = WrapperToolset(weather_toolset)
 
-test_model = TestModel() # (1)!
+@dataclass
+class ToggleableDeps:
+    active: Literal['weather', 'datetime']
+
+    def toggle(self):
+        if self.active == 'weather':
+            self.active = 'datetime'
+        else:
+            self.active = 'weather'
+
+test_model = TestModel()  # (1)!
 agent = Agent(
     test_model,
-    deps_type=WrapperToolset # (2)!
+    deps_type=ToggleableDeps  # (2)!
 )
 
-@agent.tool
-def toggle(ctx: RunContext[WrapperToolset]):
-    if ctx.deps.wrapped == weather_toolset:
-        ctx.deps.wrapped = datetime_toolset
+@agent.toolset
+def toggleable_toolset(ctx: RunContext[ToggleableDeps]):
+    if ctx.deps.active == 'weather':
+        return weather_toolset
     else:
-        ctx.deps.wrapped = weather_toolset
+        return datetime_toolset
+
+@agent.tool
+def toggle(ctx: RunContext[ToggleableDeps]):
+    ctx.deps.toggle()
 
-result = agent.run_sync('Toggle the toolset', deps=togglable_toolset, toolsets=[togglable_toolset])
-print([t.name for t in test_model.last_model_request_parameters.function_tools]) # (3)!
+deps = ToggleableDeps('weather')
+
+result = agent.run_sync('Toggle the toolset', deps=deps)
+print([t.name for t in test_model.last_model_request_parameters.function_tools])  # (3)!
 #> ['toggle', 'now']
 
-result = agent.run_sync('Toggle the toolset', deps=togglable_toolset, toolsets=[togglable_toolset])
+result = agent.run_sync('Toggle the toolset', deps=deps)
 print([t.name for t in test_model.last_model_request_parameters.function_tools])
 #> ['toggle', 'temperature_celsius', 'temperature_fahrenheit', 'conditions']
 ```
 
 1. We're using [`TestModel`][pydantic_ai.models.test.TestModel] here because it makes it easy to see which tools were available on each run.
-2. We're using the agent's dependencies to give the `toggle` tool access to the `togglable_toolset` via the `RunContext` argument.
+2. We're using the agent's dependencies to give the `toggle` tool access to the `active` via the `RunContext` argument.
 3. This shows the available tools _after_ the `toggle` tool was executed, as the "last model request" was the one that returned the `toggle` tool result to the model.
 
+_(This example is complete, it can be run "as is")_
+
 ## Building a Custom Toolset
 
 To define a fully custom toolset with its own logic to list available tools and handle them being called, you can subclass [`AbstractToolset`][pydantic_ai.toolsets.AbstractToolset] and implement the [`get_tools()`][pydantic_ai.toolsets.AbstractToolset.get_tools] and [`call_tool()`][pydantic_ai.toolsets.AbstractToolset.call_tool] methods.

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -41,6 +41,9 @@ async def build(cls, toolset: AbstractToolset[AgentDepsT], ctx: RunContext[Agent
 
     async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
         """Build a new tool manager for the next run step, carrying over the retries from the current run step."""
+        if ctx.run_step == self.ctx.run_step:
+            return self
+
         retries = {
             failed_tool_name: self.ctx.retries.get(failed_tool_name, 0) + 1 for failed_tool_name in self.failed_tools
         }

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -14,7 +14,6 @@
 from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import TypeVar, deprecated
 
-from pydantic_ai.builtin_tools import AbstractBuiltinTool
 from pydantic_graph import Graph
 
 from .. import (
@@ -30,6 +29,7 @@
 from .._agent_graph import HistoryProcessor
 from .._output import OutputToolset
 from .._tool_manager import ToolManager
+from ..builtin_tools import AbstractBuiltinTool
 from ..models.instrumented import InstrumentationSettings, InstrumentedModel, instrument_model
 from ..output import OutputDataT, OutputSpec
 from ..profiles import ModelProfile
@@ -50,6 +50,10 @@
     ToolsPrepareFunc,
 )
 from ..toolsets import AbstractToolset
+from ..toolsets._dynamic import (
+    DynamicToolset,
+    ToolsetFunc,
+)
 from ..toolsets.combined import CombinedToolset
 from ..toolsets.function import FunctionToolset
 from ..toolsets.prepared import PreparedToolset
@@ -139,7 +143,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
     )
     _function_toolset: FunctionToolset[AgentDepsT] = dataclasses.field(repr=False)
     _output_toolset: OutputToolset[AgentDepsT] | None = dataclasses.field(repr=False)
-    _user_toolsets: Sequence[AbstractToolset[AgentDepsT]] = dataclasses.field(repr=False)
+    _user_toolsets: list[AbstractToolset[AgentDepsT]] = dataclasses.field(repr=False)
     _prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _prepare_output_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
@@ -171,7 +175,7 @@ def __init__(
         builtin_tools: Sequence[AbstractBuiltinTool] = (),
         prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         prepare_output_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | ToolsetFunc[AgentDepsT]] | None = None,
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
@@ -227,7 +231,7 @@ def __init__(
         builtin_tools: Sequence[AbstractBuiltinTool] = (),
         prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         prepare_output_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
-        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+        toolsets: Sequence[AbstractToolset[AgentDepsT] | ToolsetFunc[AgentDepsT]] | None = None,
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
@@ -265,7 +269,8 @@ def __init__(
             prepare_output_tools: Custom function to prepare the tool definition of all output tools for each step.
                 This is useful if you want to customize the definition of multiple output tools or you want to register
                 a subset of output tools for a given step. See [`ToolsPrepareFunc`][pydantic_ai.tools.ToolsPrepareFunc]
-            toolsets: Toolsets to register with the agent, including MCP servers.
+            toolsets: Toolsets to register with the agent, including MCP servers and functions which take a run context
+                and return a toolset. See [`ToolsetFunc`][pydantic_ai.toolsets.ToolsetFunc] for more information.
             defer_model_check: by default, if you provide a [named][pydantic_ai.models.KnownModelName] model,
                 it's evaluated to create a [`Model`][pydantic_ai.models.Model] instance immediately,
                 which checks for the necessary environment variables. Set this to `false`
@@ -341,7 +346,12 @@ def __init__(
             self._output_toolset.max_retries = self._max_result_retries
 
         self._function_toolset = _AgentFunctionToolset(tools, max_retries=self._max_tool_retries)
-        self._user_toolsets = toolsets or ()
+        self._dynamic_toolsets = [
+            DynamicToolset[AgentDepsT](toolset_func=toolset)
+            for toolset in toolsets or []
+            if not isinstance(toolset, AbstractToolset)
+        ]
+        self._user_toolsets = [toolset for toolset in toolsets or [] if isinstance(toolset, AbstractToolset)]
 
         self.history_processors = history_processors or []
 
@@ -1138,6 +1148,53 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams
 
         return tool_decorator if func is None else tool_decorator(func)
 
+    @overload
+    def toolset(self, func: ToolsetFunc[AgentDepsT], /) -> ToolsetFunc[AgentDepsT]: ...
+
+    @overload
+    def toolset(
+        self,
+        /,
+        *,
+        per_run_step: bool = True,
+    ) -> Callable[[ToolsetFunc[AgentDepsT]], ToolsetFunc[AgentDepsT]]: ...
+
+    def toolset(
+        self,
+        func: ToolsetFunc[AgentDepsT] | None = None,
+        /,
+        *,
+        per_run_step: bool = True,
+    ) -> Any:
+        """Decorator to register a toolset function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its only argument.
+
+        Can decorate a sync or async functions.
+
+        The decorator can be used bare (`agent.toolset`).
+
+        Example:
+        ```python
+        from pydantic_ai import Agent, RunContext
+        from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
+
+        agent = Agent('test', deps_type=str)
+
+        @agent.toolset
+        async def simple_toolset(ctx: RunContext[str]) -> AbstractToolset[str]:
+            return FunctionToolset()
+        ```
+
+        Args:
+            func: The toolset function to register.
+            per_run_step: Whether to re-evaluate the toolset for each run step. Defaults to True.
+        """
+
+        def toolset_decorator(func_: ToolsetFunc[AgentDepsT]) -> ToolsetFunc[AgentDepsT]:
+            self._dynamic_toolsets.append(DynamicToolset(func_, per_run_step=per_run_step))
+            return func_
+
+        return toolset_decorator if func is None else toolset_decorator(func)
+
     def _get_model(self, model: models.Model | models.KnownModelName | str | None) -> models.Model:
         """Create a model configured for this agent.
 
@@ -1197,10 +1254,10 @@ def _get_toolset(
 
         if some_user_toolsets := self._override_toolsets.get():
             user_toolsets = some_user_toolsets.value
-        elif additional is not None:
-            user_toolsets = [*self._user_toolsets, *additional]
         else:
-            user_toolsets = self._user_toolsets
+            # Copy the dynamic toolsets to ensure each run has its own instances
+            dynamic_toolsets = [dataclasses.replace(toolset) for toolset in self._dynamic_toolsets]
+            user_toolsets = [*self._user_toolsets, *dynamic_toolsets, *(additional or [])]
 
         if user_toolsets:
             toolset = CombinedToolset([function_toolset, *user_toolsets])

pydantic_ai_slim/pydantic_ai/tools.py

@@ -118,7 +118,6 @@ async def turn_on_strict_if_openai(
 Usage `ToolsPrepareFunc[AgentDepsT]`.
 """
 
-
 DocstringFormat = Literal['google', 'numpy', 'sphinx', 'auto']
 """Supported docstring formats.
 

pydantic_ai_slim/pydantic_ai/toolsets/__init__.py

@@ -1,3 +1,4 @@
+from ._dynamic import ToolsetFunc
 from .abstract import AbstractToolset, ToolsetTool
 from .combined import CombinedToolset
 from .deferred import DeferredToolset
@@ -10,6 +11,7 @@
 
 __all__ = (
     'AbstractToolset',
+    'ToolsetFunc',
     'ToolsetTool',
     'CombinedToolset',
     'DeferredToolset',

pydantic_ai_slim/pydantic_ai/toolsets/_dynamic.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Awaitable
+from dataclasses import dataclass, replace
+from typing import Any, Callable, Union
+
+from typing_extensions import Self, TypeAlias
+
+from .._run_context import AgentDepsT, RunContext
+from .abstract import AbstractToolset, ToolsetTool
+
+ToolsetFunc: TypeAlias = Callable[
+    [RunContext[AgentDepsT]],
+    Union[AbstractToolset[AgentDepsT], None, Awaitable[Union[AbstractToolset[AgentDepsT], None]]],
+]
+"""A sync/async function which takes a run context and returns a toolset."""
+
+
+@dataclass
+class DynamicToolset(AbstractToolset[AgentDepsT]):
+    """A toolset that dynamically builds a toolset using a function that takes the run context.
+
+    It should only be used during a single agent run as it stores the generated toolset.
+    To use it multiple times, copy it using `dataclasses.replace`.
+    """
+
+    toolset_func: ToolsetFunc[AgentDepsT]
+    per_run_step: bool = True
+
+    _toolset: AbstractToolset[AgentDepsT] | None = None
+    _run_step: int | None = None
+
+    @property
+    def id(self) -> str | None:
+        return None  # pragma: no cover
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        try:
+            if self._toolset is not None:
+                return await self._toolset.__aexit__(*args)
+        finally:
+            self._toolset = None
+            self._run_step = None
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        if self._toolset is None or (self.per_run_step and ctx.run_step != self._run_step):
+            if self._toolset is not None:
+                await self._toolset.__aexit__()
+
+            toolset = self.toolset_func(ctx)
+            if inspect.isawaitable(toolset):
+                toolset = await toolset
+
+            if toolset is not None:
+                await toolset.__aenter__()
+
+            self._toolset = toolset
+            self._run_step = ctx.run_step
+
+        if self._toolset is None:
+            return {}
+
+        return await self._toolset.get_tools(ctx)
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        assert self._toolset is not None
+        return await self._toolset.call_tool(name, tool_args, ctx, tool)
+
+    def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None:
+        if self._toolset is not None:
+            self._toolset.apply(visitor)
+
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
+    ) -> AbstractToolset[AgentDepsT]:
+        if self._toolset is None:
+            return super().visit_and_replace(visitor)
+        else:
+            return replace(self, _toolset=self._toolset.visit_and_replace(visitor))