diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index c8208ac9e6..6d9c49596f 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -1142,6 +1142,7 @@ def tool_plain( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, + defer_loading: bool = False, ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ... def tool_plain( @@ -1160,6 +1161,7 @@ def tool_plain( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, + defer_loading: bool = False, ) -> Any: """Decorator to register a tool function which DOES NOT take `RunContext` as an argument. @@ -1209,6 +1211,8 @@ async def spam(ctx: RunContext[str]) -> float: requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False. See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info. metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization. + defer_loading: Whether to defer loading this tool until discovered via tool search. Defaults to False. + See [`ToolDefinition.defer_loading`][pydantic_ai.tools.ToolDefinition.defer_loading] for more info. """ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]: @@ -1227,6 +1231,7 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams sequential=sequential, requires_approval=requires_approval, metadata=metadata, + defer_loading=defer_loading, ) return func_ diff --git a/pydantic_ai_slim/pydantic_ai/builtin_tools.py b/pydantic_ai_slim/pydantic_ai/builtin_tools.py index 1baaf09f3a..2df0144111 100644 --- a/pydantic_ai_slim/pydantic_ai/builtin_tools.py +++ b/pydantic_ai_slim/pydantic_ai/builtin_tools.py @@ -18,6 +18,7 @@ 'ImageGenerationTool', 'MemoryTool', 'MCPServerTool', + 'ToolSearchTool', ) _BUILTIN_TOOL_TYPES: dict[str, type[AbstractBuiltinTool]] = {} @@ -395,6 +396,32 @@ def unique_id(self) -> str: return ':'.join([self.kind, self.id]) +@dataclass(kw_only=True) +class ToolSearchTool(AbstractBuiltinTool): + """A builtin tool that enables dynamic tool discovery without loading all definitions upfront. + + Instead of consuming tokens on many tool definitions, the model searches for relevant tools + on-demand. Only matching tools get expanded into full definitions in the model's context. + + You should mark tools with `defer_loading=True` to make them discoverable on-demand while + keeping critical tools always loaded (with `defer_loading=False`). + + Supported by: + + * [Anthropic](https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/tool-search-tool) + """ + + search_type: Literal['regex', 'bm25'] | None = None + """The type of search to use for tool discovery. If None, uses the provider's default. + + - `'regex'`: Constructs regex patterns (for Anthropic: Python `re.search()`). Case-sensitive by default. + - `'bm25'`: Uses natural language queries with semantic matching across tool metadata. + """ + + kind: str = 'tool_search' + """The kind of tool.""" + + def _tool_discriminator(tool_data: dict[str, Any] | AbstractBuiltinTool) -> str: if isinstance(tool_data, dict): return tool_data.get('kind', AbstractBuiltinTool.kind) diff --git a/pydantic_ai_slim/pydantic_ai/models/anthropic.py b/pydantic_ai_slim/pydantic_ai/models/anthropic.py index ecdb9fe61f..0741098c9e 100644 --- a/pydantic_ai_slim/pydantic_ai/models/anthropic.py +++ b/pydantic_ai_slim/pydantic_ai/models/anthropic.py @@ -13,7 +13,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 +114,9 @@ BetaToolParam, BetaToolResultBlockParam, BetaToolUnionParam, + BetaToolSearchToolBm25_20251119Param, + BetaToolSearchToolRegex20251119Param, + BetaToolSearchToolResultBlock, BetaToolUseBlock, BetaToolUseBlockParam, BetaWebFetchTool20250910Param, @@ -574,6 +577,22 @@ def _add_builtin_tools( ) -> tuple[list[BetaToolUnionParam], list[BetaRequestMCPServerURLDefinitionParam], set[str]]: beta_features: set[str] = set() mcp_servers: list[BetaRequestMCPServerURLDefinitionParam] = [] + + # Check if any tools use defer_loading and auto-add ToolSearchTool if needed + has_defer_loading = any( + tool_def.defer_loading for tool_def in model_request_parameters.tool_defs.values() + ) + has_tool_search = any(isinstance(tool, ToolSearchTool) for tool in model_request_parameters.builtin_tools) + if has_defer_loading and not has_tool_search: + # Auto-add ToolSearchTool with default (regex) when defer_loading is used + tools.append( + BetaToolSearchToolRegex20251119Param( + name='tool_search_tool_regex', + type='tool_search_tool_regex_20251119', + ) + ) + beta_features.add('advanced-tool-use-2025-11-20') + 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 @@ -626,6 +645,23 @@ 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): # pragma: no branch + if tool.search_type == 'bm25': + tools.append( + BetaToolSearchToolBm25_20251119Param( + name='tool_search_tool_bm25', + type='tool_search_tool_bm25_20251119', + ) + ) + else: + # Default to regex if search_type is None or 'regex' + tools.append( + BetaToolSearchToolRegex20251119Param( + name='tool_search_tool_regex', + type='tool_search_tool_regex_20251119', + ) + ) + beta_features.add('advanced-tool-use-2025-11-20') else: # pragma: no cover raise UserError( f'`{tool.__class__.__name__}` is not supported by `AnthropicModel`. If it should be, please file an issue.' @@ -1041,6 +1077,8 @@ def _map_tool_definition(self, f: ToolDefinition) -> BetaToolParam: } if f.strict and self.profile.supports_json_schema_output: tool_param['strict'] = f.strict + if f.defer_loading: + tool_param['defer_loading'] = True return tool_param @staticmethod diff --git a/pydantic_ai_slim/pydantic_ai/tools.py b/pydantic_ai_slim/pydantic_ai/tools.py index e54b829bfb..d95c610756 100644 --- a/pydantic_ai_slim/pydantic_ai/tools.py +++ b/pydantic_ai_slim/pydantic_ai/tools.py @@ -262,6 +262,7 @@ class Tool(Generic[ToolAgentDepsT]): sequential: bool requires_approval: bool metadata: dict[str, Any] | None + defer_loading: bool function_schema: _function_schema.FunctionSchema """ The base JSON schema for the tool's parameters. @@ -285,6 +286,7 @@ def __init__( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, + defer_loading: bool = False, function_schema: _function_schema.FunctionSchema | None = None, ): """Create a new tool instance. @@ -341,6 +343,8 @@ async def prep_my_tool( requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False. See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info. metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization. + defer_loading: Whether to defer loading this tool until discovered via tool search. Defaults to False. + See [`ToolDefinition.defer_loading`][pydantic_ai.tools.ToolDefinition.defer_loading] for more info. function_schema: The function schema to use for the tool. If not provided, it will be generated. """ self.function = function @@ -362,6 +366,7 @@ async def prep_my_tool( self.sequential = sequential self.requires_approval = requires_approval self.metadata = metadata + self.defer_loading = defer_loading @classmethod def from_schema( @@ -418,6 +423,7 @@ def tool_def(self): sequential=self.sequential, metadata=self.metadata, kind='unapproved' if self.requires_approval else 'function', + defer_loading=self.defer_loading, ) async def prepare_tool_def(self, ctx: RunContext[ToolAgentDepsT]) -> ToolDefinition | None: @@ -503,6 +509,19 @@ class ToolDefinition: For MCP tools, this contains the `meta`, `annotations`, and `output_schema` fields from the tool definition. """ + defer_loading: bool = False + """Whether to defer loading this tool until it's discovered via tool search. + + When `True`, this tool will not be loaded into the model's context initially. Instead, the model + will use the Tool Search tool to discover it on-demand when needed, reducing token usage. + + Requires the [`ToolSearchTool`][pydantic_ai.builtin_tools.ToolSearchTool] builtin tool to be enabled. + + Supported by: + + * [Anthropic](https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/tool-search-tool) + """ + @property def defer(self) -> bool: """Whether calls to this tool will be deferred. diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/function.py b/pydantic_ai_slim/pydantic_ai/toolsets/function.py index e185ed0273..2916b4ed6d 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/function.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/function.py @@ -119,6 +119,7 @@ def tool( sequential: bool | None = None, requires_approval: bool | None = None, metadata: dict[str, Any] | None = None, + defer_loading: bool | None = None, ) -> Callable[[ToolFuncEither[AgentDepsT, ToolParams]], ToolFuncEither[AgentDepsT, ToolParams]]: ... def tool( @@ -137,6 +138,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 +195,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: Whether to defer loading this tool until discovered via tool search. Defaults to False. + See [`ToolDefinition.defer_loading`][pydantic_ai.tools.ToolDefinition.defer_loading] for more info. """ def tool_decorator( @@ -213,6 +217,7 @@ def tool_decorator( sequential=sequential, requires_approval=requires_approval, metadata=metadata, + defer_loading=defer_loading, ) return func_ @@ -233,6 +238,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 +273,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: Whether to defer loading this tool until discovered via tool search. Defaults to False. + See [`ToolDefinition.defer_loading`][pydantic_ai.tools.ToolDefinition.defer_loading] for more info. """ if docstring_format is None: docstring_format = self.docstring_format @@ -295,6 +303,7 @@ def add_function( sequential=sequential, requires_approval=requires_approval, metadata=metadata, + defer_loading=defer_loading or False, ) self.add_tool(tool) diff --git a/tests/models/anthropic/test_tool_search.py b/tests/models/anthropic/test_tool_search.py new file mode 100644 index 0000000000..83af904d6a --- /dev/null +++ b/tests/models/anthropic/test_tool_search.py @@ -0,0 +1,180 @@ +"""Tests for Anthropic ToolSearchTool and defer_loading features. + +These features enable dynamic tool discovery without loading all definitions upfront. +""" + +from __future__ import annotations + +from typing import Any, cast + +import pytest + +from pydantic_ai import Agent, Tool +from pydantic_ai.tools import ToolDefinition + +from ...conftest import try_import +from ..test_anthropic import MockAnthropic, completion_message + +with try_import() as imports_successful: + from anthropic.types.beta import BetaTextBlock, BetaUsage + + from pydantic_ai.builtin_tools import ToolSearchTool + from pydantic_ai.models.anthropic import AnthropicModel + from pydantic_ai.providers.anthropic import AnthropicProvider + + +pytestmark = [ + pytest.mark.skipif(not imports_successful(), reason='anthropic not installed'), +] + + +class TestToolDefinitionDeferLoading: + """Tests for ToolDefinition.defer_loading field.""" + + def test_defer_loading_defaults_to_false(self): + """Test that defer_loading defaults to False.""" + tool_def = ToolDefinition(name='test_tool') + assert tool_def.defer_loading is False + + def test_defer_loading_can_be_set(self): + """Test setting defer_loading.""" + tool_def = ToolDefinition(name='test_tool', defer_loading=True) + assert tool_def.defer_loading is True + + +class TestToolDeferLoading: + """Tests for Tool class with defer_loading.""" + + def test_tool_defer_loading_default(self): + """Test that Tool.defer_loading defaults to False.""" + + def my_tool(x: int) -> str: + return str(x) + + tool = Tool(my_tool) + assert tool.defer_loading is False + assert tool.tool_def.defer_loading is False + + def test_tool_defer_loading_set(self): + """Test setting Tool.defer_loading.""" + + def my_tool(x: int) -> str: + return str(x) + + tool = Tool(my_tool, defer_loading=True) + assert tool.defer_loading is True + assert tool.tool_def.defer_loading is True + + +class TestToolSearchTool: + """Tests for ToolSearchTool builtin tool.""" + + def test_tool_search_tool_defaults(self): + """Test ToolSearchTool defaults.""" + tool = ToolSearchTool() + assert tool.kind == 'tool_search' + assert tool.search_type is None + + def test_tool_search_tool_with_regex(self): + """Test ToolSearchTool with regex search type.""" + tool = ToolSearchTool(search_type='regex') + assert tool.search_type == 'regex' + + def test_tool_search_tool_with_bm25(self): + """Test ToolSearchTool with bm25 search type.""" + tool = ToolSearchTool(search_type='bm25') + assert tool.search_type == 'bm25' + + +class TestAnthropicMapToolDefinitionDeferLoading: + """Tests for AnthropicModel._map_tool_definition with defer_loading.""" + + def test_map_tool_definition_without_defer_loading(self): + """Test tool definition mapping without defer_loading.""" + tool_def = ToolDefinition( + name='test_tool', + description='A test tool', + ) + c = completion_message( + [BetaTextBlock(text='Hello', type='text')], + BetaUsage(input_tokens=5, output_tokens=10), + ) + mock_client = MockAnthropic.create_mock(c) + model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(anthropic_client=mock_client)) + result = model._map_tool_definition(tool_def) # pyright: ignore[reportPrivateUsage] + result_dict = cast(dict[str, Any], result) + assert 'defer_loading' not in result_dict + + def test_map_tool_definition_with_defer_loading(self): + """Test tool definition mapping with defer_loading.""" + tool_def = ToolDefinition( + name='test_tool', + description='A test tool', + defer_loading=True, + ) + c = completion_message( + [BetaTextBlock(text='Hello', type='text')], + BetaUsage(input_tokens=5, output_tokens=10), + ) + mock_client = MockAnthropic.create_mock(c) + model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(anthropic_client=mock_client)) + result = model._map_tool_definition(tool_def) # pyright: ignore[reportPrivateUsage] + result_dict = cast(dict[str, Any], result) + assert result_dict['defer_loading'] is True + + +class TestAgentWithDeferLoading: + """Tests for Agent with defer_loading on tools.""" + + def test_agent_with_defer_loading_tool(self): + """Test creating an agent with a tool that has defer_loading.""" + + def my_tool(x: int) -> str: + """A test tool.""" + return str(x) + + agent = Agent( + 'test', + tools=[Tool(my_tool, defer_loading=True)], + ) + + # Verify the tool was registered with defer_loading + tool = agent._function_toolset.tools.get('my_tool') # pyright: ignore[reportPrivateUsage] + assert tool is not None + assert tool.defer_loading is True + + def test_agent_tool_plain_decorator_with_defer_loading(self): + """Test the @agent.tool_plain decorator with defer_loading.""" + agent: Agent[None, str] = Agent('test') + + @agent.tool_plain(defer_loading=True) + def my_deferred_tool(x: int) -> str: + """A tool with defer_loading.""" + return str(x) + + tool = agent._function_toolset.tools.get('my_deferred_tool') # pyright: ignore[reportPrivateUsage] + assert tool is not None + assert tool.defer_loading is True + + +class TestAgentWithToolSearchTool: + """Tests for Agent with ToolSearchTool.""" + + def test_agent_with_tool_search_tool(self): + """Test creating an agent with ToolSearchTool.""" + agent = Agent( + 'test', + builtin_tools=[ToolSearchTool()], + ) + + assert any(isinstance(t, ToolSearchTool) for t in agent._builtin_tools) # pyright: ignore[reportPrivateUsage] + + def test_agent_with_tool_search_tool_bm25(self): + """Test creating an agent with ToolSearchTool using bm25.""" + agent = Agent( + 'test', + builtin_tools=[ToolSearchTool(search_type='bm25')], + ) + + tool_search = next(t for t in agent._builtin_tools if isinstance(t, ToolSearchTool)) # pyright: ignore[reportPrivateUsage] + assert tool_search.search_type == 'bm25' diff --git a/tests/models/test_model_request_parameters.py b/tests/models/test_model_request_parameters.py index 78a8dffa79..50fef0d9aa 100644 --- a/tests/models/test_model_request_parameters.py +++ b/tests/models/test_model_request_parameters.py @@ -68,6 +68,7 @@ def test_model_request_parameters_are_serializable(): 'sequential': False, 'kind': 'function', 'metadata': None, + 'defer_loading': False, } ], 'builtin_tools': [ @@ -131,6 +132,7 @@ def test_model_request_parameters_are_serializable(): 'sequential': False, 'kind': 'function', 'metadata': None, + 'defer_loading': False, } ], 'prompted_output_template': None, diff --git a/tests/test_logfire.py b/tests/test_logfire.py index dadb930dd0..f24f4e7290 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -547,6 +547,7 @@ async def my_ret(x: int) -> str: 'sequential': False, 'kind': 'function', 'metadata': None, + 'defer_loading': False, } ], 'builtin_tools': [], @@ -994,6 +995,7 @@ class MyOutput: 'sequential': False, 'kind': 'output', 'metadata': None, + 'defer_loading': False, } ], 'prompted_output_template': None, diff --git a/tests/test_tools.py b/tests/test_tools.py index bcdf537994..347267470c 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -150,6 +150,7 @@ def test_docstring_google(docstring_format: Literal['google', 'auto']): 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -184,6 +185,7 @@ def test_docstring_sphinx(docstring_format: Literal['sphinx', 'auto']): 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -226,6 +228,7 @@ def test_docstring_numpy(docstring_format: Literal['numpy', 'auto']): 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -268,6 +271,7 @@ def my_tool(x: int) -> str: # pragma: no cover 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -308,6 +312,7 @@ def my_tool(x: int) -> str: # pragma: no cover 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -354,6 +359,7 @@ def my_tool(x: int) -> str: # pragma: no cover 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -388,6 +394,7 @@ def test_only_returns_type(): 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -413,6 +420,7 @@ def test_docstring_unknown(): 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -456,6 +464,7 @@ def test_docstring_google_no_body(docstring_format: Literal['google', 'auto']): 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -492,6 +501,7 @@ def takes_just_model(model: Foo) -> str: 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -537,6 +547,7 @@ def takes_just_model(model: Foo, z: int) -> str: 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -902,6 +913,7 @@ def test_suppress_griffe_logging(caplog: LogCaptureFixture): 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } ) @@ -974,6 +986,7 @@ def my_tool_plain(*, a: int = 1, b: int) -> int: 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, }, { 'description': None, @@ -989,6 +1002,7 @@ def my_tool_plain(*, a: int = 1, b: int) -> int: 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, }, ] ) @@ -1077,6 +1091,7 @@ def my_tool(x: Annotated[str | None, WithJsonSchema({'type': 'string'})] = None, 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, }, { 'description': None, @@ -1090,6 +1105,7 @@ def my_tool(x: Annotated[str | None, WithJsonSchema({'type': 'string'})] = None, 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, }, ] ) @@ -1127,6 +1143,7 @@ def get_score(data: Data) -> int: ... # pragma: no branch 'kind': 'function', 'sequential': False, 'metadata': None, + 'defer_loading': False, } )