feat: native tool search on Anthropic and OpenAI, with custom search strategies on any provider

Merged
DO
6mo ago
+11973-1371

Summary

Adds tool search to Pydantic AI: agents with large toolsets can mark tools as defer_loading=True to keep them out of the model's initial context until they're discovered by keyword. The PR delivers two related features:

  1. Native provider tool search on supporting models. Deferred tools are sent with defer_loading on the wire and the provider manages their visibility — no search_tools function tool round-trips, no extra latency.
  2. Custom search strategies, independent of provider support. Pass any (ctx, queries, tools) -> names callable to ToolSearch(strategy=...) and it runs anywhere: locally as a search_tools function tool on providers without native search, or through the provider's client-executed mode (Anthropic, OpenAI Responses) where it still benefits from defer_loading on the wire.

Both surface through the new auto-injected [ToolSearch][pydantic_ai.capabilities.ToolSearch] capability, so the default behaviour "just works": mark a tool or toolset deferred and the agent picks the best path for the active model.

Provider support for native tool search

ProviderModelsNotes
AnthropicSonnet 4.5+, Opus 4.5+, Haiku 4.5+Server-executed bm25 (default) or regex. Custom callables ride the tool_reference block path.
OpenAI ResponsesGPT-5.4+Server-executed tool_search by default. Custom callables ride the execution='client' path. Named 'bm25'/'regex' strategies are Anthropic-only.
All others—Local fallback: an auto-injected search_tools function tool (keyword-overlap by default, or the user's callable).

Public API

  • defer_loading=True on [Tool][pydantic_ai.tools.Tool], [@agent.tool][pydantic_ai.agent.Agent.tool], and [@agent.tool_plain][pydantic_ai.agent.Agent.tool_plain]; [AbstractToolset.defer_loading()][pydantic_ai.toolsets.AbstractToolset.defer_loading] for whole toolsets (incl. MCP, FastMCP).
  • [ToolSearch][pydantic_ai.capabilities.ToolSearch] capability with:
    • strategy=None (default): use the best provider-native option, fall back locally.
    • strategy='keywords': force the local keyword-overlap algorithm regardless of provider.
    • strategy='bm25' / 'regex': force a specific Anthropic-native strategy. Raises UserError on providers that can't honor it (no silent fallback to keyword-overlap).
    • strategy=<callable>: custom search function with signature (ctx, queries: Sequence[str], tools) -> names. Always invoked locally as the source of truth; on Anthropic/OpenAI Responses it also wires into the provider's client-executed mode so deferred tools still ship with defer_loading.
    • tool_description, parameter_description, max_results: customize the local search_tools fallback.
  • Cross-provider typed message parts: [ToolSearchCallPart][pydantic_ai.messages.ToolSearchCallPart] / [ToolSearchReturnPart][pydantic_ai.messages.ToolSearchReturnPart] (local-fallback path), [NativeToolSearchCallPart][pydantic_ai.messages.NativeToolSearchCallPart] / [NativeToolSearchReturnPart][pydantic_ai.messages.NativeToolSearchReturnPart] (native server-side path). Typed accessors: typed_args, queries, discovered_tools, message. Discriminate cross-path via part.tool_kind == 'tool-search'.

Internal pieces

  • ToolDefinition.with_native — marks corpus tools that stay in the request when their managing native tool is supported (adapter applies the defer_loading wire format) and drop when not. Complementary to unless_native (drops them when the native tool is supported).
  • ToolDefinition.tool_kind — typed ToolPartKind = Literal['tool-search'] discriminator that drives promotion of ToolCallPart/ToolReturnPart to the typed subclasses above. Registry-driven; future typed native tools (web_search, etc.) extend the Literal and register a narrower.
  • AbstractNativeTool.optional — when True, the native tool is silently dropped on models without native support instead of raising. Set by the ToolSearch capability at construction for the None / 'keywords' / callable strategies (which have a local fallback); left False for 'bm25' / 'regex' (no local equivalent).
  • ToolSearchToolset wraps the agent's toolset, exposing the deferred-tool corpus and serving the local search_tools function-tool fallback when needed. Model.prepare_request picks the right path based on profile support.
  • Cross-provider history replay: synthesize_local_tool_search_messages translates NativeToolSearch*Part history into local-shape typed parts when the next turn runs against a provider without native tool-search, so previously discovered tools stay accessible across provider boundaries.

Test plan

  • Unit tests for ToolSearchToolset (corpus emission, custom strategy, dispatch)
  • Unit tests for prepare_request swap logic with with_native (supported / unsupported / optional-drop / named-native error path)
  • Live VCR cassettes against Anthropic (BM25 + regex) and OpenAI Responses (server-executed tool_search)
  • Custom-callable native paths covered for Anthropic (tool_reference) and OpenAI (execution='client')
  • Cross-provider history replay tests (native → local synthesis)
  • Typed message parts and accessors covered (discriminator dispatch, typed_args / queries / discovered_tools / message)
  • Existing tool-search, Anthropic, and OpenAI model tests still pass; auto-injection is transparent on models without native support

Checklist

  • Any AI generated code has been reviewed line-by-line by the human PR author, who stands by it.
  • No breaking changes in accordance with the version policy.
  • PR title is fit for the release changelog.