pydantic/pydantic-ai
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai.git synced 2026-10-12 02:21:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

Rehaul Tavily search tool (#4158)

Browse Source 
Co-authored-by: Claude Opus 4.5 <[email protected]>
Co-authored-by: Douwe Maan <[email protected]>
This commit is contained in:
Sarth Frey
2026-03-06 15:10:46 -06:00
committed by GitHub
co-authored by Claude Opus 4.5 Douwe Maan
parent d33d1c3d5b
commit 7db7152bd0
10 changed files with 639 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/common-tools.md
+33
View File
@@ -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
+110 -10
View File
@@ -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 @@ except ImportError as _import_error:
__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
+1
View File
@@ -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
+89
View File
@@ -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 [email protected] 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
+66
View File
@@ -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
tests/cassettes/test_tavily/test_search_with_include_domains.yaml
+119
View File
@@ -0,0 +1,119 @@
interactions:
- request:
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate, br, zstd
connection:
- keep-alive
content-length:
- '117'
content-type:
- application/json
host:
- api.tavily.com
method: POST
parsed_body:
include_domains:
- arxiv.org
query: transformer architectures
search_depth: basic
topic: general
uri: https://api.tavily.com/search
response:
headers:
connection:
- keep-alive
content-length:
- '6906'
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: transformer architectures
request_id: 5219ceb5-6bd9-4d38-b675-2fb7b8a887f2
response_time: 1.08
results:
- content: '# Title:The Power of Architecture: Deep Dive into Transformer Architectures for Long-Term Time Series Forecasting.
Authors:Lefei Shen, Mouxiang Chen, Han Fu, Xiaoxue Ren, Xiaoyun Joy Wang, Jianling Sun, Zhuo Li, Chenghao Liu. View
a PDF of the paper titled The Power of Architecture: Deep Dive into Transformer Architectures for Long-Term Time
Series Forecasting, by Lefei Shen and 7 other authors. View a PDF of the paper titled The Power of Architecture:
Deep Dive into Transformer Architectures for Long-Term Time Series Forecasting, by Lefei Shen and 7 other authors.
> Abstract:Transformer-based models have recently become dominant in Long-term Time Series Forecasting (LTSF), yet
the variations in their architecture, such as encoder-only, encoder-decoder, and decoder-only designs, raise a crucial
question: What Transformer architecture works best for LTSF tasks? | Cite as: | arXiv:2507.13043 [cs.LG] |. | |
(or arXiv:2507.13043v1 [cs.LG] for this version) |. # Bibliographic and Citation Tools. Have an idea for a project
that will add value for arXiv''s community?'
raw_content: null
score: 0.7923522
title: Deep Dive into Transformer Architectures for Long-Term ...
url: https://arxiv.org/abs/2507.13043
- content: '# Computer Science > Machine Learning. # Title:Optimal Control for Transformer Architectures: Enhancing
Generalization, Robustness and Efficiency. | Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI);
Optimization and Control (math.OC) |. | Cite as: | arXiv:2505.13499 [cs.LG] |. | | (or arXiv:2505.13499v2 [cs.LG]
for this version) |. | | Focus to learn more arXiv-issued DOI via DataCite |. ### References & Citations. ##
BibTeX formatted citation. # Bibliographic and Citation Tools. # Code, Data and Media Associated with this Article.
# Recommenders and Search Tools. # arXivLabs: experimental projects with community collaborators. arXivLabs is a
framework that allows collaborators to develop and share new arXiv features directly on our website. Both individuals
and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence,
and user data privacy. arXiv is committed to these values and only works with partners that adhere to them. Have
an idea for a project that will add value for arXiv''s community?'
raw_content: null
score: 0.783542
title: '[2505.13499] Optimal Control for Transformer Architectures'
url: https://arxiv.org/abs/2505.13499
- content: 'Transformer-based models have recently become dominant in Long-term Time Series Forecasting (LTSF), yet
the variations in their architecture, such as encoder-only, encoder-decoder, and decoder-only designs, raise a crucial
question: What Transformer architecture works best for LTSF tasks? In recent years, Transformer-based models have
become dominant in long-term time series forecasting (LTSF) tasks (Informer, ; Autoformer, ; FEDformer, ; PatchTST,
; iTransformer, ; TimeXer, ; ARMA\_Attention, ; Pyraformer, ; TFT, ; PDFormer, ; BasisFormer, ; SAMformer, ; Scaleformer,
; Quatformer, ), demonstrating strong performance across various real-world applications (TSF\_Energy\_1, ; TSF\_Energy\_2,
; TSF\_Economics\_1, ; TSF\_Web\_1, ; TSF\_Web\_2, ; TSF\_Web\_3, ; TSF\_Weather\_1, ; TSF\_Weather\_2, ; TSF\_Finance\_1,
). We examine Transformer-based LTSF models from multiple perspectives, including attention mechanisms, forecasting
aggregation strategies, forecasting paradigms, and normalization layers. Based on the above conclusions, we construct
an optimal Transformer architecture by combining the best choices, including bi-directional attention with joint-attention,
complete forecasting aggregation, direct-mapping paradigm, and the BatchNorm layer.'
raw_content: null
score: 0.77731717
title: Deep Dive into Transformer Architectures for Long-Term ...
url: https://arxiv.org/html/2507.13043v1
- content: '# Title:Lightweight Transformer Architectures for Edge Devices in Real-Time Applications. View a PDF of
the paper titled Lightweight Transformer Architectures for Edge Devices in Real-Time Applications, by Hema Hariharan
Samson. > Abstract:The deployment of transformer-based models on resource-constrained edge devices represents a
critical challenge in enabling real-time artificial intelligence applications. This comprehensive survey examines
lightweight transformer architectures specifically designed for edge deployment, analyzing recent advances in model
compression, quantization, pruning, and knowledge distillation techniques. Experimental results demonstrate that
modern lightweight transformers can achieve 75-96% of full-model accuracy while reducing model size by 4-10x and
inference latency by 3-9x, enabling deployment on devices with as little as 2-5W power consumption. Comprehensive
study of lightweight transformer architectures for edge computing with novel findings on memory-bandwidth tradeoffs,
quantization strategies, and hardware-specific optimizations. | | (or arXiv:2601.03290v1 [cs.LG] for this version)
|. View a PDF of the paper titled Lightweight Transformer Architectures for Edge Devices in Real-Time Applications,
by Hema Hariharan Samson.'
raw_content: null
score: 0.76826453
title: Lightweight Transformer Architectures for Edge Devices in ...
url: https://www.arxiv.org/abs/2601.03290
- content: '# Title:Study of Lightweight Transformer Architectures for Single-Channel Speech Enhancement. View a PDF
of the paper titled Study of Lightweight Transformer Architectures for Single-Channel Speech Enhancement, by Haixin
Zhao and Nilesh Madhu. Networks integrating stacked temporal and spectral modelling effectively leverage improved
architectures such as transformers; however, they inevitably incur substantial computational complexity and model
expansion. The proposed lightweight, causal, transformer-based architecture with adversarial training (LCT-GAN)
yields SoTA performance on instrumental metrics among contemporary lightweight models, but with far less overhead.
| Cite as: | arXiv:2505.21057 [eess.AS] |. | | (or arXiv:2505.21057v1 [eess.AS] for this version) |. View a PDF
of the paper titled Study of Lightweight Transformer Architectures for Single-Channel Speech Enhancement, by Haixin
Zhao and Nilesh Madhu. # Bibliographic and Citation Tools. arXiv is committed to these values and only works with
partners that adhere to them. Have an idea for a project that will add value for arXiv''s community?'
raw_content: null
score: 0.7663815
title: Study of Lightweight Transformer Architectures for Single ...
url: https://arxiv.org/abs/2505.21057
status:
code: 200
message: OK
version: 1
tests/conftest.py
+5
View File
@@ -532,6 +532,11 @@ def xai_api_key() -> str:
return os.getenv('XAI_API_KEY', 'mock-api-key')
@pytest.fixture(scope='session')
def tavily_api_key() -> str:
return os.getenv('TAVILY_API_KEY', 'mock-api-key')
@pytest.fixture(scope='function') # Needs to be function scoped to get the request node name
def xai_provider(request: pytest.FixtureRequest) -> Iterator[XaiProvider]:
"""xAI provider fixture backed by protobuf cassettes.
tests/test_examples.py
+7
View File
@@ -203,6 +203,7 @@ def test_docs_examples(
env.set('PYDANTIC_AI_GATEWAY_API_KEY', 'testing')
env.set('VOYAGE_API_KEY', 'testing')
env.set('XAI_API_KEY', 'testing')
env.set('TAVILY_API_KEY', 'testing')
prefix_settings = example.prefix_settings()
opt_test = prefix_settings.get('test', '')
@@ -573,6 +574,12 @@ text_responses: dict[str, str | ToolCallPart | Sequence[ToolCallPart]] = {
args={'name': 'test', 'value': 42},
tool_call_id='pyd_ai_tool_call_id',
),
'Find recent papers about transformer architectures': (
'Here are some recent papers about transformer architectures from arxiv.org:\n'
'\n'
'1. "Attention Is All You Need" - The foundational paper on the Transformer model.\n'
'2. "FlashAttention: Fast and Memory-Efficient Exact Attention" - Proposes an IO-aware attention algorithm.'
),
}
tool_responses: dict[tuple[str, str], str] = {
tests/test_tavily.py
+207
View File
@@ -0,0 +1,207 @@
"""Tests for the Tavily search tool."""
from __future__ import annotations
import pytest
from inline_snapshot import snapshot
from tavily import AsyncTavilyClient
from pydantic_ai._run_context import RunContext
from pydantic_ai.common_tools.tavily import TavilySearchTool, tavily_search_tool
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage
from .conftest import IsStr
@pytest.mark.vcr()
async def test_basic_search(tavily_api_key: str):
"""Test basic search with default parameters."""
tool = TavilySearchTool(client=AsyncTavilyClient(tavily_api_key))
results = await tool('What is Pydantic AI?')
assert results == snapshot(
[
{
'title': 'Pydantic AI: Agent Framework',
'url': 'https://medium.com/ai-agent-insider/pydantic-ai-agent-framework-02b138e8db71',
'content': IsStr(),
'score': 0.9999875,
},
{
'title': 'Pydantic AI - Pydantic AI',
'url': 'https://ai.pydantic.dev/',
'content': IsStr(),
'score': 0.99997807,
},
{
'title': 'Build Production-Ready AI Agents in Python with Pydantic AI',
'url': 'https://www.youtube.com/watch?v=-WB0T0XmDrY',
'content': IsStr(),
'score': 0.9999398,
},
{
'title': 'What is Pydantic AI?. Build production-ready AI agents with...',
'url': 'https://medium.com/@tahirbalarabe2/what-is-pydantic-ai-15cc81dea3c3',
'content': IsStr(),
'score': 0.9999125,
},
{
'title': 'Pydantic AI : r/LLMDevs',
'url': 'https://www.reddit.com/r/LLMDevs/comments/1iih8az/pydantic_ai/',
'content': IsStr(),
'score': 0.9999001,
},
]
)
@pytest.mark.vcr()
async def test_search_with_include_domains(tavily_api_key: str):
"""Test search with include_domains filtering."""
tool = TavilySearchTool(client=AsyncTavilyClient(tavily_api_key))
results = await tool('transformer architectures', include_domains=['arxiv.org'])
assert results == snapshot(
[
{
'title': 'Deep Dive into Transformer Architectures for Long-Term ...',
'url': 'https://arxiv.org/abs/2507.13043',
'content': IsStr(),
'score': 0.7923522,
},
{
'title': '[2505.13499] Optimal Control for Transformer Architectures',
'url': 'https://arxiv.org/abs/2505.13499',
'content': IsStr(),
'score': 0.783542,
},
{
'title': 'Deep Dive into Transformer Architectures for Long-Term ...',
'url': 'https://arxiv.org/html/2507.13043v1',
'content': IsStr(),
'score': 0.77731717,
},
{
'title': 'Lightweight Transformer Architectures for Edge Devices in ...',
'url': 'https://www.arxiv.org/abs/2601.03290',
'content': IsStr(),
'score': 0.76826453,
},
{
'title': 'Study of Lightweight Transformer Architectures for Single ...',
'url': 'https://arxiv.org/abs/2505.21057',
'content': IsStr(),
'score': 0.7663815,
},
]
)
@pytest.mark.vcr()
async def test_factory_with_bound_params(tavily_api_key: str):
"""Test factory-bound params are forwarded through FunctionSchema.call."""
tool = tavily_search_tool(tavily_api_key, max_results=2, include_domains=['arxiv.org'])
ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage())
results = await tool.function_schema.call({'query': 'attention mechanisms'}, ctx)
assert results == snapshot(
[
{
'title': '[2601.03329] Attention mechanisms in neural networks',
'url': 'https://arxiv.org/abs/2601.03329',
'content': IsStr(),
'score': 0.81770587,
},
{
'title': 'A General Survey on Attention Mechanisms in Deep ...',
'url': 'https://arxiv.org/abs/2203.14263',
'content': IsStr(),
'score': 0.8138313,
},
]
)
def test_no_params_bound_exposes_all_in_schema(tavily_api_key: str):
"""Test that with no factory params, all parameters appear in the tool schema."""
tool = tavily_search_tool(tavily_api_key)
assert tool.name == snapshot('tavily_search')
assert tool.function_schema.json_schema == snapshot(
{
'additionalProperties': False,
'properties': {
'query': {
'description': 'The search query to execute with Tavily.',
'type': 'string',
},
'search_depth': {
'default': 'basic',
'description': 'The depth of the search.',
'enum': ['basic', 'advanced', 'fast', 'ultra-fast'],
'type': 'string',
},
'topic': {
'default': 'general',
'description': 'The category of the search.',
'enum': ['general', 'news', 'finance'],
'type': 'string',
},
'time_range': {
'anyOf': [
{'enum': ['day', 'week', 'month', 'year'], 'type': 'string'},
{'type': 'null'},
],
'default': None,
'description': 'The time range back from the current date to filter results.',
},
'include_domains': {
'anyOf': [{'items': {'type': 'string'}, 'type': 'array'}, {'type': 'null'}],
'default': None,
'description': 'List of domains to specifically include in the search results.',
},
'exclude_domains': {
'anyOf': [{'items': {'type': 'string'}, 'type': 'array'}, {'type': 'null'}],
'default': None,
'description': 'List of domains to specifically exclude from the search results.',
},
},
'required': ['query'],
'type': 'object',
}
)
def test_factory_requires_api_key_or_client():
"""Test that tavily_search_tool raises when neither api_key nor client is provided."""
with pytest.raises(ValueError, match='Either api_key or client must be provided'):
tavily_search_tool() # pyright: ignore[reportCallIssue]
def test_factory_with_client():
"""Test that tavily_search_tool accepts a pre-built client."""
client = AsyncTavilyClient('test-key')
tool = tavily_search_tool(client=client)
assert tool.name == 'tavily_search'
def test_bound_params_hidden_from_schema(tavily_api_key: str):
"""Test that factory-provided params are excluded from the tool schema."""
tool = tavily_search_tool(
tavily_api_key,
search_depth='advanced',
topic='news',
time_range='week',
include_domains=['arxiv.org'],
exclude_domains=['medium.com'],
)
assert tool.function_schema.json_schema == snapshot(
{
'additionalProperties': False,
'properties': {
'query': {
'description': 'The search query to execute with Tavily.',
'type': 'string',
},
},
'required': ['query'],
'type': 'object',
}
)
uv.lock
Generated
+2
View File
@@ -6673,6 +6673,7 @@ dev = [
{ name = "pytest-recording" },
{ name = "pytest-xdist" },
{ name = "strict-no-cover" },
{ name = "tavily-python" },
]
docs = [
{ name = "black" },
@@ -6735,6 +6736,7 @@ dev = [
{ name = "pytest-recording", specifier = ">=0.13.2" },
{ name = "pytest-xdist", specifier = ">=3.6.1" },
{ name = "strict-no-cover", git = "https://github.com/pydantic/strict-no-cover.git?rev=7fc59da2c4dff919db2095a0f0e47101b657131d" },
{ name = "tavily-python", specifier = ">=0.5.0" },
]
docs = [
{ name = "black", specifier = ">=24.10.0" },