Add Dynamic Toolset Decorator

Author: strawgate

pydantic_ai_slim/pydantic_ai/agent.py

@@ -51,7 +51,7 @@
     ToolPrepareFunc,
     ToolsPrepareFunc,
 )
-from .toolsets import AbstractToolset
+from .toolsets import AbstractToolset, ToolsetFunc
 from .toolsets.combined import CombinedToolset
 from .toolsets.function import FunctionToolset
 from .toolsets.prepared import PreparedToolset
@@ -165,6 +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)
     _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)
@@ -192,7 +193,7 @@ def __init__(
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
         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,
@@ -223,7 +224,7 @@ def __init__(
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
         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,
@@ -278,7 +279,7 @@ def __init__(
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
         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,
@@ -421,7 +422,8 @@ def __init__(
             self._output_toolset.max_retries = self._max_result_retries
 
         self._function_toolset = FunctionToolset(tools, max_retries=retries)
-        self._user_toolsets = toolsets or ()
+        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.history_processors = history_processors or []
 
@@ -772,7 +774,11 @@ async def main():
             run_step=state.run_step,
         )
 
-        toolset = self._get_toolset(output_toolset=output_toolset, additional_toolsets=toolsets)
+        toolset = self._get_toolset(
+            output_toolset=output_toolset,
+            additional_toolsets=[*(toolsets or []), *[func(run_context) for func in self._toolset_functions]],
+        )
+
         # This will raise errors for any name conflicts
         async with toolset:
             run_toolset = await ToolManager[AgentDepsT].build(toolset, run_context)
@@ -1624,6 +1630,49 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams
 
         return tool_decorator if func is None else tool_decorator(func)
 
+    def toolset(
+        self,
+        func: ToolsetFunc[AgentDepsT] | None = None,
+        /,
+    ) -> Callable[[ToolsetFunc[AgentDepsT]], ToolsetFunc[AgentDepsT]] | ToolsetFunc[AgentDepsT]:
+        """Decorator to register a toolset function.
+
+        Optionally 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`).
+
+        Overloads for every possible signature of `toolset` are included so the decorator doesn't obscure
+        the type of the function.
+
+        Example:
+        ```python
+        from pydantic_ai import Agent, RunContext
+
+        agent = Agent('test', deps_type=str)
+
+        @agent.toolset
+        def simple_toolset(ctx: RunContext[str]) -> AbstractToolset[str]:
+            return FunctionToolset(foobar)
+
+        @agent.toolset
+        async def async_toolset(ctx: RunContext[str]) -> AbstractToolset[str]:
+            return FunctionToolset(foobar)
+        ```
+        """
+        if func is None:
+
+            def decorator(
+                func_: ToolsetFunc[AgentDepsT],
+            ) -> ToolsetFunc[AgentDepsT]:
+                self._toolset_functions = [*self._toolset_functions, func_]
+                return func_
+
+            return decorator
+        else:
+            self._toolset_functions = [*self._toolset_functions, func]
+            return func
+
     def _get_model(self, model: models.Model | models.KnownModelName | str | None) -> models.Model:
         """Create a model configured for this agent.
 
@@ -1672,7 +1721,7 @@ def _get_toolset(
         self,
         output_toolset: AbstractToolset[AgentDepsT] | None | _utils.Unset = _utils.UNSET,
         additional_toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
-    ) -> AbstractToolset[AgentDepsT]:
+    ) -> CombinedToolset[AgentDepsT]:
         """Get the complete toolset.
 
         Args:

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,4 +1,4 @@
-from .abstract import AbstractToolset, ToolsetTool
+from .abstract import AbstractToolset, ToolsetFunc, ToolsetTool
 from .combined import CombinedToolset
 from .deferred import DeferredToolset
 from .filtered import FilteredToolset
@@ -10,6 +10,7 @@
 
 __all__ = (
     'AbstractToolset',
+    'ToolsetFunc',
     'ToolsetTool',
     'CombinedToolset',
     'DeferredToolset',

pydantic_ai_slim/pydantic_ai/toolsets/abstract.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol, TypeAlias
 
 from pydantic_core import SchemaValidator
 from typing_extensions import Self
@@ -17,6 +17,10 @@
     from .renamed import RenamedToolset
 
 
+ToolsetFunc: TypeAlias = 'Callable[[RunContext[AgentDepsT]], AbstractToolset[AgentDepsT]]'
+"""Definition of a function that returns a toolset based on the run context."""
+
+
 class SchemaValidatorProt(Protocol):
     """Protocol for a Pydantic Core `SchemaValidator` or `PluggableSchemaValidator` (which is private but API-compatible)."""
 

pydantic_ai_slim/pydantic_ai/toolsets/dynamic.py

@@ -46,6 +46,7 @@
 from pydantic_ai.profiles import ModelProfile
 from pydantic_ai.result import Usage
 from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.toolsets.abstract import AbstractToolset
 from pydantic_ai.toolsets.combined import CombinedToolset
 from pydantic_ai.toolsets.function import FunctionToolset
 from pydantic_ai.toolsets.prefixed import PrefixedToolset
@@ -3701,6 +3702,44 @@ def bar() -> str:
     assert result.output == snapshot('{"baz":"Hello from baz"}')
 
 
+def test_toolset_decorator():
+    toolset = FunctionToolset()
+
+    @toolset.tool
+    def foo() -> str:
+        return 'Hello from foo'
+
+    available_tools: list[str] = []
+
+    async def prepare_tools(ctx: RunContext[None], tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
+        nonlocal available_tools
+        available_tools = [tool_def.name for tool_def in tool_defs]
+        return tool_defs
+
+    toolset_creation_count = 0
+
+    def create_function_toolset(ctx: RunContext[None]) -> AbstractToolset[None]:
+        nonlocal toolset_creation_count
+        toolset_creation_count += 1
+        return toolset
+
+    def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        if len(messages) == 1:
+            return ModelResponse(parts=[ToolCallPart('foo')])
+        elif len(messages) == 3:
+            return ModelResponse(parts=[ToolCallPart('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
+
+
 def test_adding_tools_during_run():
     toolset = FunctionToolset()
 

tests/test_agent.py

@@ -3,7 +3,7 @@
 import re
 from collections import defaultdict
 from dataclasses import dataclass, replace
-from typing import Literal, TypeVar
+from typing import TypeVar
 from unittest.mock import AsyncMock
 
 import pytest
@@ -15,9 +15,7 @@
 from pydantic_ai.messages import ToolCallPart
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.tools import ToolDefinition
-from pydantic_ai.toolsets.abstract import AbstractToolset
 from pydantic_ai.toolsets.combined import CombinedToolset
-from pydantic_ai.toolsets.dynamic import DynamicToolset
 from pydantic_ai.toolsets.filtered import FilteredToolset
 from pydantic_ai.toolsets.function import FunctionToolset
 from pydantic_ai.toolsets.prefixed import PrefixedToolset
@@ -628,97 +626,3 @@ def tool_c(x: int) -> int:
 
     assert new_tool_manager.ctx.retries == {'tool_a': 1, 'tool_b': 1}
     assert new_tool_manager.failed_tools == set()  # reset for new run step
-
-
-async def test_dynamic_toolset():
-    """Test that the dynamic toolset correctly handles nested context managers."""
-
-    something_context: RunContext[str] = build_run_context(deps='something')
-    something_else_context: RunContext[str] = build_run_context(deps='something_else')
-    nothing_context: RunContext[str] = build_run_context(deps='nothing')
-
-    something_toolset = FunctionToolset[str]()
-
-    something_else_toolset = FunctionToolset[str]()
-
-    nothing_toolset = FunctionToolset[str]()
-
-    async def prepare_toolset(ctx: RunContext[str]) -> AbstractToolset[str]:
-        if ctx.deps == 'something':
-            return something_toolset
-        elif ctx.deps == 'something_else':
-            return something_else_toolset
-        else:
-            return nothing_toolset
-
-    dynamic_toolset: DynamicToolset[str] = DynamicToolset[str](build_toolset_fn=prepare_toolset)
-
-    # Enter the first context manager
-    async with dynamic_toolset:
-        assert dynamic_toolset.toolset is None
-
-        # The toolset is built dynamically on the first call to get_tools within the context
-        _ = await dynamic_toolset.get_tools(something_context)
-        assert dynamic_toolset.toolset == something_toolset
-        assert dynamic_toolset._toolset_stack.get() == [something_toolset]  # pyright: ignore[reportPrivateUsage]
-
-        # Enter the second context manager
-        async with dynamic_toolset:
-            # The toolset appears empty, and is built on the call to get_tools
-            assert dynamic_toolset.toolset is None
-            _ = await dynamic_toolset.get_tools(nothing_context)
-            assert dynamic_toolset.toolset == nothing_toolset
-            assert dynamic_toolset._toolset_stack.get() == [something_toolset, nothing_toolset]  # pyright: ignore[reportPrivateUsage]
-
-            # Enter the third context manager
-            async with dynamic_toolset:
-                # The toolset appears empty, and is built on the call to get_tools
-                assert dynamic_toolset.toolset is None
-                _ = await dynamic_toolset.get_tools(something_else_context)
-                assert dynamic_toolset.toolset == something_else_toolset
-                assert dynamic_toolset._toolset_stack.get() == [  # pyright: ignore[reportPrivateUsage]
-                    something_toolset,
-                    nothing_toolset,
-                    something_else_toolset,
-                ]
-
-            # Ensure the toolset reverts to the 2nd toolset
-            _ = await dynamic_toolset.get_tools(nothing_context)
-            assert dynamic_toolset.toolset == nothing_toolset
-            assert dynamic_toolset._toolset_stack.get() == [something_toolset, nothing_toolset]  # pyright: ignore[reportPrivateUsage]
-
-        # Ensure the toolset reverts to the 1st toolset
-        assert dynamic_toolset.toolset == something_toolset
-        assert dynamic_toolset._toolset_stack.get() == [something_toolset]  # pyright: ignore[reportPrivateUsage]
-
-    # Ensure the toolset is empty after exiting all context managers
-    async with dynamic_toolset:
-        assert dynamic_toolset.toolset is None
-        assert dynamic_toolset._toolset_stack.get() == [None]  # pyright: ignore[reportPrivateUsage]
-
-
-async def test_dynamic_toolset_call():
-    """Test that the dynamic toolset correctly handles nested context managers."""
-
-    something_context: RunContext[str] = build_run_context(deps='something')
-
-    def test_something(ctx: RunContext[str]) -> Literal['something']:
-        return 'something'
-
-    async def prepare_toolset(ctx: RunContext[str]) -> AbstractToolset[str]:
-        toolset = FunctionToolset[str]()
-        toolset.add_function(test_something)
-        return toolset
-
-    dynamic_toolset: DynamicToolset[str] = DynamicToolset[str](build_toolset_fn=prepare_toolset)
-
-    # Enter the first context manager
-    async with dynamic_toolset:
-        # The toolset is built dynamically on the first call to get_tools within the context
-        tools = await dynamic_toolset.get_tools(something_context)
-
-        first_tool = tools['test_something']
-        first_tool_result = await dynamic_toolset.call_tool(
-            name='test_something', tool_args={}, ctx=something_context, tool=first_tool
-        )
-        assert first_tool_result == 'something'