Toward SearchableToolSet and cross-model ToolSearch #3680

t0yv0 · 2025-12-09T13:03:41Z

Closes Support deferred loading of tools and discovery via tool search tool #3590
Related to feat: Add ToolSearchTool and defer_loading for dynamic tool discovery #3620

DouweM · 2025-12-09T17:36:33Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+class SearchableToolset(AbstractToolset[AgentDepsT]):
+    """A toolset that implements tool search and deferred tool loading."""
+
+    toolset: AbstractToolset[AgentDepsT]


Have a look at WrapperToolset which already handles this + properly forwards __aexit__ and __aenter__!

t0yv0 · 2025-12-11T12:27:35Z

test_searchable_example.py

@@ -0,0 +1,136 @@
+"""Minimal example to test SearchableToolset functionality.


It looks like proper tests need to go into:

test_toolsets.py has space for unit tests

somewhere there are VCR cassettes that record an interaction with an LLM could be useful here

I just wanted to get something quick to iterate with an actual LLM. This ended up working on Claude but took a few iterations on the prompt. The model seemed sensitive to how the "search tool" is called and the content of the description - it would either refuse to load it or start asking for user confirmation before loading it. It took some tweaking to get the current description to pass this simple test.

❯ uv run python test_searchable_example.py ============================================================ Testing SearchableToolset ============================================================ Test 1: Calculation task ------------------------------------------------------------ 2025-12-11 07:20:48,189 - root - DEBUG - SearchableToolset.get_tools 2025-12-11 07:20:48,189 - root - DEBUG - SearchableToolset.get_tools ==> ['load_tools'] Result: I can calculate that for you directly. 123 multiplied by 456 equals **56,088**. Test 2: Database task ------------------------------------------------------------ 2025-12-11 07:20:50,983 - root - DEBUG - SearchableToolset.get_tools 2025-12-11 07:20:50,984 - root - DEBUG - SearchableToolset.get_tools ==> ['load_tools'] 2025-12-11 07:20:54,254 - root - DEBUG - SearchableToolset.call_tool(load_tools, {'regex': 'database|sql|table|query'}) ==> ['fetch_user_data', 'list_database_tables'] 2025-12-11 07:20:54,255 - root - DEBUG - SearchableToolset.get_tools 2025-12-11 07:20:54,255 - root - DEBUG - SearchableToolset.get_tools ==> ['load_tools', 'fetch_user_data', 'list_database_tables'] 2025-12-11 07:20:57,735 - root - DEBUG - SearchableToolset.call_tool(list_database_tables, {}) ==> ['users', 'orders', 'products', 'reviews'] 2025-12-11 07:20:57,735 - root - DEBUG - SearchableToolset.call_tool(fetch_user_data, {'user_id': 42}) ==> {'id': 42, 'name': 'John Doe', 'email': 'john@example.com'} 2025-12-11 07:20:57,735 - root - DEBUG - SearchableToolset.get_tools 2025-12-11 07:20:57,736 - root - DEBUG - SearchableToolset.get_tools ==> ['load_tools', 'fetch_user_data', 'list_database_tables'] Result: Perfect! Here are the results: **Database Tables:** - users - orders - products - reviews **User 42 Data:** - ID: 42 - Name: John Doe - Email: john@example.com Test 3: Weather task ------------------------------------------------------------ 2025-12-11 07:21:00,605 - root - DEBUG - SearchableToolset.get_tools 2025-12-11 07:21:00,607 - root - DEBUG - SearchableToolset.get_tools ==> ['load_tools', 'fetch_user_data', 'list_database_tables'] 2025-12-11 07:21:04,597 - root - DEBUG - SearchableToolset.call_tool(load_tools, {'regex': 'weather'}) ==> ['get_weather'] 2025-12-11 07:21:04,598 - root - DEBUG - SearchableToolset.get_tools 2025-12-11 07:21:04,599 - root - DEBUG - SearchableToolset.get_tools ==> ['load_tools', 'get_weather', 'fetch_user_data', 'list_database_tables'] 2025-12-11 07:21:07,769 - root - DEBUG - SearchableToolset.call_tool(get_weather, {'city': 'San Francisco'}) ==> The weather in San Francisco is sunny and 72°F 2025-12-11 07:21:07,770 - root - DEBUG - SearchableToolset.get_tools 2025-12-11 07:21:07,771 - root - DEBUG - SearchableToolset.get_tools ==> ['load_tools', 'get_weather', 'fetch_user_data', 'list_database_tables'] Result: The weather in San Francisco is currently sunny and 72°F - a beautiful day!

t0yv0 · 2025-12-11T12:29:33Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+from ..tools import ToolDefinition
+from .abstract import AbstractToolset, SchemaValidatorProt, ToolsetTool
+
+_SEARCH_TOOL_NAME = 'load_tools'


Another curious bit is that when tool was called "more_tools", I hit a crash:

Traceback (most recent call last): File "/Users/anton/code/pydantic-ai/test_searchable_example.py", line 136, in <module> asyncio.run(main()) File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/runners.py", line 195, in run return runner.run(main) ^^^^^^^^^^^^^^^^ File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/runners.py", line 118, in run return self._loop.run_until_complete(task) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete return future.result() ^^^^^^^^^^^^^^^ File "/Users/anton/code/pydantic-ai/test_searchable_example.py", line 123, in main result = await agent.run("Can you list the database tables and then fetch user 42?") ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/agent/abstract.py", line 226, in run async with self.iter( ^^^^^^^^^^ File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/contextlib.py", line 231, in __aexit__ await self.gen.athrow(value) File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/agent/__init__.py", line 658, in iter async with graph.iter( ^^^^^^^^^^^ File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/contextlib.py", line 231, in __aexit__ await self.gen.athrow(value) File "/Users/anton/code/pydantic-ai/pydantic_graph/pydantic_graph/beta/graph.py", line 270, in iter async with GraphRun[StateT, DepsT, OutputT]( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/anton/code/pydantic-ai/pydantic_graph/pydantic_graph/beta/graph.py", line 423, in __aexit__ await self._async_exit_stack.__aexit__(exc_type, exc_val, exc_tb) File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/contextlib.py", line 754, in __aexit__ raise exc_details[1] File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/contextlib.py", line 735, in __aexit__ cb_suppress = cb(*exc_details) ^^^^^^^^^^^^^^^^ File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/contextlib.py", line 158, in __exit__ self.gen.throw(value) File "/Users/anton/code/pydantic-ai/pydantic_graph/pydantic_graph/beta/graph.py", line 978, in _unwrap_exception_groups raise exception File "/Users/anton/code/pydantic-ai/pydantic_graph/pydantic_graph/beta/graph.py", line 750, in _run_tracked_task result = await self._run_task(t_) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/anton/code/pydantic-ai/pydantic_graph/pydantic_graph/beta/graph.py", line 779, in _run_task output = await node.call(step_context) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/anton/code/pydantic-ai/pydantic_graph/pydantic_graph/beta/step.py", line 253, in _call_node return await node.run(GraphRunContext(state=ctx.state, deps=ctx.deps)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py", line 576, in run async with self.stream(ctx): ^^^^^^^^^^^^^^^^ File "/Users/anton/.local/share/uv/python/cpython-3.12.11-macos-aarch64-none/lib/python3.12/contextlib.py", line 217, in __aexit__ await anext(self.gen) File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py", line 590, in stream async for _event in stream: File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py", line 716, in _run_stream async for event in self._events_iterator: File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py", line 677, in _run_stream async for event in self._handle_tool_calls(ctx, tool_calls): File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py", line 732, in _handle_tool_calls async for event in process_tool_calls( File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py", line 925, in process_tool_calls ctx.state.increment_retries(ctx.deps.max_result_retries, model_settings=ctx.deps.model_settings) File "/Users/anton/code/pydantic-ai/pydantic_ai_slim/pydantic_ai/_agent_graph.py", line 127, in increment_retries raise exceptions.UnexpectedModelBehavior(message) pydantic_ai.exceptions.UnexpectedModelBehavior: Exceeded maximum retries (1) for output validation

Interesting, that suggests that the model was not calling it correctly (wrong args possibly). I suggest adding https://ai.pydantic.dev/logfire/ so you can easily see what's happening behind the scenes in an agent run.

DouweM · 2025-12-11T23:06:14Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+from ..tools import ToolDefinition
+from .abstract import AbstractToolset, SchemaValidatorProt, ToolsetTool
+
+_SEARCH_TOOL_NAME = 'load_tools'


Interesting, that suggests that the model was not calling it correctly (wrong args possibly). I suggest adding https://ai.pydantic.dev/logfire/ so you can easily see what's happening behind the scenes in an agent run.

DouweM · 2025-12-11T23:08:25Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+    regex: str
+
+
+def _search_tool_def() -> ToolDefinition:


Check out Tool.from_schema and the Tool constructor that takes a function (as used by FunctionToolset) for easier ways to construct a single tool. The function approach is the easiest by far

DouweM · 2025-12-11T23:10:37Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+        description="""Search and load additional tools to make them available to the agent.
+
+DO call this to find and load more tools needed for a task.
+NEVER ask the user if you should try loading tools, just try.


Hmm, I see you explained below that this was needed to pass the tests, even for Sonnet 4.5, but tokens are expensive so it'll be worth another iteration on this.

DouweM · 2025-12-11T23:11:51Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+        parameters_json_schema={
+            'type': 'object',
+            'properties': {
+                'regex': {


I like pattern slightly better as an argument name, as we may at some point support different ones. Although it is very helpful to the model in knowing what to put here, in case we remove/shorted the description.

DouweM · 2025-12-11T23:13:58Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+        all_tools: dict[str, ToolsetTool[AgentDepsT]] = {}
+        all_tools[_SEARCH_TOOL_NAME] = _SearchTool(
+            toolset=self,
+            max_retries=1,


We may want to increase this, to give the model a few chances to fix its regex, if it submitted an invalid one the first time

DouweM · 2025-12-11T23:15:20Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+    ) -> Any:
+        if isinstance(tool, _SearchTool):
+            adapter = TypeAdapter(_SearchToolArgs)
+            typed_args = adapter.validate_python(tool_args)


Arguments will/should already have been validated by this point when used through ToolManager/Agent!

DouweM · 2025-12-11T23:16:17Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+        matching_tool_names: list[str] = []
+
+        for tool_name, tool in toolset_tools.items():
+            rx = re.compile(args['regex'])


This'll be more efficient one line up :)

DouweM · 2025-12-11T23:16:40Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+
+        for tool_name, tool in toolset_tools.items():
+            rx = re.compile(args['regex'])
+            if rx.search(tool.tool_def.name) or rx.search(tool.tool_def.description):


For error handling, check out the ModelRetry exception

DouweM · 2025-12-11T23:19:11Z

pydantic_ai_slim/pydantic_ai/toolsets/searchable.py

+    """A toolset that implements tool search and deferred tool loading."""
+
+    toolset: AbstractToolset[AgentDepsT]
+    _active_tool_names: set[str] = field(default_factory=set)


The fact that this has instance variables means that it can't be reused across multiple agent runs, even though the same instance is registered to an agent just once... We had a similar issue with DynamicToolset, I suggest having a look at how we handle it there

Sketch SearchableToolSet

df4dea0

DouweM self-assigned this Dec 9, 2025

DouweM reviewed Dec 9, 2025

View reviewed changes

t0yv0 added 4 commits December 11, 2025 06:35

WIP on searchable tool

980187b

Format with unsafe fixes

35a65e9

Simple fixes

364a58e

Debug search tools and iterate on the description prompt

8ffdf17

t0yv0 commented Dec 11, 2025

View reviewed changes

DouweM requested changes Dec 11, 2025

View reviewed changes

DouweM added the awaiting author revision label Dec 11, 2025

		@@ -0,0 +1,136 @@
		"""Minimal example to test SearchableToolset functionality.

Toward SearchableToolSet and cross-model ToolSearch #3680

Are you sure you want to change the base?

Toward SearchableToolSet and cross-model ToolSearch #3680

Conversation

t0yv0 commented Dec 9, 2025 • edited by DouweM Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

t0yv0 commented Dec 9, 2025 •

edited by DouweM

Loading