-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Toward SearchableToolSet and cross-model ToolSearch #3680
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
df4dea0
980187b
35a65e9
364a58e
8ffdf17
0f754c2
4d38f43
309e442
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -395,6 +395,33 @@ def unique_id(self) -> str: | |
| return ':'.join([self.kind, self.id]) | ||
|
|
||
|
|
||
| @dataclass(kw_only=True) | ||
| class ToolSearchTool(AbstractBuiltinTool): | ||
| """A builtin tool that searches for tools during dynamic tool discovery. | ||
|
|
||
| To defer loading a tool's definition until the model finds it, mark it as `defer_loading=True`. | ||
|
|
||
| Note that only models with `ModelProfile.supports_tool_search` use this builtin tool. These models receive all tool | ||
| definitions and natively implement search and loading. All other models rely on `SearchableToolset` instead. | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'd rather explain this in the future Tools docs section on |
||
|
|
||
| Supported by: | ||
|
|
||
| * Anthropic | ||
|
|
||
| """ | ||
|
|
||
| search_type: Literal['regex', 'bm25'] | None = None | ||
| """Custom search type to use for tool discovery. Currently only supported by Anthropic models. | ||
|
|
||
| See https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/tool-search-tool for more info. | ||
|
|
||
| - `'regex'`: Constructs Python `re.search()` patterns. Max 200 characters per query. Case-sensitive by default. | ||
| - `'bm25'`: Uses natural language queries with semantic matching across tool metadata. | ||
| """ | ||
|
|
||
| kind: str = 'tool_search' | ||
|
|
||
|
|
||
| def _tool_discriminator(tool_data: dict[str, Any] | AbstractBuiltinTool) -> str: | ||
| if isinstance(tool_data, dict): | ||
| return tool_data.get('kind', AbstractBuiltinTool.kind) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,6 @@ | ||
| from __future__ import annotations as _annotations | ||
|
|
||
| import logging | ||
| import io | ||
| from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator | ||
| from contextlib import asynccontextmanager | ||
|
|
@@ -13,7 +14,7 @@ | |
| from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage | ||
| from .._run_context import RunContext | ||
| from .._utils import guard_tool_call_id as _guard_tool_call_id | ||
| from ..builtin_tools import CodeExecutionTool, MCPServerTool, MemoryTool, WebFetchTool, WebSearchTool | ||
| from ..builtin_tools import CodeExecutionTool, MCPServerTool, MemoryTool, ToolSearchTool, WebFetchTool, WebSearchTool | ||
| from ..exceptions import ModelAPIError, UserError | ||
| from ..messages import ( | ||
| BinaryContent, | ||
|
|
@@ -114,6 +115,9 @@ | |
| BetaToolParam, | ||
| BetaToolResultBlockParam, | ||
| BetaToolUnionParam, | ||
| BetaToolSearchToolBm25_20251119Param, | ||
| BetaToolSearchToolRegex20251119Param, | ||
| BetaToolSearchToolResultBlock, | ||
| BetaToolUseBlock, | ||
| BetaToolUseBlockParam, | ||
| BetaWebFetchTool20250910Param, | ||
|
|
@@ -125,6 +129,7 @@ | |
| BetaWebSearchToolResultBlockParam, | ||
| BetaWebSearchToolResultBlockParamContentParam, | ||
| ) | ||
| from anthropic.types.beta.beta_tool_search_tool_result_block import BetaToolSearchToolResultBlock | ||
| from anthropic.types.beta.beta_web_fetch_tool_result_block_param import ( | ||
| Content as WebFetchToolResultBlockParamContent, | ||
| ) | ||
|
|
@@ -511,6 +516,9 @@ def _process_response(self, response: BetaMessage) -> ModelResponse: | |
| elif isinstance(item, BetaMCPToolResultBlock): | ||
| call_part = builtin_tool_calls.get(item.tool_use_id) | ||
| items.append(_map_mcp_server_result_block(item, call_part, self.system)) | ||
| elif isinstance(item, BetaToolSearchToolResultBlock): | ||
| call_part = builtin_tool_calls.get(item.tool_use_id) | ||
| items.append(_map_mcp_server_result_block(item, call_part, self.system)) | ||
| else: | ||
| assert isinstance(item, BetaToolUseBlock), f'unexpected item type {type(item)}' | ||
| items.append( | ||
|
|
@@ -577,6 +585,8 @@ def _add_builtin_tools( | |
| ) -> tuple[list[BetaToolUnionParam], list[BetaRequestMCPServerURLDefinitionParam], set[str]]: | ||
| beta_features: set[str] = set() | ||
| mcp_servers: list[BetaRequestMCPServerURLDefinitionParam] = [] | ||
| tool_search_type: Literal['regex', 'bm25'] | None = None | ||
|
|
||
| for tool in model_request_parameters.builtin_tools: | ||
| if isinstance(tool, WebSearchTool): | ||
| user_location = UserLocation(type='approximate', **tool.user_location) if tool.user_location else None | ||
|
|
@@ -629,10 +639,32 @@ def _add_builtin_tools( | |
| mcp_server_url_definition_param['authorization_token'] = tool.authorization_token | ||
| mcp_servers.append(mcp_server_url_definition_param) | ||
| beta_features.add('mcp-client-2025-04-04') | ||
| elif isinstance(tool, ToolSearchTool): | ||
| tool_search_type = tool.search_type | ||
| else: # pragma: no cover | ||
| raise UserError( | ||
| f'`{tool.__class__.__name__}` is not supported by `AnthropicModel`. If it should be, please file an issue.' | ||
| ) | ||
|
|
||
| needs_tool_search = any(tool.get('defer_loading') for tool in tools) | ||
|
|
||
| if needs_tool_search: | ||
| beta_features.add('advanced-tool-use-2025-11-20') | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Note that different providers use different headers: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool |
||
| if tool_search_type == 'bm25': | ||
| tools.append( | ||
| BetaToolSearchToolBm25_20251119Param( | ||
| name='tool_search_tool_bm25', | ||
| type='tool_search_tool_bm25_20251119', | ||
| ) | ||
| ) | ||
| else: | ||
| tools.append( | ||
| BetaToolSearchToolRegex20251119Param( | ||
| name='tool_search_tool_regex', | ||
| type='tool_search_tool_regex_20251119', | ||
| ) | ||
| ) | ||
|
|
||
| return tools, mcp_servers, beta_features | ||
|
|
||
| def _infer_tool_choice( | ||
|
|
@@ -1062,6 +1094,8 @@ def _map_tool_definition(self, f: ToolDefinition) -> BetaToolParam: | |
| 'description': f.description or '', | ||
| 'input_schema': f.parameters_json_schema, | ||
| } | ||
| if f.defer_loading: | ||
| tool_param['defer_loading'] = True | ||
| if f.strict and self.profile.supports_json_schema_output: | ||
| tool_param['strict'] = f.strict | ||
| return tool_param | ||
|
|
@@ -1297,8 +1331,12 @@ def _map_server_tool_use_block(item: BetaServerToolUseBlock, provider_name: str) | |
| elif item.name in ('bash_code_execution', 'text_editor_code_execution'): # pragma: no cover | ||
| raise NotImplementedError(f'Anthropic built-in tool {item.name!r} is not currently supported.') | ||
| elif item.name in ('tool_search_tool_regex', 'tool_search_tool_bm25'): # pragma: no cover | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's move this up so that the |
||
| # NOTE this is being implemented in https://github.com/pydantic/pydantic-ai/pull/3550 | ||
| raise NotImplementedError(f'Anthropic built-in tool {item.name!r} is not currently supported.') | ||
| return BuiltinToolCallPart( | ||
| provider_name=provider_name, | ||
| tool_name=ToolSearchTool.kind, | ||
| args=cast(dict[str, Any], item.input) or None, | ||
| tool_call_id=item.id, | ||
| ) | ||
| else: | ||
| assert_never(item.name) | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -65,6 +65,9 @@ class ModelProfile: | |
| This is currently only used by `OpenAIChatModel`, `HuggingFaceModel`, and `GroqModel`. | ||
| """ | ||
|
|
||
| supports_tool_search: bool = False | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In #3456, we're getting |
||
| """Whether the model has native support for tool search (builtin ToolSearchTool) and defer loading tools.""" | ||
|
|
||
| @classmethod | ||
| def from_profile(cls, profile: ModelProfile | None) -> Self: | ||
| """Build a ModelProfile subclass instance from a ModelProfile instance.""" | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -23,6 +23,7 @@ def anthropic_model_profile(model_name: str) -> ModelProfile | None: | |
| thinking_tags=('<thinking>', '</thinking>'), | ||
| supports_json_schema_output=supports_json_schema_output, | ||
| json_schema_transformer=AnthropicJsonSchemaTransformer, | ||
| supports_tool_search=True, | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Note that this is not actually supported by all models: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool |
||
| ) | ||
|
|
||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -39,6 +39,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]): | |
| docstring_format: DocstringFormat | ||
| require_parameter_descriptions: bool | ||
| schema_generator: type[GenerateJsonSchema] | ||
| defer_loading: bool | ||
|
|
||
| def __init__( | ||
| self, | ||
|
|
@@ -53,6 +54,7 @@ def __init__( | |
| requires_approval: bool = False, | ||
| metadata: dict[str, Any] | None = None, | ||
| id: str | None = None, | ||
| defer_loading: bool = False, | ||
| ): | ||
| """Build a new function toolset. | ||
|
|
||
|
|
@@ -78,6 +80,7 @@ def __init__( | |
| Applies to all tools, unless overridden when adding a tool, which will be merged with the toolset's metadata. | ||
| id: An optional unique ID for the toolset. A toolset needs to have an ID in order to be used in a durable execution environment like Temporal, | ||
| in which case the ID will be used to identify the toolset's activities within the workflow. | ||
| defer_loading: If True, default to hiding each tool and only activate it when the model searches for tools. | ||
| """ | ||
| self.max_retries = max_retries | ||
| self._id = id | ||
|
|
@@ -88,6 +91,7 @@ def __init__( | |
| self.sequential = sequential | ||
| self.requires_approval = requires_approval | ||
| self.metadata = metadata | ||
| self.defer_loading = defer_loading | ||
|
|
||
| self.tools = {} | ||
| for tool in tools: | ||
|
|
@@ -137,6 +141,7 @@ def tool( | |
| sequential: bool | None = None, | ||
| requires_approval: bool | None = None, | ||
| metadata: dict[str, Any] | None = None, | ||
| defer_loading: bool | None = None, | ||
| ) -> Any: | ||
| """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument. | ||
|
|
||
|
|
@@ -193,6 +198,8 @@ async def spam(ctx: RunContext[str], y: float) -> float: | |
| If `None`, the default value is determined by the toolset. | ||
| metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization. | ||
| If `None`, the default value is determined by the toolset. If provided, it will be merged with the toolset's metadata. | ||
| defer_loading: If True, hide the tool by default and only activate it when the model searches for tools. | ||
| If `None`, the default value is determined by the toolset. | ||
| """ | ||
|
|
||
| def tool_decorator( | ||
|
|
@@ -213,6 +220,7 @@ def tool_decorator( | |
| sequential=sequential, | ||
| requires_approval=requires_approval, | ||
| metadata=metadata, | ||
| defer_loading=defer_loading, | ||
| ) | ||
| return func_ | ||
|
|
||
|
|
@@ -233,6 +241,7 @@ def add_function( | |
| sequential: bool | None = None, | ||
| requires_approval: bool | None = None, | ||
| metadata: dict[str, Any] | None = None, | ||
| defer_loading: bool | None = None, | ||
| ) -> None: | ||
| """Add a function as a tool to the toolset. | ||
|
|
||
|
|
@@ -267,6 +276,8 @@ def add_function( | |
| If `None`, the default value is determined by the toolset. | ||
| metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization. | ||
| If `None`, the default value is determined by the toolset. If provided, it will be merged with the toolset's metadata. | ||
| defer_loading: If True, hide the tool by default and only activate it when the model searches for tools. | ||
| If `None`, the default value is determined by the toolset. | ||
| """ | ||
| if docstring_format is None: | ||
| docstring_format = self.docstring_format | ||
|
|
@@ -280,6 +291,8 @@ def add_function( | |
| sequential = self.sequential | ||
| if requires_approval is None: | ||
| requires_approval = self.requires_approval | ||
| if defer_loading is None: | ||
| defer_loading = self.defer_loading | ||
|
|
||
| tool = Tool[AgentDepsT]( | ||
| func, | ||
|
|
@@ -295,6 +308,7 @@ def add_function( | |
| sequential=sequential, | ||
| requires_approval=requires_approval, | ||
| metadata=metadata, | ||
| defer_loading=defer_loading, | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We need to also add this field to the |
||
| ) | ||
| self.add_tool(tool) | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As mentioned in the penultimate paragraph in #3666 (comment), I envision that users will primarily interact with this feature via the
defer_loading=Truefield on tools (rather than throughToolSearchToolorSearchableToolset), which we'd detect automatically and handle by wrapping those tools inSearchableToolset. Then, as implemented, that would fall back on Anthropic's native functionality when appropriate.So I'd love to see a test where we register a few
@agent.tool(defer_loading=True)s and then test that it works as expected with both Anthropic and OpenAI.