add dynamic toolset example

Author: strawgate

pydantic_ai_slim/pydantic_ai/toolsets/dynamic.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from abc import ABC
+from collections.abc import Awaitable, Callable
+from contextvars import ContextVar, Token
+from typing import TYPE_CHECKING, Any, Self
+
+from pydantic_ai.tools import AgentDepsT, RunContext
+from pydantic_ai.toolsets import AbstractToolset
+from pydantic_ai.toolsets.abstract import ToolsetTool
+from pydantic_ai.toolsets.combined import CombinedToolset
+
+if TYPE_CHECKING:
+    pass
+
+
+BuildToolsetFunc = Callable[[RunContext[AgentDepsT]], Awaitable[AbstractToolset[AgentDepsT]]]
+
+
+class DynamicToolset(AbstractToolset[AgentDepsT], ABC):
+    """A Toolset that is dynamically built during an Agent run based on the first available run context."""
+
+    _build_toolset_fn: BuildToolsetFunc[AgentDepsT]
+
+    _dynamic_toolset: ContextVar[CombinedToolset[AgentDepsT]] = ContextVar('_toolset', default=CombinedToolset[AgentDepsT](toolsets=[]))
+    _token: Token[CombinedToolset[AgentDepsT]] | None = None
+    #_toolset_deps: ContextVar[AgentDepsT | None] = ContextVar('_toolset_deps', default=None)
+
+    def __init__(self, build_toolset_fn: BuildToolsetFunc[AgentDepsT]):
+        self._build_toolset_fn = build_toolset_fn
+
+    async def __aenter__(self) -> Self:
+        # Store the current toolset in a token, so that it can be reset when the context is exited
+        self._token = self._dynamic_toolset.set(CombinedToolset[AgentDepsT](toolsets=[]))
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        # Reset the toolset to the previous toolset, so that it can be used again
+        if self._token:
+            self._dynamic_toolset.reset(self._token)
+        self._token = None
+        return None
+
+    @property
+    def _toolset(self) -> CombinedToolset[AgentDepsT]:
+        if not (toolset := self._dynamic_toolset.get()):
+            msg = 'Toolset not initialized. Use the `async with` context manager to initialize the toolset.'
+            raise RuntimeError(msg)
+
+        return toolset
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        if len(self._toolset.toolsets) == 0 or ctx.run_step == 0:
+            toolset = await self._build_toolset_fn(ctx)
+            self._toolset.toolsets = [toolset]
+
+        return await self._toolset.get_tools(ctx=ctx)
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        return await self._toolset.call_tool(name=name, tool_args=tool_args, ctx=ctx, tool=tool)

tests/test_toolsets.py

@@ -2,8 +2,10 @@
 
 import re
 from dataclasses import dataclass, replace
-from typing import TypeVar
+from pathlib import Path
+from typing import Literal, TypeVar
 
+from pydantic_ai.agent import Agent
 import pytest
 from inline_snapshot import snapshot
 
@@ -13,7 +15,9 @@
 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
@@ -469,3 +473,71 @@ async def test_context_manager():
         async with toolset:
             assert server1.is_running
             assert server2.is_running
+
+
+async def test_dynamic_toolset():
+    run_context = build_run_context(Path())
+
+    def test_function(ctx: RunContext[Path]) -> Literal['nothing']:
+        return 'nothing'
+
+    function_toolset = FunctionToolset[Path]()
+    function_toolset.add_function(test_function)
+
+    async def prepare_toolset(ctx: RunContext[Path]) -> AbstractToolset[Path]:
+        return function_toolset
+
+    dynamic_toolset: DynamicToolset[Path] = DynamicToolset[Path](build_toolset_fn=prepare_toolset)
+
+    # The toolset is unique per context manager
+    async with dynamic_toolset:
+
+        # The toolset starts empty
+        assert dynamic_toolset._dynamic_toolset.get().toolsets == []
+
+        # The toolset is built dynamically on the first call to get_tools within the context
+        _ = await dynamic_toolset.get_tools(run_context)
+        assert dynamic_toolset._dynamic_toolset.get().toolsets == [function_toolset]
+
+        # Any time the context is entered again, the toolsets are reset, to be generated again
+        async with dynamic_toolset:
+            assert dynamic_toolset._dynamic_toolset.get().toolsets == []
+
+        assert dynamic_toolset._dynamic_toolset.get().toolsets == [function_toolset]
+
+    async with dynamic_toolset:
+        assert dynamic_toolset._dynamic_toolset.get().toolsets == []
+
+async def test_dynamic_toolset_with_agent():
+    run_context = build_run_context(Path())
+
+    def test_function(ctx: RunContext[Path]) -> Literal['nothing']:
+        return 'nothing'
+
+    
+    def test_function_two(ctx: RunContext[Path]) -> Literal['nothing']:
+        return 'nothing'
+
+
+    function_toolset = FunctionToolset[Path]()
+    function_toolset.add_function(test_function)
+    function_toolset.add_function(test_function_two)
+
+    async def prepare_toolset(ctx: RunContext[Path]) -> AbstractToolset[Path]:
+        return function_toolset
+
+    dynamic_toolset: DynamicToolset[Path] = DynamicToolset[Path](build_toolset_fn=prepare_toolset)
+
+    agent = Agent[Path, str](
+        model=TestModel(),
+        toolsets=[dynamic_toolset],
+        deps_type=Path,
+        output_type=str,
+    )
+
+    async with agent:
+        result = await agent.run(deps=Path("."), user_prompt="Please call each tool you have access to and tell me what it returns")
+        print(result.output)
+
+        result = await agent.run(deps=Path("./tomato"), user_prompt="Please call each tool you have access to and tell me what it returns.")
+        print(result.output)