From df4dea0b38f72eb3bb0407d613580598c4a4c9de Mon Sep 17 00:00:00 2001 From: Anton Tayanovskyy Date: Tue, 9 Dec 2025 08:02:43 -0500 Subject: [PATCH 1/6] Sketch SearchableToolSet --- .../pydantic_ai/toolsets/searchable.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 pydantic_ai_slim/pydantic_ai/toolsets/searchable.py diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py new file mode 100644 index 0000000000..e0882ceb8f --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py @@ -0,0 +1,67 @@ +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Any + +from .._run_context import AgentDepsT, RunContext +from .abstract import AbstractToolset, ToolsetTool + + +SEARCH_TOOL_NAME = "search_tool" + + +@dataclass +class _SearchTool(ToolsetTool[AgentDepsT]): + """A tool that searches for more relevant tools from a SearchableToolSet""" + pass + + +@dataclass +class SearchableToolset(AbstractToolset[AgentDepsT]): + """A toolset that implements tool search and deferred tool loading.""" + + toolset: AbstractToolset[AgentDepsT] + _active_tool_names: set[str] = {} + + @property + def id(self) -> str | None: + return None # pragma: no cover + + @property + def label(self) -> str: + return f'{self.__class__.__name__}({self.toolset.label})' # pragma: no cover + + async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]: + + all_tools: dict[str, ToolsetTool[AgentDepsT]] = {} + all_tools[SEARCH_TOOL_NAME] = _SearchTool( + toolset=self, + tool_def=None, + max_retries=1, + args_validator=None, + ) + + toolset_tools = await self.toolset.get_tools(ctx) + for tool in toolset_tools: + + # TODO proper error handling + assert tool.name != SEARCH_TOOL_NAME + + if tool.name in self._active_tool_names: + all_tools[tool.name] = tool + return all_tools + + async def call_tool( + self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT] + ) -> Any: + if isinstance(tool, _SearchTool): + raise Exception("TODO call search tool") + else: + return await self.toolset.call_tool(name, tool_args, ctx, tool) + + def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None: + self.toolset.apply(visitor) + + def visit_and_replace( + self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]] + ) -> AbstractToolset[AgentDepsT]: + return replace(self, toolset=self.toolset.visit_and_replace(visitor)) From 980187b0b8c8b8f7a8bdcd479c91ecfd97c7f580 Mon Sep 17 00:00:00 2001 From: Anton Tayanovskyy Date: Thu, 11 Dec 2025 06:35:32 -0500 Subject: [PATCH 2/6] WIP on searchable tool --- .../pydantic_ai/toolsets/searchable.py | 75 +++++++++++++++---- 1 file changed, 62 insertions(+), 13 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py index e0882ceb8f..a7bff22886 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py @@ -1,18 +1,49 @@ +import re from collections.abc import Callable -from dataclasses import dataclass, replace -from typing import Any +from dataclasses import dataclass, field, replace +from typing import Any, TypedDict + +from pydantic import TypeAdapter +from typing_extensions import Self from .._run_context import AgentDepsT, RunContext -from .abstract import AbstractToolset, ToolsetTool +from ..tools import ToolDefinition +from .abstract import AbstractToolset, SchemaValidatorProt, ToolsetTool + +_SEARCH_TOOL_NAME = 'search_tool' + + +class _SearchToolArgs(TypedDict): + regex: str -SEARCH_TOOL_NAME = "search_tool" +def _search_tool_def() -> ToolDefinition: + return ToolDefinition( + name=_SEARCH_TOOL_NAME, + description='Search for additional tools', + parameters_json_schema={ + 'type': 'object', + 'properties': { + 'regex': { + 'type': 'string', + 'description': 'Regex pattern to search for relevant tools', + } + }, + 'required': ['regex'], + }, + ) + + +def _search_tool_validator() -> SchemaValidatorProt: + return TypeAdapter(_SearchToolArgs).validator @dataclass class _SearchTool(ToolsetTool[AgentDepsT]): """A tool that searches for more relevant tools from a SearchableToolSet""" - pass + + tool_def: ToolDefinition = field(default_factory=_search_tool_def) + args_validator: SchemaValidatorProt = field(default_factory=_search_tool_validator) @dataclass @@ -31,20 +62,16 @@ def label(self) -> str: return f'{self.__class__.__name__}({self.toolset.label})' # pragma: no cover async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]: - all_tools: dict[str, ToolsetTool[AgentDepsT]] = {} - all_tools[SEARCH_TOOL_NAME] = _SearchTool( + all_tools[_SEARCH_TOOL_NAME] = _SearchTool( toolset=self, - tool_def=None, max_retries=1, - args_validator=None, ) toolset_tools = await self.toolset.get_tools(ctx) - for tool in toolset_tools: - + for tool_name, tool in toolset_tools.items(): # TODO proper error handling - assert tool.name != SEARCH_TOOL_NAME + assert tool_name != _SEARCH_TOOL_NAME if tool.name in self._active_tool_names: all_tools[tool.name] = tool @@ -54,10 +81,25 @@ async def call_tool( self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT] ) -> Any: if isinstance(tool, _SearchTool): - raise Exception("TODO call search tool") + adapter = TypeAdapter(_SearchToolArgs) + typed_args = adapter.validate_python(tool_args) + return await self.call_search_tool(tool, typed_args) else: return await self.toolset.call_tool(name, tool_args, ctx, tool) + async def call_search_tool(self, args: _SearchToolArgs, ctx: RunContext[AgentDepsT]) -> list[str]: + """Searches for tools matching the query, activates them and returns their names.""" + toolset_tools = await self.toolset.get_tools(ctx) + matching_tool_names: list[str] = [] + + 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): + matching_tool_names.append(tool.tool_def.name) + + self._active_tool_names.update(matching_tool_names) + return matching_tool_names + def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None: self.toolset.apply(visitor) @@ -65,3 +107,10 @@ def visit_and_replace( self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]] ) -> AbstractToolset[AgentDepsT]: return replace(self, toolset=self.toolset.visit_and_replace(visitor)) + + async def __aenter__(self) -> Self: + await self.wrapped.__aenter__() + return self + + async def __aexit__(self, *args: Any) -> bool | None: + return await self.wrapped.__aexit__(*args) From 35a65e9b18af41a4525e5c12313ead9beb02e510 Mon Sep 17 00:00:00 2001 From: Anton Tayanovskyy Date: Thu, 11 Dec 2025 06:36:15 -0500 Subject: [PATCH 3/6] Format with unsafe fixes --- pydantic_ai_slim/pydantic_ai/toolsets/searchable.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py index a7bff22886..93587d5282 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py @@ -40,7 +40,7 @@ def _search_tool_validator() -> SchemaValidatorProt: @dataclass class _SearchTool(ToolsetTool[AgentDepsT]): - """A tool that searches for more relevant tools from a SearchableToolSet""" + """A tool that searches for more relevant tools from a SearchableToolSet.""" tool_def: ToolDefinition = field(default_factory=_search_tool_def) args_validator: SchemaValidatorProt = field(default_factory=_search_tool_validator) From 364a58e8591a4a771c3738f7c652103b86bc8ed6 Mon Sep 17 00:00:00 2001 From: Anton Tayanovskyy Date: Thu, 11 Dec 2025 06:54:53 -0500 Subject: [PATCH 4/6] Simple fixes --- pydantic_ai_slim/pydantic_ai/toolsets/searchable.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py index 93587d5282..881381b60f 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py @@ -51,7 +51,7 @@ class SearchableToolset(AbstractToolset[AgentDepsT]): """A toolset that implements tool search and deferred tool loading.""" toolset: AbstractToolset[AgentDepsT] - _active_tool_names: set[str] = {} + _active_tool_names: set[str] = field(default_factory=set) @property def id(self) -> str | None: @@ -73,8 +73,8 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[ # TODO proper error handling assert tool_name != _SEARCH_TOOL_NAME - if tool.name in self._active_tool_names: - all_tools[tool.name] = tool + if tool_name in self._active_tool_names: + all_tools[tool_name] = tool return all_tools async def call_tool( @@ -83,7 +83,7 @@ async def call_tool( if isinstance(tool, _SearchTool): adapter = TypeAdapter(_SearchToolArgs) typed_args = adapter.validate_python(tool_args) - return await self.call_search_tool(tool, typed_args) + return await self.call_search_tool(typed_args, ctx) else: return await self.toolset.call_tool(name, tool_args, ctx, tool) @@ -109,8 +109,8 @@ def visit_and_replace( return replace(self, toolset=self.toolset.visit_and_replace(visitor)) async def __aenter__(self) -> Self: - await self.wrapped.__aenter__() + await self.toolset.__aenter__() return self async def __aexit__(self, *args: Any) -> bool | None: - return await self.wrapped.__aexit__(*args) + return await self.toolset.__aexit__(*args) From 8ffdf179c679b37e30559d5f993700db1683cf22 Mon Sep 17 00:00:00 2001 From: Anton Tayanovskyy Date: Thu, 11 Dec 2025 07:22:00 -0500 Subject: [PATCH 5/6] Debug search tools and iterate on the description prompt --- .../pydantic_ai/toolsets/searchable.py | 20 ++- test_searchable_example.py | 136 ++++++++++++++++++ 2 files changed, 152 insertions(+), 4 deletions(-) create mode 100644 test_searchable_example.py diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py index 881381b60f..7702104688 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/searchable.py @@ -1,3 +1,4 @@ +import logging import re from collections.abc import Callable from dataclasses import dataclass, field, replace @@ -10,7 +11,7 @@ from ..tools import ToolDefinition from .abstract import AbstractToolset, SchemaValidatorProt, ToolsetTool -_SEARCH_TOOL_NAME = 'search_tool' +_SEARCH_TOOL_NAME = 'load_tools' class _SearchToolArgs(TypedDict): @@ -20,7 +21,11 @@ class _SearchToolArgs(TypedDict): def _search_tool_def() -> ToolDefinition: return ToolDefinition( name=_SEARCH_TOOL_NAME, - description='Search for additional tools', + 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. +""", parameters_json_schema={ 'type': 'object', 'properties': { @@ -62,6 +67,7 @@ def label(self) -> str: return f'{self.__class__.__name__}({self.toolset.label})' # pragma: no cover async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]: + logging.debug("SearchableToolset.get_tools") all_tools: dict[str, ToolsetTool[AgentDepsT]] = {} all_tools[_SEARCH_TOOL_NAME] = _SearchTool( toolset=self, @@ -75,6 +81,8 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[ if tool_name in self._active_tool_names: all_tools[tool_name] = tool + + logging.debug(f"SearchableToolset.get_tools ==> {[t for t in all_tools]}") return all_tools async def call_tool( @@ -83,9 +91,13 @@ async def call_tool( if isinstance(tool, _SearchTool): adapter = TypeAdapter(_SearchToolArgs) typed_args = adapter.validate_python(tool_args) - return await self.call_search_tool(typed_args, ctx) + result = await self.call_search_tool(typed_args, ctx) + logging.debug(f"SearchableToolset.call_tool({name}, {tool_args}) ==> {result}") + return result else: - return await self.toolset.call_tool(name, tool_args, ctx, tool) + result = await self.toolset.call_tool(name, tool_args, ctx, tool) + logging.debug(f"SearchableToolset.call_tool({name}, {tool_args}) ==> {result}") + return result async def call_search_tool(self, args: _SearchToolArgs, ctx: RunContext[AgentDepsT]) -> list[str]: """Searches for tools matching the query, activates them and returns their names.""" diff --git a/test_searchable_example.py b/test_searchable_example.py new file mode 100644 index 0000000000..2f9a0fe96c --- /dev/null +++ b/test_searchable_example.py @@ -0,0 +1,136 @@ +"""Minimal example to test SearchableToolset functionality. + +Run with: uv run python test_searchable_example.py +Make sure you have ANTHROPIC_API_KEY set in your environment. +""" + +import asyncio +import logging +import sys + +# Configure logging to print to stdout +logging.basicConfig( + level=logging.DEBUG, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + stream=sys.stdout +) + +# Silence noisy loggers +logging.getLogger('asyncio').setLevel(logging.WARNING) +logging.getLogger('httpx').setLevel(logging.WARNING) +logging.getLogger('httpcore.connection').setLevel(logging.WARNING) +logging.getLogger('httpcore.http11').setLevel(logging.WARNING) +logging.getLogger('anthropic._base_client').setLevel(logging.WARNING) + +from pydantic_ai import Agent +from pydantic_ai.toolsets import FunctionToolset +from pydantic_ai.toolsets.searchable import SearchableToolset + + +# Create a toolset with various tools +toolset = FunctionToolset() + + +@toolset.tool +def get_weather(city: str) -> str: + """Get the current weather for a given city. + + Args: + city: The name of the city to get weather for. + """ + return f"The weather in {city} is sunny and 72°F" + + +@toolset.tool +def calculate_sum(a: float, b: float) -> float: + """Add two numbers together. + + Args: + a: The first number. + b: The second number. + """ + return a + b + + +@toolset.tool +def calculate_product(a: float, b: float) -> float: + """Multiply two numbers together. + + Args: + a: The first number. + b: The second number. + """ + return a * b + + +@toolset.tool +def fetch_user_data(user_id: int) -> dict: + """Fetch user data from the database. + + Args: + user_id: The ID of the user to fetch. + """ + return {"id": user_id, "name": "John Doe", "email": "john@example.com"} + + +@toolset.tool +def send_email(recipient: str, subject: str, body: str) -> str: + """Send an email to a recipient. + + Args: + recipient: The email address of the recipient. + subject: The subject line of the email. + body: The body content of the email. + """ + return f"Email sent to {recipient} with subject '{subject}'" + + +@toolset.tool +def list_database_tables() -> list[str]: + """List all tables in the database.""" + return ["users", "orders", "products", "reviews"] + + +# Wrap the toolset with SearchableToolset +searchable_toolset = SearchableToolset(toolset=toolset) + +# Create an agent with the searchable toolset +agent = Agent( + 'anthropic:claude-sonnet-4-5', + toolsets=[searchable_toolset], + system_prompt=( + "You are a helpful assistant." + ), +) + + +async def main(): + print("=" * 60) + print("Testing SearchableToolset") + print("=" * 60) + print() + + # Test 1: Ask something that requires searching for calculation tools + print("Test 1: Calculation task") + print("-" * 60) + result = await agent.run("What is 123 multiplied by 456?") + print(f"Result: {result.output}") + print() + + # Test 2: Ask something that requires searching for database tools + print("\nTest 2: Database task") + print("-" * 60) + result = await agent.run("Can you list the database tables and then fetch user 42?") + print(f"Result: {result.output}") + print() + + # Test 3: Ask something that requires weather tool + print("\nTest 3: Weather task") + print("-" * 60) + result = await agent.run("What's the weather like in San Francisco?") + print(f"Result: {result.output}") + print() + + +if __name__ == "__main__": + asyncio.run(main()) From 0f754c2cca268f866bfc3d059ef44187a8333a05 Mon Sep 17 00:00:00 2001 From: Anton Tayanovskyy Date: Sat, 13 Dec 2025 10:03:03 -0500 Subject: [PATCH 6/6] Add supports_tool_search to ModelProfile --- pydantic_ai_slim/pydantic_ai/profiles/__init__.py | 3 +++ pydantic_ai_slim/pydantic_ai/profiles/anthropic.py | 1 + 2 files changed, 4 insertions(+) diff --git a/pydantic_ai_slim/pydantic_ai/profiles/__init__.py b/pydantic_ai_slim/pydantic_ai/profiles/__init__.py index 84a1c04012..8c0d773832 100644 --- a/pydantic_ai_slim/pydantic_ai/profiles/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/profiles/__init__.py @@ -65,6 +65,9 @@ class ModelProfile: This is currently only used by `OpenAIChatModel`, `HuggingFaceModel`, and `GroqModel`. """ + supports_tool_search: bool = False + """Whether the model has native support for tool search and defer loading tools.""" + @classmethod def from_profile(cls, profile: ModelProfile | None) -> Self: """Build a ModelProfile subclass instance from a ModelProfile instance.""" diff --git a/pydantic_ai_slim/pydantic_ai/profiles/anthropic.py b/pydantic_ai_slim/pydantic_ai/profiles/anthropic.py index 6a59ab2dec..bc76b4d5a9 100644 --- a/pydantic_ai_slim/pydantic_ai/profiles/anthropic.py +++ b/pydantic_ai_slim/pydantic_ai/profiles/anthropic.py @@ -23,6 +23,7 @@ def anthropic_model_profile(model_name: str) -> ModelProfile | None: thinking_tags=('', ''), supports_json_schema_output=supports_json_schema_output, json_schema_transformer=AnthropicJsonSchemaTransformer, + supports_tool_search=True, )