Let toolset factory be registered per run step or for entire run

Author: DouweM

pydantic_ai_slim/pydantic_ai/agent.py

@@ -16,6 +16,7 @@
 from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Literal, Never, Self, TypeIs, TypeVar, deprecated
 
+from pydantic_ai.toolsets._dynamic import DynamicToolset
 from pydantic_graph import End, Graph, GraphRun, GraphRunContext
 from pydantic_graph._utils import get_event_loop
 
@@ -164,8 +165,7 @@ class Agent(Generic[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)
-    _toolset_functions: Sequence[ToolsetFunc[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)
@@ -422,8 +422,9 @@ def __init__(
             self._output_toolset.max_retries = self._max_result_retries
 
         self._function_toolset = FunctionToolset(tools, max_retries=retries)
-        self._user_toolsets = [toolset for toolset in toolsets or [] if isinstance(toolset, AbstractToolset)]
-        self._toolset_functions = [toolset for toolset in toolsets or [] if not isinstance(toolset, AbstractToolset)]
+        self._user_toolsets = [
+            toolset if isinstance(toolset, AbstractToolset) else DynamicToolset(toolset) for toolset in toolsets or []
+        ]
 
         self.history_processors = history_processors or []
 
@@ -774,11 +775,9 @@ async def main():
             run_step=state.run_step,
         )
 
-        toolsets_from_functions = await self._materialize_toolset_functions(run_context)
-
         toolset = self._get_toolset(
             output_toolset=output_toolset,
-            additional_toolsets=[*(toolsets or []), *toolsets_from_functions],
+            additional_toolsets=toolsets,
         )
 
         # This will raise errors for any name conflicts
@@ -1632,11 +1631,24 @@ 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,
-        func: ToolsetFunc[AgentDepsT],
         /,
-    ) -> Callable[[ToolsetFunc[AgentDepsT]], ToolsetFunc[AgentDepsT]] | ToolsetFunc[AgentDepsT]:
+        *,
+        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.
 
         Optionally takes [`RunContext`][pydantic_ai.tools.RunContext] as its only argument.
@@ -1656,9 +1668,17 @@ 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.
         """
-        self._toolset_functions = [*self._toolset_functions, func]
-        return func
+
+        def toolset_decorator(func_: ToolsetFunc[AgentDepsT]) -> ToolsetFunc[AgentDepsT]:
+            self._user_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.
@@ -1780,20 +1800,6 @@ def _prepare_output_schema(
 
         return schema  # pyright: ignore[reportReturnType]
 
-    async def _materialize_toolset_functions(
-        self, run_context: RunContext[AgentDepsT]
-    ) -> list[AbstractToolset[AgentDepsT]]:
-        materialized_toolsets: list[AbstractToolset[AgentDepsT]] = []
-
-        for toolset_function in self._toolset_functions:
-            toolset = toolset_function(run_context)
-            if inspect.isawaitable(toolset):
-                materialized_toolsets.append(await toolset)
-            else:
-                materialized_toolsets.append(toolset)
-
-        return materialized_toolsets
-
     @staticmethod
     def is_model_request_node(
         node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]],

pydantic_ai_slim/pydantic_ai/toolsets/__init__.py

@@ -1,4 +1,5 @@
-from .abstract import AbstractToolset, ToolsetFunc, ToolsetTool
+from ._dynamic import ToolsetFunc
+from .abstract import AbstractToolset, ToolsetTool
 from .combined import CombinedToolset
 from .deferred import DeferredToolset
 from .filtered import FilteredToolset

pydantic_ai_slim/pydantic_ai/toolsets/abstract.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol
 
 from pydantic_core import SchemaValidator
-from typing_extensions import Self, TypeAlias
+from typing_extensions import Self
 
 from .._run_context import AgentDepsT, RunContext
 from ..tools import ToolDefinition, ToolsPrepareFunc
@@ -17,11 +16,6 @@
     from .prepared import PreparedToolset
     from .renamed import RenamedToolset
 
-ToolsetFunc: TypeAlias = (
-    'Callable[[RunContext[AgentDepsT]], AbstractToolset[AgentDepsT] | Awaitable[AbstractToolset[AgentDepsT]]]'
-)
-"""An sync/async function which takes a run context and returns a toolset."""
-
 
 class SchemaValidatorProt(Protocol):
     """Protocol for a Pydantic Core `SchemaValidator` or `PluggableSchemaValidator` (which is private but API-compatible)."""

tests/test_agent.py

@@ -1,6 +1,7 @@
 import json
 import re
 import sys
+from collections import defaultdict
 from dataclasses import dataclass
 from datetime import timezone
 from typing import Any, Callable, Union
@@ -3716,64 +3717,47 @@ async def prepare_tools(ctx: RunContext[None], tool_defs: list[ToolDefinition])
         available_tools = [tool_def.name for tool_def in tool_defs]
         return tool_defs
 
-    toolset_creation_count = 0
+    toolset_creation_counts: dict[str, int] = defaultdict(int)
 
-    def create_function_toolset(ctx: RunContext[None]) -> AbstractToolset[None]:
-        nonlocal toolset_creation_count
-        toolset_creation_count += 1
-        return toolset
+    def via_toolsets_arg(ctx: RunContext[None]) -> AbstractToolset[None]:
+        nonlocal toolset_creation_counts
+        toolset_creation_counts['via_toolsets_arg'] += 1
+        return toolset.prefixed('via_toolsets_arg')
 
     def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
         if len(messages) == 1:
-            return ModelResponse(parts=[ToolCallPart('foo')])
+            return ModelResponse(parts=[ToolCallPart('via_toolsets_arg_foo')])
         elif len(messages) == 3:
-            return ModelResponse(parts=[ToolCallPart('foo')])
+            return ModelResponse(parts=[ToolCallPart('via_toolset_decorator_foo')])
         else:
             return ModelResponse(parts=[TextPart('Done')])
 
-    agent = Agent(FunctionModel(respond), toolsets=[create_function_toolset], prepare_tools=prepare_tools)
-
-    run_result = agent.run_sync('Hello')
-
-    assert run_result._state.run_step == 3  # pyright: ignore[reportPrivateUsage]
-    assert len(available_tools) == 1
-    assert toolset_creation_count == 1
-
-
-async def test_toolset_decorator():
-    toolset = FunctionToolset()
-
-    agent: Agent[None, str] = Agent('test')
+    agent = Agent(FunctionModel(respond), toolsets=[via_toolsets_arg], prepare_tools=prepare_tools)
 
     @agent.toolset
-    def create_function_toolset(ctx: RunContext[None]) -> AbstractToolset[None]:
-        return toolset
-
-    async def create_function_toolset_async(ctx: RunContext[None]) -> AbstractToolset[None]:
-        return toolset
+    def via_toolset_decorator(ctx: RunContext[None]) -> AbstractToolset[None]:
+        nonlocal toolset_creation_counts
+        toolset_creation_counts['via_toolset_decorator'] += 1
+        return toolset.prefixed('via_toolset_decorator')
 
-    agent.toolset(create_function_toolset_async)
+    @agent.toolset(per_run_step=False)
+    async def via_toolset_decorator_for_entire_run(ctx: RunContext[None]) -> AbstractToolset[None]:
+        nonlocal toolset_creation_counts
+        toolset_creation_counts['via_toolset_decorator_for_entire_run'] += 1
+        return toolset.prefixed('via_toolset_decorator_for_entire_run')
 
-    agent_toolset_functions = agent._toolset_functions  # pyright: ignore[reportPrivateUsage]
-
-    assert len(agent_toolset_functions) == 2
-    assert agent_toolset_functions[0] is create_function_toolset
-    assert agent_toolset_functions[1] is create_function_toolset_async
+    run_result = agent.run_sync('Hello')
 
-    fake_run_context = RunContext(
-        deps=None,
-        model=TestModel(),
-        usage=Usage(),
-        prompt=None,
-        messages=[],
-        run_step=0,
+    assert run_result._state.run_step == 3  # pyright: ignore[reportPrivateUsage]
+    assert len(available_tools) == 3
+    assert toolset_creation_counts == snapshot(
+        {
+            'via_toolsets_arg': 4,
+            'via_toolset_decorator': 4,
+            'via_toolset_decorator_for_entire_run': 1,
+        }
     )
 
-    toolsets = await agent._materialize_toolset_functions(run_context=fake_run_context)  # pyright: ignore[reportPrivateUsage]
-    assert len(toolsets) == 2
-    assert isinstance(toolsets[0], AbstractToolset)
-    assert isinstance(toolsets[1], AbstractToolset)
-
 
 def test_adding_tools_during_run():
     toolset = FunctionToolset()