Rehaul Tavily search tool (#4158)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Douwe Maan <douwe@pydantic.dev>

Author: 3 people

docs/common-tools.md

@@ -139,6 +139,39 @@ Feel free to click on the links to dive deeper into each story!
 """
 ```
 
+### Configuring Parameters
+
+The `tavily_search_tool` factory accepts optional parameters that control search behavior. `max_results` is always developer-controlled and never appears in the LLM tool schema. Other parameters, when provided, are fixed for all searches and hidden from the LLM's tool schema. Parameters left unset remain available for the LLM to set per-call.
+
+For example, you can lock in `max_results` and `include_domains` at tool creation time while still letting the LLM control `exclude_domains`:
+
+```py {title="tavily_domain_filtering.py"}
+import os
+
+from pydantic_ai import Agent
+from pydantic_ai.common_tools.tavily import tavily_search_tool
+
+api_key = os.getenv('TAVILY_API_KEY')
+assert api_key is not None
+
+agent = Agent(
+    'openai:gpt-5.2',
+    tools=[tavily_search_tool(api_key, max_results=5, include_domains=['arxiv.org'])],
+    instructions='Search for information and return the results.',
+)
+
+result = agent.run_sync(
+    'Find recent papers about transformer architectures'
+)
+print(result.output)
+"""
+Here are some recent papers about transformer architectures from arxiv.org:
+
+1. "Attention Is All You Need" - The foundational paper on the Transformer model.
+2. "FlashAttention: Fast and Memory-Efficient Exact Attention" - Proposes an IO-aware attention algorithm.
+"""
+```
+
 ## Exa Search Tool
 
 !!! info

pydantic_ai_slim/pydantic_ai/common_tools/tavily.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass
-from typing import Literal
+from dataclasses import KW_ONLY, dataclass
+from functools import partial
+from inspect import signature
+from typing import Literal, overload
 
 from pydantic import TypeAdapter
 from typing_extensions import Any, TypedDict
@@ -16,6 +18,9 @@
 
 __all__ = ('tavily_search_tool',)
 
+_UNSET: Any = object()
+"""Sentinel to distinguish "not provided" from None in factory kwargs."""
+
 
 class TavilySearchResult(TypedDict):
     """A Tavily search result.
@@ -44,38 +49,133 @@ class TavilySearchTool:
     client: AsyncTavilyClient
     """The Tavily search client."""
 
+    _: KW_ONLY
+
+    max_results: int | None = None
+    """The maximum number of results. If None, the Tavily default is used."""
+
     async def __call__(
         self,
         query: str,
-        search_deep: Literal['basic', 'advanced'] = 'basic',
-        topic: Literal['general', 'news'] = 'general',
-        time_range: Literal['day', 'week', 'month', 'year', 'd', 'w', 'm', 'y'] | None = None,
+        search_depth: Literal['basic', 'advanced', 'fast', 'ultra-fast'] = 'basic',
+        topic: Literal['general', 'news', 'finance'] = 'general',
+        time_range: Literal['day', 'week', 'month', 'year'] | None = None,
+        include_domains: list[str] | None = None,
+        exclude_domains: list[str] | None = None,
     ) -> list[TavilySearchResult]:
         """Searches Tavily for the given query and returns the results.
 
         Args:
             query: The search query to execute with Tavily.
-            search_deep: The depth of the search.
+            search_depth: The depth of the search.
             topic: The category of the search.
             time_range: The time range back from the current date to filter results.
+            include_domains: List of domains to specifically include in the search results.
+            exclude_domains: List of domains to specifically exclude from the search results.
 
         Returns:
             A list of search results from Tavily.
         """
-        results = await self.client.search(query, search_depth=search_deep, topic=topic, time_range=time_range)  # type: ignore[reportUnknownMemberType]
+        results: dict[str, Any] = await self.client.search(  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            query,
+            search_depth=search_depth,
+            topic=topic,
+            time_range=time_range,  # pyright: ignore[reportArgumentType]
+            max_results=self.max_results,  # pyright: ignore[reportArgumentType]
+            include_domains=include_domains,  # pyright: ignore[reportArgumentType]
+            exclude_domains=exclude_domains,  # pyright: ignore[reportArgumentType]
+        )
         return tavily_search_ta.validate_python(results['results'])
 
 
-def tavily_search_tool(api_key: str):
+@overload
+def tavily_search_tool(
+    api_key: str,
+    *,
+    max_results: int | None = None,
+    search_depth: Literal['basic', 'advanced', 'fast', 'ultra-fast'] = _UNSET,
+    topic: Literal['general', 'news', 'finance'] = _UNSET,
+    time_range: Literal['day', 'week', 'month', 'year'] | None = _UNSET,
+    include_domains: list[str] | None = _UNSET,
+    exclude_domains: list[str] | None = _UNSET,
+) -> Tool[Any]: ...
+
+
+@overload
+def tavily_search_tool(
+    *,
+    client: AsyncTavilyClient,
+    max_results: int | None = None,
+    search_depth: Literal['basic', 'advanced', 'fast', 'ultra-fast'] = _UNSET,
+    topic: Literal['general', 'news', 'finance'] = _UNSET,
+    time_range: Literal['day', 'week', 'month', 'year'] | None = _UNSET,
+    include_domains: list[str] | None = _UNSET,
+    exclude_domains: list[str] | None = _UNSET,
+) -> Tool[Any]: ...
+
+
+def tavily_search_tool(
+    api_key: str | None = None,
+    *,
+    client: AsyncTavilyClient | None = None,
+    max_results: int | None = None,
+    search_depth: Literal['basic', 'advanced', 'fast', 'ultra-fast'] = _UNSET,
+    topic: Literal['general', 'news', 'finance'] = _UNSET,
+    time_range: Literal['day', 'week', 'month', 'year'] | None = _UNSET,
+    include_domains: list[str] | None = _UNSET,
+    exclude_domains: list[str] | None = _UNSET,
+) -> Tool[Any]:
     """Creates a Tavily search tool.
 
+    `max_results` is always developer-controlled and does not appear in the LLM tool schema.
+    Other parameters, when provided, are fixed for all searches and hidden from the LLM's
+    tool schema. Parameters left unset remain available for the LLM to set per-call.
+
     Args:
-        api_key: The Tavily API key.
+        api_key: The Tavily API key. Required if `client` is not provided.
 
             You can get one by signing up at [https://app.tavily.com/home](https://app.tavily.com/home).
+        client: An existing AsyncTavilyClient. If provided, `api_key` is ignored.
+            This is useful for sharing a client across multiple tool instances.
+        max_results: The maximum number of results. If None, the Tavily default is used.
+        search_depth: The depth of the search.
+        topic: The category of the search.
+        time_range: The time range back from the current date to filter results.
+        include_domains: List of domains to specifically include in the search results.
+        exclude_domains: List of domains to specifically exclude from the search results.
     """
+    if client is None:
+        if api_key is None:
+            raise ValueError('Either api_key or client must be provided')
+        client = AsyncTavilyClient(api_key)
+    func = TavilySearchTool(client=client, max_results=max_results).__call__
+
+    kwargs: dict[str, Any] = {}
+    if search_depth is not _UNSET:
+        kwargs['search_depth'] = search_depth
+    if topic is not _UNSET:
+        kwargs['topic'] = topic
+    if time_range is not _UNSET:
+        kwargs['time_range'] = time_range
+    if include_domains is not _UNSET:
+        kwargs['include_domains'] = include_domains
+    if exclude_domains is not _UNSET:
+        kwargs['exclude_domains'] = exclude_domains
+
+    if kwargs:
+        original = func
+        func = partial(func, **kwargs)
+        func.__name__ = original.__name__  # type: ignore[union-attr]
+        func.__qualname__ = original.__qualname__
+        # partial with keyword args only updates defaults, not removes params.
+        # Set __signature__ explicitly to exclude bound params from the tool schema.
+        orig_sig = signature(original)
+        func.__signature__ = orig_sig.replace(  # type: ignore[attr-defined]
+            parameters=[p for name, p in orig_sig.parameters.items() if name not in kwargs]
+        )
+
     return Tool[Any](
-        TavilySearchTool(client=AsyncTavilyClient(api_key)).__call__,
+        func,  # pyright: ignore[reportArgumentType]
         name='tavily_search',
         description='Searches Tavily for the given query and returns the results.',
     )

pyproject.toml

@@ -101,6 +101,7 @@ dev = [
     "dirty-equals>=0.9.0",
     "duckduckgo-search>=7.0.0",
     "exa-py>=2.0.0",
+    "tavily-python>=0.5.0",
     "inline-snapshot>=0.19.3",
     "pytest>=9.0.0",
     "pytest-examples>=0.0.18",

tests/cassettes/test_tavily/test_basic_search.yaml

@@ -0,0 +1,89 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate, br, zstd
+      connection:
+      - keep-alive
+      content-length:
+      - '78'
+      content-type:
+      - application/json
+      host:
+      - api.tavily.com
+    method: POST
+    parsed_body:
+      query: What is Pydantic AI?
+      search_depth: basic
+      topic: general
+    uri: https://api.tavily.com/search
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '3705'
+      content-security-policy:
+      - default-src 'none'; script-src 'self'; connect-src 'self'; img-src 'self'; style-src 'self';base-uri 'self';form-action
+        'self'; require-trusted-types-for 'script'; upgrade-insecure-requests;
+      content-type:
+      - application/json
+    parsed_body:
+      answer: null
+      follow_up_questions: null
+      images: []
+      query: What is Pydantic AI?
+      request_id: 3b2f7385-b01d-479e-9825-1ce0f8c59b93
+      response_time: 0.89
+      results:
+      - content: '## AI Agent Insider. # Pydantic AI: Agent Framework. Introducing **Pydantic AI**-- a groundbreaking Python
+          framework specifically designed to simplify the creation of production-grade AI agents. **Pydantic AI** is a Python
+          framework that acts as a bridge between developers and LLMs, providing tools to create **agents** -- entities that
+          execute specific tasks based on system prompts, functions, and structured outputs. Here''s a basic example of using
+          Pydantic AI to create an agent that responds to user queries:. from pydantic_ai import Agentagent = Agent("openai:gpt-4",
+          system_prompt="Be a helpful assistant.")result = await agent.run("Hello, how are you?")print(result.data) # Outputs
+          the response. from pydantic_ai import ModelRetry@agent.tooldef validate_data(ctx): if not ctx.input_data: raise
+          ModelRetry("Data missing, retrying..."). Pydantic AI is transforming how developers build AI agents. Install Pydantic
+          AI today and build your first agent!**. ### Agent Framework / shim to use Pydantic with LLMs. Contribute to pydantic/pydantic-ai
+          development by creating an account.... ## GitHub - pydantic/pydantic-ai: Agent Framework / shim to use Pydantic
+          with LLMs. ## Published in AI Agent Insider.'
+        raw_content: null
+        score: 0.9999875
+        title: 'Pydantic AI: Agent Framework'
+        url: https://medium.com/ai-agent-insider/pydantic-ai-agent-framework-02b138e8db71
+      - content: Pydantic AI is a Python agent framework designed to help you quickly, confidently, and painlessly build production
+          grade applications and workflows with
+        raw_content: null
+        score: 0.99997807
+        title: Pydantic AI - Pydantic AI
+        url: https://ai.pydantic.dev/
+      - content: Pydantic AI lets you integrate large language models like GPT-5 **directly into your Python applications**.
+        raw_content: null
+        score: 0.9999398
+        title: Build Production-Ready AI Agents in Python with Pydantic AI
+        url: https://www.youtube.com/watch?v=-WB0T0XmDrY
+      - content: '*"Pydantic AI is a Python agent framework designed to help you quickly, confidently, and painlessly build
+          production grade applications and workflows with Generative AI."*. So they built Pydantic AI with a single, simple
+          aim, to bring that FastAPI feeling to building applications and workflows with generative AI. A: Pydantic AI is
+          a Python agent framework from the creators of Pydantic Validation. Q: How is Pydantic AI different from other agent
+          frameworks like LangChain or LlamaIndex? Q: What LLM providers does Pydantic AI support? Q: Can I use my own custom
+          models with Pydantic AI? Q: Can I use Pydantic AI with existing FastAPI applications? A: Yes, Pydantic AI is designed
+          to integrate well with FastAPI and other Python web frameworks. A: Pydantic AI uses Pydantic models to define structured
+          output types, ensuring LLM responses are validated and type-safe. A: Yes, Pydantic AI is designed specifically for
+          production-grade applications with features like durable execution, observability integration, and type safety.'
+        raw_content: null
+        score: 0.9999125
+        title: What is Pydantic AI?. Build production-ready AI agents with...
+        url: https://medium.com/@tahirbalarabe2/what-is-pydantic-ai-15cc81dea3c3
+      - content: I've been using Pydantic AI to build some basic agents and multi agents and it seems quite straight forward
+          and I'm quite pleased with it.
+        raw_content: null
+        score: 0.9999001
+        title: 'Pydantic AI : r/LLMDevs'
+        url: https://www.reddit.com/r/LLMDevs/comments/1iih8az/pydantic_ai/
+    status:
+      code: 200
+      message: OK
+version: 1

tests/cassettes/test_tavily/test_factory_with_bound_params.yaml

@@ -0,0 +1,66 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate, br, zstd
+      connection:
+      - keep-alive
+      content-length:
+      - '130'
+      content-type:
+      - application/json
+      host:
+      - api.tavily.com
+    method: POST
+    parsed_body:
+      include_domains:
+      - arxiv.org
+      max_results: 2
+      query: attention mechanisms
+      search_depth: basic
+      topic: general
+    uri: https://api.tavily.com/search
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '1668'
+      content-security-policy:
+      - default-src 'none'; script-src 'self'; connect-src 'self'; img-src 'self'; style-src 'self';base-uri 'self';form-action
+        'self'; require-trusted-types-for 'script'; upgrade-insecure-requests;
+      content-type:
+      - application/json
+    parsed_body:
+      answer: null
+      follow_up_questions: null
+      images: []
+      query: attention mechanisms
+      request_id: 139680ef-e03d-48c8-a3bc-950081750288
+      response_time: 1.31
+      results:
+      - content: by H Hays · 2026 · Cited by 1 -- Attention mechanisms represent a fundamental paradigm shift in neural network
+          architectures, enabling models to selectively focus on relevant
+        raw_content: null
+        score: 0.81770587
+        title: '[2601.03329] Attention mechanisms in neural networks'
+        url: https://arxiv.org/abs/2601.03329
+      - content: '# Title:A General Survey on Attention Mechanisms in Deep Learning. View a PDF of the paper titled A General
+          Survey on Attention Mechanisms in Deep Learning, by Gianni Brauwers and Flavius Frasincar. > Abstract:Attention
+          is an important mechanism that can be employed for a variety of deep learning models across many different domains
+          and tasks. The various attention mechanisms are explained by means of a framework consisting of a general attention
+          model, uniform notation, and a comprehensive taxonomy of attention mechanisms. | Subjects: | Machine Learning (cs.LG)
+          |. | Cite as: | arXiv:2203.14263 [cs.LG] |. |  | (or  arXiv:2203.14263v1 [cs.LG] for this version) |. View a PDF
+          of the paper titled A General Survey on Attention Mechanisms in Deep Learning, by Gianni Brauwers and Flavius Frasincar.
+          ### References & Citations. # Bibliographic and Citation Tools. # Recommenders and Search Tools. Have an idea for
+          a project that will add value for arXiv''s community?'
+        raw_content: null
+        score: 0.8138313
+        title: A General Survey on Attention Mechanisms in Deep ...
+        url: https://arxiv.org/abs/2203.14263
+    status:
+      code: 200
+      message: OK
+version: 1