tavily-python 0.7.22__tar.gz → 0.7.24__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {tavily_python-0.7.22 → tavily_python-0.7.24}/PKG-INFO +72 -2
- {tavily_python-0.7.22 → tavily_python-0.7.24}/README.md +71 -1
- {tavily_python-0.7.22 → tavily_python-0.7.24}/setup.py +1 -1
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/async_tavily.py +77 -33
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/hybrid_rag/hybrid_rag.py +6 -3
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/tavily.py +63 -18
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/PKG-INFO +72 -2
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/SOURCES.txt +1 -0
- tavily_python-0.7.24/tests/test_custom_session.py +410 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/LICENSE +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/setup.cfg +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/__init__.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/config.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/errors.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/hybrid_rag/__init__.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/utils.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/dependency_links.txt +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/requires.txt +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/top_level.txt +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_crawl.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_errors.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_map.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_research.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_search.py +0 -0
- {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_session_pooling.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: tavily-python
|
|
3
|
-
Version: 0.7.
|
|
3
|
+
Version: 0.7.24
|
|
4
4
|
Summary: Python wrapper for the Tavily API
|
|
5
5
|
Home-page: https://github.com/tavily-ai/tavily-python
|
|
6
6
|
Author: Tavily AI
|
|
@@ -25,7 +25,7 @@ Dynamic: requires-dist
|
|
|
25
25
|
Dynamic: requires-python
|
|
26
26
|
Dynamic: summary
|
|
27
27
|
|
|
28
|
-
# Tavily Python
|
|
28
|
+
# Tavily Python SDK
|
|
29
29
|
|
|
30
30
|
[](https://github.com/tavily-ai/tavily-python/stargazers)
|
|
31
31
|
[](https://pypi.org/project/tavily-python/)
|
|
@@ -273,6 +273,76 @@ for chunk in stream:
|
|
|
273
273
|
print(chunk.decode('utf-8'))
|
|
274
274
|
```
|
|
275
275
|
|
|
276
|
+
## Advanced: Custom Session / Client Injection
|
|
277
|
+
|
|
278
|
+
For enterprise environments that proxy Tavily traffic through an API gateway (e.g., for centralized auth, logging, or policy enforcement), you can pass a pre-configured HTTP session instead of a Tavily API key.
|
|
279
|
+
|
|
280
|
+
### Sync (custom `requests.Session`)
|
|
281
|
+
|
|
282
|
+
```python
|
|
283
|
+
import requests
|
|
284
|
+
from tavily import TavilyClient
|
|
285
|
+
|
|
286
|
+
# Pre-configure a session with your gateway's auth
|
|
287
|
+
session = requests.Session()
|
|
288
|
+
session.headers["Authorization"] = "Bearer your-gateway-token"
|
|
289
|
+
session.headers["X-Subscription-Key"] = "your-subscription-key"
|
|
290
|
+
|
|
291
|
+
# No Tavily API key needed — auth is handled by the session
|
|
292
|
+
client = TavilyClient(
|
|
293
|
+
session=session,
|
|
294
|
+
api_base_url="https://your-gateway.com/tavily",
|
|
295
|
+
)
|
|
296
|
+
|
|
297
|
+
response = client.search("latest AI research")
|
|
298
|
+
```
|
|
299
|
+
|
|
300
|
+
### Async (custom `httpx.AsyncClient`)
|
|
301
|
+
|
|
302
|
+
```python
|
|
303
|
+
import httpx
|
|
304
|
+
from tavily import AsyncTavilyClient
|
|
305
|
+
|
|
306
|
+
# Pre-configure an async client with your gateway's auth
|
|
307
|
+
custom_client = httpx.AsyncClient(
|
|
308
|
+
headers={"Authorization": "Bearer your-gateway-token"},
|
|
309
|
+
base_url="https://your-gateway.com/tavily",
|
|
310
|
+
)
|
|
311
|
+
|
|
312
|
+
client = AsyncTavilyClient(client=custom_client)
|
|
313
|
+
|
|
314
|
+
response = await client.search("latest AI research")
|
|
315
|
+
```
|
|
316
|
+
|
|
317
|
+
**Key behaviors:**
|
|
318
|
+
- If a custom session/client is provided, `api_key` is optional
|
|
319
|
+
- Custom session headers take precedence over SDK defaults (e.g., your `Authorization` won't be overwritten)
|
|
320
|
+
- Custom session proxies take precedence over SDK proxy settings
|
|
321
|
+
- The SDK will **not** close externally-provided sessions — you manage the lifecycle
|
|
322
|
+
|
|
323
|
+
## Session & User Tracking
|
|
324
|
+
|
|
325
|
+
`session_id`, `human_id`, and `client_name` are optional identifiers that help attribute requests to a logical session, an end user, and a named client. All three are sent as HTTP headers (`X-Session-Id`, `X-Human-Id`, `X-Client-Name`) and are never persisted in raw form — `human_id` is hashed server-side.
|
|
326
|
+
|
|
327
|
+
Set them once at client init, or per-call (per-call wins):
|
|
328
|
+
|
|
329
|
+
```python
|
|
330
|
+
from tavily import TavilyClient
|
|
331
|
+
|
|
332
|
+
# Client-level — applied to every request
|
|
333
|
+
client = TavilyClient(
|
|
334
|
+
api_key="tvly-YOUR_API_KEY",
|
|
335
|
+
session_id="my-session-123",
|
|
336
|
+
human_id="internal-user-id-42",
|
|
337
|
+
client_name="my-app",
|
|
338
|
+
)
|
|
339
|
+
|
|
340
|
+
# Per-call override
|
|
341
|
+
client.search("hello", session_id="ad-hoc-session")
|
|
342
|
+
```
|
|
343
|
+
|
|
344
|
+
All three are opt-in. Leave them unset and the SDK sends nothing — behavior is identical to earlier versions.
|
|
345
|
+
|
|
276
346
|
## Documentation
|
|
277
347
|
|
|
278
348
|
For a complete guide on how to use the different endpoints and their parameters, please head to our [Python API Reference](https://docs.tavily.com/sdk/python/reference).
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
# Tavily Python
|
|
1
|
+
# Tavily Python SDK
|
|
2
2
|
|
|
3
3
|
[](https://github.com/tavily-ai/tavily-python/stargazers)
|
|
4
4
|
[](https://pypi.org/project/tavily-python/)
|
|
@@ -246,6 +246,76 @@ for chunk in stream:
|
|
|
246
246
|
print(chunk.decode('utf-8'))
|
|
247
247
|
```
|
|
248
248
|
|
|
249
|
+
## Advanced: Custom Session / Client Injection
|
|
250
|
+
|
|
251
|
+
For enterprise environments that proxy Tavily traffic through an API gateway (e.g., for centralized auth, logging, or policy enforcement), you can pass a pre-configured HTTP session instead of a Tavily API key.
|
|
252
|
+
|
|
253
|
+
### Sync (custom `requests.Session`)
|
|
254
|
+
|
|
255
|
+
```python
|
|
256
|
+
import requests
|
|
257
|
+
from tavily import TavilyClient
|
|
258
|
+
|
|
259
|
+
# Pre-configure a session with your gateway's auth
|
|
260
|
+
session = requests.Session()
|
|
261
|
+
session.headers["Authorization"] = "Bearer your-gateway-token"
|
|
262
|
+
session.headers["X-Subscription-Key"] = "your-subscription-key"
|
|
263
|
+
|
|
264
|
+
# No Tavily API key needed — auth is handled by the session
|
|
265
|
+
client = TavilyClient(
|
|
266
|
+
session=session,
|
|
267
|
+
api_base_url="https://your-gateway.com/tavily",
|
|
268
|
+
)
|
|
269
|
+
|
|
270
|
+
response = client.search("latest AI research")
|
|
271
|
+
```
|
|
272
|
+
|
|
273
|
+
### Async (custom `httpx.AsyncClient`)
|
|
274
|
+
|
|
275
|
+
```python
|
|
276
|
+
import httpx
|
|
277
|
+
from tavily import AsyncTavilyClient
|
|
278
|
+
|
|
279
|
+
# Pre-configure an async client with your gateway's auth
|
|
280
|
+
custom_client = httpx.AsyncClient(
|
|
281
|
+
headers={"Authorization": "Bearer your-gateway-token"},
|
|
282
|
+
base_url="https://your-gateway.com/tavily",
|
|
283
|
+
)
|
|
284
|
+
|
|
285
|
+
client = AsyncTavilyClient(client=custom_client)
|
|
286
|
+
|
|
287
|
+
response = await client.search("latest AI research")
|
|
288
|
+
```
|
|
289
|
+
|
|
290
|
+
**Key behaviors:**
|
|
291
|
+
- If a custom session/client is provided, `api_key` is optional
|
|
292
|
+
- Custom session headers take precedence over SDK defaults (e.g., your `Authorization` won't be overwritten)
|
|
293
|
+
- Custom session proxies take precedence over SDK proxy settings
|
|
294
|
+
- The SDK will **not** close externally-provided sessions — you manage the lifecycle
|
|
295
|
+
|
|
296
|
+
## Session & User Tracking
|
|
297
|
+
|
|
298
|
+
`session_id`, `human_id`, and `client_name` are optional identifiers that help attribute requests to a logical session, an end user, and a named client. All three are sent as HTTP headers (`X-Session-Id`, `X-Human-Id`, `X-Client-Name`) and are never persisted in raw form — `human_id` is hashed server-side.
|
|
299
|
+
|
|
300
|
+
Set them once at client init, or per-call (per-call wins):
|
|
301
|
+
|
|
302
|
+
```python
|
|
303
|
+
from tavily import TavilyClient
|
|
304
|
+
|
|
305
|
+
# Client-level — applied to every request
|
|
306
|
+
client = TavilyClient(
|
|
307
|
+
api_key="tvly-YOUR_API_KEY",
|
|
308
|
+
session_id="my-session-123",
|
|
309
|
+
human_id="internal-user-id-42",
|
|
310
|
+
client_name="my-app",
|
|
311
|
+
)
|
|
312
|
+
|
|
313
|
+
# Per-call override
|
|
314
|
+
client.search("hello", session_id="ad-hoc-session")
|
|
315
|
+
```
|
|
316
|
+
|
|
317
|
+
All three are opt-in. Leave them unset and the SDK sends nothing — behavior is identical to earlier versions.
|
|
318
|
+
|
|
249
319
|
## Documentation
|
|
250
320
|
|
|
251
321
|
For a complete guide on how to use the different endpoints and their parameters, please head to our [Python API Reference](https://docs.tavily.com/sdk/python/reference).
|
|
@@ -19,48 +19,69 @@ class AsyncTavilyClient:
|
|
|
19
19
|
proxies: Optional[dict[str, str]] = None,
|
|
20
20
|
api_base_url: Optional[str] = None,
|
|
21
21
|
client_source: Optional[str] = None,
|
|
22
|
-
project_id: Optional[str] = None
|
|
22
|
+
project_id: Optional[str] = None,
|
|
23
|
+
session_id: Optional[str] = None,
|
|
24
|
+
human_id: Optional[str] = None,
|
|
25
|
+
client_name: Optional[str] = None,
|
|
26
|
+
client: Optional[httpx.AsyncClient] = None):
|
|
23
27
|
if api_key is None:
|
|
24
28
|
api_key = os.getenv("TAVILY_API_KEY")
|
|
25
29
|
|
|
26
|
-
if not api_key:
|
|
30
|
+
if not api_key and client is None:
|
|
27
31
|
raise MissingAPIKeyError()
|
|
28
32
|
|
|
29
|
-
|
|
33
|
+
tavily_project = project_id or os.getenv("TAVILY_PROJECT")
|
|
34
|
+
|
|
35
|
+
self._api_base_url = api_base_url or "https://api.tavily.com"
|
|
36
|
+
self._company_info_tags = company_info_tags
|
|
30
37
|
|
|
31
|
-
|
|
32
|
-
"
|
|
33
|
-
"
|
|
38
|
+
default_headers = {
|
|
39
|
+
"Content-Type": "application/json",
|
|
40
|
+
**({"Authorization": f"Bearer {api_key}"} if api_key else {}),
|
|
41
|
+
"X-Client-Source": client_source or "tavily-python",
|
|
42
|
+
**({"X-Project-ID": tavily_project} if tavily_project else {}),
|
|
43
|
+
**({"X-Session-Id": session_id} if session_id else {}),
|
|
44
|
+
**({"X-Human-Id": human_id} if human_id else {}),
|
|
45
|
+
**({"X-Client-Name": client_name} if client_name else {}),
|
|
34
46
|
}
|
|
35
47
|
|
|
36
|
-
|
|
48
|
+
self._external_client = client is not None
|
|
49
|
+
|
|
50
|
+
if client is not None:
|
|
51
|
+
self._client = client
|
|
52
|
+
# Only set headers that aren't already configured on the external client
|
|
53
|
+
for key, value in default_headers.items():
|
|
54
|
+
if key not in self._client.headers:
|
|
55
|
+
self._client.headers[key] = value
|
|
56
|
+
# Set base_url if the external client doesn't have one
|
|
57
|
+
if not str(self._client.base_url):
|
|
58
|
+
self._client.base_url = self._api_base_url
|
|
59
|
+
else:
|
|
60
|
+
proxies = proxies or {}
|
|
37
61
|
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
)
|
|
62
|
+
mapped_proxies = {
|
|
63
|
+
"http://": proxies.get("http", os.getenv("TAVILY_HTTP_PROXY")),
|
|
64
|
+
"https://": proxies.get("https", os.getenv("TAVILY_HTTPS_PROXY")),
|
|
65
|
+
}
|
|
43
66
|
|
|
44
|
-
|
|
67
|
+
mapped_proxies = {key: value for key, value in mapped_proxies.items() if value}
|
|
45
68
|
|
|
46
|
-
|
|
69
|
+
proxy_mounts = (
|
|
70
|
+
{scheme: httpx.AsyncHTTPTransport(proxy=proxy) for scheme, proxy in mapped_proxies.items()}
|
|
71
|
+
if mapped_proxies
|
|
72
|
+
else None
|
|
73
|
+
)
|
|
47
74
|
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
"X-Client-Source": client_source or "tavily-python",
|
|
54
|
-
**({"X-Project-ID": tavily_project} if tavily_project else {})
|
|
55
|
-
},
|
|
56
|
-
base_url=self._api_base_url,
|
|
57
|
-
mounts=proxy_mounts
|
|
58
|
-
)
|
|
59
|
-
self._company_info_tags = company_info_tags
|
|
75
|
+
self._client = httpx.AsyncClient(
|
|
76
|
+
headers=default_headers,
|
|
77
|
+
base_url=self._api_base_url,
|
|
78
|
+
mounts=proxy_mounts
|
|
79
|
+
)
|
|
60
80
|
|
|
61
81
|
async def close(self):
|
|
62
82
|
"""Close the client and release connection pool resources."""
|
|
63
|
-
|
|
83
|
+
if not self._external_client:
|
|
84
|
+
await self._client.aclose()
|
|
64
85
|
|
|
65
86
|
async def __aenter__(self):
|
|
66
87
|
return self
|
|
@@ -68,6 +89,23 @@ class AsyncTavilyClient:
|
|
|
68
89
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
69
90
|
await self.close()
|
|
70
91
|
|
|
92
|
+
@staticmethod
|
|
93
|
+
def _pop_request_headers(kwargs: dict) -> Optional[dict]:
|
|
94
|
+
"""Pop session_id, human_id, and client_name from kwargs and return them as headers.
|
|
95
|
+
|
|
96
|
+
Returns None when no overrides are provided so callers can omit the headers kwarg.
|
|
97
|
+
"""
|
|
98
|
+
overrides = {}
|
|
99
|
+
for key, header_name in (
|
|
100
|
+
("session_id", "X-Session-Id"),
|
|
101
|
+
("human_id", "X-Human-Id"),
|
|
102
|
+
("client_name", "X-Client-Name"),
|
|
103
|
+
):
|
|
104
|
+
value = kwargs.pop(key, None)
|
|
105
|
+
if value is not None:
|
|
106
|
+
overrides[header_name] = str(value)
|
|
107
|
+
return overrides or None
|
|
108
|
+
|
|
71
109
|
async def _search(
|
|
72
110
|
self,
|
|
73
111
|
query: str,
|
|
@@ -117,13 +155,14 @@ class AsyncTavilyClient:
|
|
|
117
155
|
|
|
118
156
|
data = {k: v for k, v in data.items() if v is not None}
|
|
119
157
|
|
|
158
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
120
159
|
if kwargs:
|
|
121
160
|
data.update(kwargs)
|
|
122
161
|
|
|
123
162
|
timeout = min(timeout, 120)
|
|
124
163
|
|
|
125
164
|
try:
|
|
126
|
-
response = await self._client.post("/search", content=json.dumps(data), timeout=timeout)
|
|
165
|
+
response = await self._client.post("/search", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
127
166
|
except httpx.TimeoutException:
|
|
128
167
|
raise TimeoutError(timeout)
|
|
129
168
|
|
|
@@ -232,11 +271,12 @@ class AsyncTavilyClient:
|
|
|
232
271
|
|
|
233
272
|
data = {k: v for k, v in data.items() if v is not None}
|
|
234
273
|
|
|
274
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
235
275
|
if kwargs:
|
|
236
276
|
data.update(kwargs)
|
|
237
277
|
|
|
238
278
|
try:
|
|
239
|
-
response = await self._client.post("/extract", content=json.dumps(data), timeout=timeout)
|
|
279
|
+
response = await self._client.post("/extract", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
240
280
|
except httpx.TimeoutException:
|
|
241
281
|
raise TimeoutError(timeout)
|
|
242
282
|
|
|
@@ -340,13 +380,14 @@ class AsyncTavilyClient:
|
|
|
340
380
|
"chunks_per_source": chunks_per_source,
|
|
341
381
|
}
|
|
342
382
|
|
|
383
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
343
384
|
if kwargs:
|
|
344
385
|
data.update(kwargs)
|
|
345
386
|
|
|
346
387
|
data = {k: v for k, v in data.items() if v is not None}
|
|
347
388
|
|
|
348
389
|
try:
|
|
349
|
-
response = await self._client.post("/crawl", content=json.dumps(data), timeout=timeout)
|
|
390
|
+
response = await self._client.post("/crawl", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
350
391
|
except httpx.TimeoutException:
|
|
351
392
|
raise TimeoutError(timeout)
|
|
352
393
|
|
|
@@ -450,13 +491,14 @@ class AsyncTavilyClient:
|
|
|
450
491
|
"include_usage": include_usage,
|
|
451
492
|
}
|
|
452
493
|
|
|
494
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
453
495
|
if kwargs:
|
|
454
496
|
data.update(kwargs)
|
|
455
497
|
|
|
456
498
|
data = {k: v for k, v in data.items() if v is not None}
|
|
457
499
|
|
|
458
500
|
try:
|
|
459
|
-
response = await self._client.post("/map", content=json.dumps(data), timeout=timeout)
|
|
501
|
+
response = await self._client.post("/map", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
460
502
|
except httpx.TimeoutException:
|
|
461
503
|
raise TimeoutError(timeout)
|
|
462
504
|
|
|
@@ -644,6 +686,7 @@ class AsyncTavilyClient:
|
|
|
644
686
|
|
|
645
687
|
data = {k: v for k, v in data.items() if v is not None}
|
|
646
688
|
|
|
689
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
647
690
|
if kwargs:
|
|
648
691
|
data.update(kwargs)
|
|
649
692
|
|
|
@@ -654,7 +697,8 @@ class AsyncTavilyClient:
|
|
|
654
697
|
"POST",
|
|
655
698
|
"/research",
|
|
656
699
|
content=json.dumps(data),
|
|
657
|
-
timeout=timeout
|
|
700
|
+
timeout=timeout,
|
|
701
|
+
**({"headers": override_headers} if override_headers else {})
|
|
658
702
|
) as response:
|
|
659
703
|
if response.status_code != 200:
|
|
660
704
|
try:
|
|
@@ -686,7 +730,7 @@ class AsyncTavilyClient:
|
|
|
686
730
|
else:
|
|
687
731
|
async def _make_request():
|
|
688
732
|
try:
|
|
689
|
-
response = await self._client.post("/research", content=json.dumps(data), timeout=timeout)
|
|
733
|
+
response = await self._client.post("/research", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
690
734
|
except httpx.TimeoutException:
|
|
691
735
|
raise TimeoutError(timeout)
|
|
692
736
|
|
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
import os
|
|
2
2
|
from typing import Union, Optional, Literal
|
|
3
3
|
|
|
4
|
+
import requests
|
|
4
5
|
from tavily import TavilyClient
|
|
5
6
|
|
|
6
7
|
try:
|
|
@@ -76,7 +77,8 @@ class TavilyHybridClient():
|
|
|
76
77
|
embeddings_field: str = 'embeddings',
|
|
77
78
|
content_field: str = 'content',
|
|
78
79
|
embedding_function: Optional[callable] = None,
|
|
79
|
-
ranking_function: Optional[callable] = None
|
|
80
|
+
ranking_function: Optional[callable] = None,
|
|
81
|
+
session: Optional[requests.Session] = None
|
|
80
82
|
):
|
|
81
83
|
'''
|
|
82
84
|
A client for performing hybrid RAG using both the Tavily API and a local database collection.
|
|
@@ -90,9 +92,10 @@ class TavilyHybridClient():
|
|
|
90
92
|
content_field (str): The name of the field in the collection that contains the content.
|
|
91
93
|
embedding_function (callable): If provided, this function will be used to generate embeddings for the search query and documents.
|
|
92
94
|
ranking_function (callable): If provided, this function will be used to rerank the combined results.
|
|
95
|
+
session (requests.Session): If provided, this pre-configured session will be used for HTTP requests. When set, api_key is optional.
|
|
93
96
|
'''
|
|
94
|
-
|
|
95
|
-
self.tavily = TavilyClient(api_key)
|
|
97
|
+
|
|
98
|
+
self.tavily = TavilyClient(api_key, session=session)
|
|
96
99
|
|
|
97
100
|
if db_provider != 'mongodb':
|
|
98
101
|
raise ValueError("Only MongoDB is currently supported as a database provider.")
|
|
@@ -11,11 +11,22 @@ class TavilyClient:
|
|
|
11
11
|
Tavily API client class.
|
|
12
12
|
"""
|
|
13
13
|
|
|
14
|
-
def __init__(
|
|
14
|
+
def __init__(
|
|
15
|
+
self,
|
|
16
|
+
api_key: Optional[str] = None,
|
|
17
|
+
proxies: Optional[dict[str, str]] = None,
|
|
18
|
+
api_base_url: Optional[str] = None,
|
|
19
|
+
client_source: Optional[str] = None,
|
|
20
|
+
project_id: Optional[str] = None,
|
|
21
|
+
session_id: Optional[str] = None,
|
|
22
|
+
human_id: Optional[str] = None,
|
|
23
|
+
client_name: Optional[str] = None,
|
|
24
|
+
session: Optional[requests.Session] = None,
|
|
25
|
+
):
|
|
15
26
|
if api_key is None:
|
|
16
27
|
api_key = os.getenv("TAVILY_API_KEY")
|
|
17
28
|
|
|
18
|
-
if not api_key:
|
|
29
|
+
if not api_key and session is None:
|
|
19
30
|
raise MissingAPIKeyError()
|
|
20
31
|
|
|
21
32
|
resolved_proxies = {
|
|
@@ -25,26 +36,36 @@ class TavilyClient:
|
|
|
25
36
|
|
|
26
37
|
resolved_proxies = {k: v for k, v in resolved_proxies.items() if v} or None
|
|
27
38
|
tavily_project = project_id or os.getenv("TAVILY_PROJECT")
|
|
28
|
-
|
|
39
|
+
|
|
29
40
|
self.base_url = api_base_url or "https://api.tavily.com"
|
|
30
41
|
self.api_key = api_key
|
|
31
42
|
self.proxies = resolved_proxies
|
|
32
|
-
|
|
43
|
+
|
|
33
44
|
self.headers = {
|
|
34
45
|
"Content-Type": "application/json",
|
|
35
|
-
"Authorization": f"Bearer {self.api_key}",
|
|
46
|
+
**({"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}),
|
|
36
47
|
"X-Client-Source": client_source or "tavily-python",
|
|
37
|
-
**({"X-Project-ID": tavily_project} if tavily_project else {})
|
|
48
|
+
**({"X-Project-ID": tavily_project} if tavily_project else {}),
|
|
49
|
+
**({"X-Session-Id": session_id} if session_id else {}),
|
|
50
|
+
**({"X-Human-Id": human_id} if human_id else {}),
|
|
51
|
+
**({"X-Client-Name": client_name} if client_name else {}),
|
|
38
52
|
}
|
|
39
53
|
|
|
40
|
-
self.
|
|
41
|
-
self.session.
|
|
54
|
+
self._external_session = session is not None
|
|
55
|
+
self.session = session if session is not None else requests.Session()
|
|
56
|
+
# For external sessions, only set headers that aren't already configured
|
|
57
|
+
for key, value in self.headers.items():
|
|
58
|
+
if key not in self.session.headers:
|
|
59
|
+
self.session.headers[key] = value
|
|
42
60
|
if self.proxies:
|
|
43
|
-
self.
|
|
61
|
+
for protocol, url in self.proxies.items():
|
|
62
|
+
if protocol not in self.session.proxies:
|
|
63
|
+
self.session.proxies[protocol] = url
|
|
44
64
|
|
|
45
65
|
def close(self):
|
|
46
66
|
"""Close the session and release resources."""
|
|
47
|
-
self.
|
|
67
|
+
if not self._external_session:
|
|
68
|
+
self.session.close()
|
|
48
69
|
|
|
49
70
|
def __enter__(self):
|
|
50
71
|
return self
|
|
@@ -52,6 +73,23 @@ class TavilyClient:
|
|
|
52
73
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
53
74
|
self.close()
|
|
54
75
|
|
|
76
|
+
@staticmethod
|
|
77
|
+
def _pop_request_headers(kwargs: dict) -> Optional[dict]:
|
|
78
|
+
"""Pop session_id, human_id, and client_name from kwargs and return them as headers.
|
|
79
|
+
|
|
80
|
+
Returns None when no overrides are provided so callers can omit the headers kwarg.
|
|
81
|
+
"""
|
|
82
|
+
overrides = {}
|
|
83
|
+
for key, header_name in (
|
|
84
|
+
("session_id", "X-Session-Id"),
|
|
85
|
+
("human_id", "X-Human-Id"),
|
|
86
|
+
("client_name", "X-Client-Name"),
|
|
87
|
+
):
|
|
88
|
+
value = kwargs.pop(key, None)
|
|
89
|
+
if value is not None:
|
|
90
|
+
overrides[header_name] = str(value)
|
|
91
|
+
return overrides or None
|
|
92
|
+
|
|
55
93
|
def _search(self,
|
|
56
94
|
query: str,
|
|
57
95
|
search_depth: Literal["basic", "advanced", "fast", "ultra-fast"] = None,
|
|
@@ -101,6 +139,7 @@ class TavilyClient:
|
|
|
101
139
|
|
|
102
140
|
data = {k: v for k, v in data.items() if v is not None}
|
|
103
141
|
|
|
142
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
104
143
|
if kwargs:
|
|
105
144
|
data.update(kwargs)
|
|
106
145
|
|
|
@@ -109,7 +148,7 @@ class TavilyClient:
|
|
|
109
148
|
payload = json.dumps(data)
|
|
110
149
|
|
|
111
150
|
try:
|
|
112
|
-
response = self.session.post(url, data=payload, timeout=timeout)
|
|
151
|
+
response = self.session.post(url, data=payload, timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
113
152
|
except requests.exceptions.Timeout:
|
|
114
153
|
raise TimeoutError(timeout)
|
|
115
154
|
|
|
@@ -212,11 +251,12 @@ class TavilyClient:
|
|
|
212
251
|
|
|
213
252
|
data = {k: v for k, v in data.items() if v is not None}
|
|
214
253
|
|
|
254
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
215
255
|
if kwargs:
|
|
216
256
|
data.update(kwargs)
|
|
217
257
|
|
|
218
258
|
try:
|
|
219
|
-
response = self.session.post(self.base_url + "/extract", data=json.dumps(data), timeout=timeout)
|
|
259
|
+
response = self.session.post(self.base_url + "/extract", data=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
220
260
|
except requests.exceptions.Timeout:
|
|
221
261
|
raise TimeoutError(timeout)
|
|
222
262
|
|
|
@@ -313,13 +353,14 @@ class TavilyClient:
|
|
|
313
353
|
"chunks_per_source": chunks_per_source,
|
|
314
354
|
}
|
|
315
355
|
|
|
356
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
316
357
|
if kwargs:
|
|
317
358
|
data.update(kwargs)
|
|
318
|
-
|
|
359
|
+
|
|
319
360
|
data = {k: v for k, v in data.items() if v is not None}
|
|
320
361
|
|
|
321
362
|
try:
|
|
322
|
-
response = self.session.post(self.base_url + "/crawl", data=json.dumps(data), timeout=timeout)
|
|
363
|
+
response = self.session.post(self.base_url + "/crawl", data=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
323
364
|
except requests.exceptions.Timeout:
|
|
324
365
|
raise TimeoutError(timeout)
|
|
325
366
|
|
|
@@ -421,13 +462,14 @@ class TavilyClient:
|
|
|
421
462
|
"include_usage": include_usage,
|
|
422
463
|
}
|
|
423
464
|
|
|
465
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
424
466
|
if kwargs:
|
|
425
467
|
data.update(kwargs)
|
|
426
|
-
|
|
468
|
+
|
|
427
469
|
data = {k: v for k, v in data.items() if v is not None}
|
|
428
470
|
|
|
429
471
|
try:
|
|
430
|
-
response = self.session.post(self.base_url + "/map", data=json.dumps(data), timeout=timeout)
|
|
472
|
+
response = self.session.post(self.base_url + "/map", data=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {}))
|
|
431
473
|
except requests.exceptions.Timeout:
|
|
432
474
|
raise TimeoutError(timeout)
|
|
433
475
|
|
|
@@ -588,6 +630,7 @@ class TavilyClient:
|
|
|
588
630
|
|
|
589
631
|
data = {k: v for k, v in data.items() if v is not None}
|
|
590
632
|
|
|
633
|
+
override_headers = self._pop_request_headers(kwargs)
|
|
591
634
|
if kwargs:
|
|
592
635
|
data.update(kwargs)
|
|
593
636
|
|
|
@@ -597,7 +640,8 @@ class TavilyClient:
|
|
|
597
640
|
self.base_url + "/research",
|
|
598
641
|
data=json.dumps(data),
|
|
599
642
|
timeout=timeout,
|
|
600
|
-
stream=True
|
|
643
|
+
stream=True,
|
|
644
|
+
**({"headers": override_headers} if override_headers else {})
|
|
601
645
|
)
|
|
602
646
|
except requests.exceptions.Timeout:
|
|
603
647
|
raise TimeoutError(timeout)
|
|
@@ -634,7 +678,8 @@ class TavilyClient:
|
|
|
634
678
|
response = self.session.post(
|
|
635
679
|
self.base_url + "/research",
|
|
636
680
|
data=json.dumps(data),
|
|
637
|
-
timeout=timeout
|
|
681
|
+
timeout=timeout,
|
|
682
|
+
**({"headers": override_headers} if override_headers else {})
|
|
638
683
|
)
|
|
639
684
|
except requests.exceptions.Timeout:
|
|
640
685
|
raise TimeoutError(timeout)
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: tavily-python
|
|
3
|
-
Version: 0.7.
|
|
3
|
+
Version: 0.7.24
|
|
4
4
|
Summary: Python wrapper for the Tavily API
|
|
5
5
|
Home-page: https://github.com/tavily-ai/tavily-python
|
|
6
6
|
Author: Tavily AI
|
|
@@ -25,7 +25,7 @@ Dynamic: requires-dist
|
|
|
25
25
|
Dynamic: requires-python
|
|
26
26
|
Dynamic: summary
|
|
27
27
|
|
|
28
|
-
# Tavily Python
|
|
28
|
+
# Tavily Python SDK
|
|
29
29
|
|
|
30
30
|
[](https://github.com/tavily-ai/tavily-python/stargazers)
|
|
31
31
|
[](https://pypi.org/project/tavily-python/)
|
|
@@ -273,6 +273,76 @@ for chunk in stream:
|
|
|
273
273
|
print(chunk.decode('utf-8'))
|
|
274
274
|
```
|
|
275
275
|
|
|
276
|
+
## Advanced: Custom Session / Client Injection
|
|
277
|
+
|
|
278
|
+
For enterprise environments that proxy Tavily traffic through an API gateway (e.g., for centralized auth, logging, or policy enforcement), you can pass a pre-configured HTTP session instead of a Tavily API key.
|
|
279
|
+
|
|
280
|
+
### Sync (custom `requests.Session`)
|
|
281
|
+
|
|
282
|
+
```python
|
|
283
|
+
import requests
|
|
284
|
+
from tavily import TavilyClient
|
|
285
|
+
|
|
286
|
+
# Pre-configure a session with your gateway's auth
|
|
287
|
+
session = requests.Session()
|
|
288
|
+
session.headers["Authorization"] = "Bearer your-gateway-token"
|
|
289
|
+
session.headers["X-Subscription-Key"] = "your-subscription-key"
|
|
290
|
+
|
|
291
|
+
# No Tavily API key needed — auth is handled by the session
|
|
292
|
+
client = TavilyClient(
|
|
293
|
+
session=session,
|
|
294
|
+
api_base_url="https://your-gateway.com/tavily",
|
|
295
|
+
)
|
|
296
|
+
|
|
297
|
+
response = client.search("latest AI research")
|
|
298
|
+
```
|
|
299
|
+
|
|
300
|
+
### Async (custom `httpx.AsyncClient`)
|
|
301
|
+
|
|
302
|
+
```python
|
|
303
|
+
import httpx
|
|
304
|
+
from tavily import AsyncTavilyClient
|
|
305
|
+
|
|
306
|
+
# Pre-configure an async client with your gateway's auth
|
|
307
|
+
custom_client = httpx.AsyncClient(
|
|
308
|
+
headers={"Authorization": "Bearer your-gateway-token"},
|
|
309
|
+
base_url="https://your-gateway.com/tavily",
|
|
310
|
+
)
|
|
311
|
+
|
|
312
|
+
client = AsyncTavilyClient(client=custom_client)
|
|
313
|
+
|
|
314
|
+
response = await client.search("latest AI research")
|
|
315
|
+
```
|
|
316
|
+
|
|
317
|
+
**Key behaviors:**
|
|
318
|
+
- If a custom session/client is provided, `api_key` is optional
|
|
319
|
+
- Custom session headers take precedence over SDK defaults (e.g., your `Authorization` won't be overwritten)
|
|
320
|
+
- Custom session proxies take precedence over SDK proxy settings
|
|
321
|
+
- The SDK will **not** close externally-provided sessions — you manage the lifecycle
|
|
322
|
+
|
|
323
|
+
## Session & User Tracking
|
|
324
|
+
|
|
325
|
+
`session_id`, `human_id`, and `client_name` are optional identifiers that help attribute requests to a logical session, an end user, and a named client. All three are sent as HTTP headers (`X-Session-Id`, `X-Human-Id`, `X-Client-Name`) and are never persisted in raw form — `human_id` is hashed server-side.
|
|
326
|
+
|
|
327
|
+
Set them once at client init, or per-call (per-call wins):
|
|
328
|
+
|
|
329
|
+
```python
|
|
330
|
+
from tavily import TavilyClient
|
|
331
|
+
|
|
332
|
+
# Client-level — applied to every request
|
|
333
|
+
client = TavilyClient(
|
|
334
|
+
api_key="tvly-YOUR_API_KEY",
|
|
335
|
+
session_id="my-session-123",
|
|
336
|
+
human_id="internal-user-id-42",
|
|
337
|
+
client_name="my-app",
|
|
338
|
+
)
|
|
339
|
+
|
|
340
|
+
# Per-call override
|
|
341
|
+
client.search("hello", session_id="ad-hoc-session")
|
|
342
|
+
```
|
|
343
|
+
|
|
344
|
+
All three are opt-in. Leave them unset and the SDK sends nothing — behavior is identical to earlier versions.
|
|
345
|
+
|
|
276
346
|
## Documentation
|
|
277
347
|
|
|
278
348
|
For a complete guide on how to use the different endpoints and their parameters, please head to our [Python API Reference](https://docs.tavily.com/sdk/python/reference).
|
|
@@ -0,0 +1,410 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import httpx
|
|
3
|
+
from tests.request_intercept import intercept_requests, clear_interceptor, MockSession
|
|
4
|
+
import tavily.tavily as sync_tavily
|
|
5
|
+
import tavily.async_tavily as async_tavily
|
|
6
|
+
import pytest
|
|
7
|
+
from tavily.errors import MissingAPIKeyError
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@pytest.fixture
|
|
11
|
+
def sync_interceptor():
|
|
12
|
+
yield intercept_requests(sync_tavily)
|
|
13
|
+
clear_interceptor(sync_tavily)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@pytest.fixture
|
|
17
|
+
def async_interceptor():
|
|
18
|
+
yield intercept_requests(async_tavily)
|
|
19
|
+
clear_interceptor(async_tavily)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
# --- Sync TavilyClient tests ---
|
|
23
|
+
|
|
24
|
+
class TestSyncCustomSession:
|
|
25
|
+
def test_default_session_created_when_none_provided(self, sync_interceptor):
|
|
26
|
+
client = sync_tavily.TavilyClient(api_key="tvly-test")
|
|
27
|
+
assert not client._external_session
|
|
28
|
+
|
|
29
|
+
def test_custom_session_used(self, sync_interceptor):
|
|
30
|
+
custom_session = MockSession(sync_interceptor)
|
|
31
|
+
client = sync_tavily.TavilyClient(api_key="tvly-test", session=custom_session)
|
|
32
|
+
assert client._external_session
|
|
33
|
+
assert client.session is custom_session
|
|
34
|
+
|
|
35
|
+
def test_custom_session_preserves_existing_headers(self, sync_interceptor):
|
|
36
|
+
custom_session = MockSession(sync_interceptor)
|
|
37
|
+
custom_session.headers["Authorization"] = "Bearer apim-token-123"
|
|
38
|
+
custom_session.headers["X-Custom"] = "custom-value"
|
|
39
|
+
|
|
40
|
+
client = sync_tavily.TavilyClient(api_key="tvly-test", session=custom_session)
|
|
41
|
+
|
|
42
|
+
# Custom Authorization should be preserved (not overwritten by Tavily's)
|
|
43
|
+
assert client.session.headers["Authorization"] == "Bearer apim-token-123"
|
|
44
|
+
# Custom header should be preserved
|
|
45
|
+
assert client.session.headers["X-Custom"] == "custom-value"
|
|
46
|
+
# Tavily defaults should fill in missing headers
|
|
47
|
+
assert client.session.headers["Content-Type"] == "application/json"
|
|
48
|
+
assert client.session.headers["X-Client-Source"] == "tavily-python"
|
|
49
|
+
|
|
50
|
+
def test_custom_session_gets_default_headers_when_empty(self, sync_interceptor):
|
|
51
|
+
custom_session = MockSession(sync_interceptor)
|
|
52
|
+
client = sync_tavily.TavilyClient(api_key="tvly-test", session=custom_session)
|
|
53
|
+
|
|
54
|
+
assert client.session.headers["Authorization"] == "Bearer tvly-test"
|
|
55
|
+
assert client.session.headers["Content-Type"] == "application/json"
|
|
56
|
+
|
|
57
|
+
def test_custom_session_preserves_existing_proxies(self, sync_interceptor):
|
|
58
|
+
custom_session = MockSession(sync_interceptor)
|
|
59
|
+
custom_session.proxies["https"] = "http://my-proxy:8080"
|
|
60
|
+
|
|
61
|
+
client = sync_tavily.TavilyClient(
|
|
62
|
+
api_key="tvly-test",
|
|
63
|
+
session=custom_session,
|
|
64
|
+
proxies={"https": "http://tavily-proxy:9090"},
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
# Custom session proxy should take precedence
|
|
68
|
+
assert client.session.proxies["https"] == "http://my-proxy:8080"
|
|
69
|
+
|
|
70
|
+
def test_close_does_not_close_external_session(self, sync_interceptor):
|
|
71
|
+
closed = []
|
|
72
|
+
custom_session = MockSession(sync_interceptor)
|
|
73
|
+
custom_session.close = lambda: closed.append(True)
|
|
74
|
+
|
|
75
|
+
client = sync_tavily.TavilyClient(api_key="tvly-test", session=custom_session)
|
|
76
|
+
client.close()
|
|
77
|
+
assert len(closed) == 0
|
|
78
|
+
|
|
79
|
+
def test_close_closes_internal_session(self, sync_interceptor):
|
|
80
|
+
client = sync_tavily.TavilyClient(api_key="tvly-test")
|
|
81
|
+
# Should not raise — just verifies close() is called on internal session
|
|
82
|
+
client.close()
|
|
83
|
+
|
|
84
|
+
def test_context_manager_does_not_close_external_session(self, sync_interceptor):
|
|
85
|
+
closed = []
|
|
86
|
+
custom_session = MockSession(sync_interceptor)
|
|
87
|
+
custom_session.close = lambda: closed.append(True)
|
|
88
|
+
|
|
89
|
+
with sync_tavily.TavilyClient(api_key="tvly-test", session=custom_session):
|
|
90
|
+
pass
|
|
91
|
+
assert len(closed) == 0
|
|
92
|
+
|
|
93
|
+
def test_custom_session_sends_request(self, sync_interceptor):
|
|
94
|
+
sync_interceptor.set_response(200, json={"results": []})
|
|
95
|
+
custom_session = MockSession(sync_interceptor)
|
|
96
|
+
custom_session.headers["Authorization"] = "Bearer apim-token"
|
|
97
|
+
|
|
98
|
+
client = sync_tavily.TavilyClient(api_key="tvly-test", session=custom_session)
|
|
99
|
+
client.search("test query")
|
|
100
|
+
|
|
101
|
+
req = sync_interceptor.get_request()
|
|
102
|
+
assert req is not None
|
|
103
|
+
assert req.headers["Authorization"] == "Bearer apim-token"
|
|
104
|
+
|
|
105
|
+
# --- API key validation edge cases ---
|
|
106
|
+
|
|
107
|
+
def test_no_api_key_no_session_raises(self, monkeypatch):
|
|
108
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
109
|
+
with pytest.raises(MissingAPIKeyError):
|
|
110
|
+
sync_tavily.TavilyClient()
|
|
111
|
+
|
|
112
|
+
def test_no_api_key_with_session_allowed(self, sync_interceptor, monkeypatch):
|
|
113
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
114
|
+
custom_session = MockSession(sync_interceptor)
|
|
115
|
+
custom_session.headers["Authorization"] = "Bearer apim-token"
|
|
116
|
+
client = sync_tavily.TavilyClient(session=custom_session)
|
|
117
|
+
assert client.api_key is None
|
|
118
|
+
assert "Authorization" not in client.headers
|
|
119
|
+
assert client.session.headers["Authorization"] == "Bearer apim-token"
|
|
120
|
+
|
|
121
|
+
def test_no_api_key_with_session_no_auth_header_on_defaults(self, sync_interceptor, monkeypatch):
|
|
122
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
123
|
+
custom_session = MockSession(sync_interceptor)
|
|
124
|
+
client = sync_tavily.TavilyClient(session=custom_session)
|
|
125
|
+
# No api_key means no Authorization in defaults
|
|
126
|
+
assert "Authorization" not in client.headers
|
|
127
|
+
# Session shouldn't get an Authorization header either
|
|
128
|
+
assert "Authorization" not in client.session.headers
|
|
129
|
+
|
|
130
|
+
def test_no_api_key_with_session_sends_request(self, sync_interceptor, monkeypatch):
|
|
131
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
132
|
+
sync_interceptor.set_response(200, json={"results": []})
|
|
133
|
+
custom_session = MockSession(sync_interceptor)
|
|
134
|
+
custom_session.headers["Authorization"] = "Bearer apim-token"
|
|
135
|
+
|
|
136
|
+
client = sync_tavily.TavilyClient(session=custom_session)
|
|
137
|
+
client.search("test query")
|
|
138
|
+
|
|
139
|
+
req = sync_interceptor.get_request()
|
|
140
|
+
assert req is not None
|
|
141
|
+
assert req.headers["Authorization"] == "Bearer apim-token"
|
|
142
|
+
|
|
143
|
+
def test_empty_string_api_key_no_session_raises(self, monkeypatch):
|
|
144
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
145
|
+
with pytest.raises(MissingAPIKeyError):
|
|
146
|
+
sync_tavily.TavilyClient(api_key="")
|
|
147
|
+
|
|
148
|
+
def test_empty_string_api_key_with_session_allowed(self, sync_interceptor):
|
|
149
|
+
custom_session = MockSession(sync_interceptor)
|
|
150
|
+
client = sync_tavily.TavilyClient(api_key="", session=custom_session)
|
|
151
|
+
assert "Authorization" not in client.headers
|
|
152
|
+
|
|
153
|
+
def test_api_key_and_session_both_provided(self, sync_interceptor):
|
|
154
|
+
custom_session = MockSession(sync_interceptor)
|
|
155
|
+
client = sync_tavily.TavilyClient(api_key="tvly-test", session=custom_session)
|
|
156
|
+
# api_key provided and session has no Authorization, so default fills it in
|
|
157
|
+
assert client.session.headers["Authorization"] == "Bearer tvly-test"
|
|
158
|
+
|
|
159
|
+
def test_custom_session_with_all_endpoints(self, sync_interceptor):
|
|
160
|
+
custom_session = MockSession(sync_interceptor)
|
|
161
|
+
custom_session.headers["Authorization"] = "Bearer apim-token"
|
|
162
|
+
|
|
163
|
+
client = sync_tavily.TavilyClient(session=custom_session)
|
|
164
|
+
|
|
165
|
+
# search
|
|
166
|
+
sync_interceptor.set_response(200, json={"results": []})
|
|
167
|
+
client.search("test")
|
|
168
|
+
assert sync_interceptor.get_request().headers["Authorization"] == "Bearer apim-token"
|
|
169
|
+
|
|
170
|
+
# extract
|
|
171
|
+
sync_interceptor.set_response(200, json={"results": [], "failed_results": []})
|
|
172
|
+
client.extract(urls=["https://example.com"])
|
|
173
|
+
assert sync_interceptor.get_request().headers["Authorization"] == "Bearer apim-token"
|
|
174
|
+
|
|
175
|
+
# crawl
|
|
176
|
+
sync_interceptor.set_response(200, json={"results": []})
|
|
177
|
+
client.crawl(url="https://example.com")
|
|
178
|
+
assert sync_interceptor.get_request().headers["Authorization"] == "Bearer apim-token"
|
|
179
|
+
|
|
180
|
+
# map
|
|
181
|
+
sync_interceptor.set_response(200, json={"results": []})
|
|
182
|
+
client.map(url="https://example.com")
|
|
183
|
+
assert sync_interceptor.get_request().headers["Authorization"] == "Bearer apim-token"
|
|
184
|
+
|
|
185
|
+
def test_custom_session_with_custom_base_url(self, sync_interceptor):
|
|
186
|
+
sync_interceptor.set_response(200, json={"results": []})
|
|
187
|
+
custom_session = MockSession(sync_interceptor)
|
|
188
|
+
|
|
189
|
+
client = sync_tavily.TavilyClient(
|
|
190
|
+
api_key="tvly-test",
|
|
191
|
+
session=custom_session,
|
|
192
|
+
api_base_url="https://apim.corp.com/tavily",
|
|
193
|
+
)
|
|
194
|
+
client.search("test")
|
|
195
|
+
|
|
196
|
+
req = sync_interceptor.get_request()
|
|
197
|
+
assert req.url == "https://apim.corp.com/tavily/search"
|
|
198
|
+
|
|
199
|
+
def test_custom_session_proxies_fill_missing_protocols(self, sync_interceptor):
|
|
200
|
+
custom_session = MockSession(sync_interceptor)
|
|
201
|
+
custom_session.proxies["http"] = "http://session-proxy:8080"
|
|
202
|
+
|
|
203
|
+
client = sync_tavily.TavilyClient(
|
|
204
|
+
api_key="tvly-test",
|
|
205
|
+
session=custom_session,
|
|
206
|
+
proxies={"http": "http://arg-proxy:9090", "https": "http://arg-proxy:9091"},
|
|
207
|
+
)
|
|
208
|
+
|
|
209
|
+
# http: session proxy wins
|
|
210
|
+
assert client.session.proxies["http"] == "http://session-proxy:8080"
|
|
211
|
+
# https: session didn't have it, so arg fills it in
|
|
212
|
+
assert client.session.proxies["https"] == "http://arg-proxy:9091"
|
|
213
|
+
|
|
214
|
+
def test_custom_session_project_id_header(self, sync_interceptor):
|
|
215
|
+
custom_session = MockSession(sync_interceptor)
|
|
216
|
+
client = sync_tavily.TavilyClient(
|
|
217
|
+
api_key="tvly-test",
|
|
218
|
+
session=custom_session,
|
|
219
|
+
project_id="my-project",
|
|
220
|
+
)
|
|
221
|
+
assert client.session.headers["X-Project-ID"] == "my-project"
|
|
222
|
+
|
|
223
|
+
def test_shared_session_across_multiple_clients(self, sync_interceptor):
|
|
224
|
+
sync_interceptor.set_response(200, json={"results": []})
|
|
225
|
+
shared_session = MockSession(sync_interceptor)
|
|
226
|
+
shared_session.headers["Authorization"] = "Bearer shared-token"
|
|
227
|
+
|
|
228
|
+
client1 = sync_tavily.TavilyClient(session=shared_session)
|
|
229
|
+
client2 = sync_tavily.TavilyClient(session=shared_session)
|
|
230
|
+
|
|
231
|
+
assert client1.session is client2.session
|
|
232
|
+
|
|
233
|
+
client1.search("query1")
|
|
234
|
+
assert sync_interceptor.get_request().headers["Authorization"] == "Bearer shared-token"
|
|
235
|
+
|
|
236
|
+
client2.search("query2")
|
|
237
|
+
assert sync_interceptor.get_request().headers["Authorization"] == "Bearer shared-token"
|
|
238
|
+
|
|
239
|
+
# Closing one client should not close the shared session
|
|
240
|
+
client1.close()
|
|
241
|
+
client2.search("query3")
|
|
242
|
+
assert sync_interceptor.get_request() is not None
|
|
243
|
+
|
|
244
|
+
|
|
245
|
+
# --- Async AsyncTavilyClient tests ---
|
|
246
|
+
|
|
247
|
+
class TestAsyncCustomClient:
|
|
248
|
+
def test_default_client_created_when_none_provided(self):
|
|
249
|
+
client = async_tavily.AsyncTavilyClient(api_key="tvly-test")
|
|
250
|
+
assert not client._external_client
|
|
251
|
+
|
|
252
|
+
def test_custom_client_used(self):
|
|
253
|
+
custom_client = httpx.AsyncClient()
|
|
254
|
+
client = async_tavily.AsyncTavilyClient(api_key="tvly-test", client=custom_client)
|
|
255
|
+
assert client._external_client
|
|
256
|
+
assert client._client is custom_client
|
|
257
|
+
|
|
258
|
+
def test_custom_client_preserves_existing_headers(self):
|
|
259
|
+
custom_client = httpx.AsyncClient(headers={
|
|
260
|
+
"Authorization": "Bearer apim-token-123",
|
|
261
|
+
"X-Custom": "custom-value",
|
|
262
|
+
})
|
|
263
|
+
client = async_tavily.AsyncTavilyClient(api_key="tvly-test", client=custom_client)
|
|
264
|
+
|
|
265
|
+
assert client._client.headers["Authorization"] == "Bearer apim-token-123"
|
|
266
|
+
assert client._client.headers["X-Custom"] == "custom-value"
|
|
267
|
+
assert client._client.headers["Content-Type"] == "application/json"
|
|
268
|
+
assert client._client.headers["X-Client-Source"] == "tavily-python"
|
|
269
|
+
|
|
270
|
+
def test_custom_client_gets_default_headers_when_empty(self):
|
|
271
|
+
custom_client = httpx.AsyncClient()
|
|
272
|
+
client = async_tavily.AsyncTavilyClient(api_key="tvly-test", client=custom_client)
|
|
273
|
+
|
|
274
|
+
assert client._client.headers["Authorization"] == "Bearer tvly-test"
|
|
275
|
+
assert client._client.headers["Content-Type"] == "application/json"
|
|
276
|
+
|
|
277
|
+
def test_custom_client_base_url_set_when_missing(self):
|
|
278
|
+
custom_client = httpx.AsyncClient()
|
|
279
|
+
client = async_tavily.AsyncTavilyClient(api_key="tvly-test", client=custom_client)
|
|
280
|
+
assert "api.tavily.com" in str(client._client.base_url)
|
|
281
|
+
|
|
282
|
+
def test_custom_client_base_url_preserved_when_set(self):
|
|
283
|
+
custom_client = httpx.AsyncClient(base_url="https://apim.example.com/tavily")
|
|
284
|
+
client = async_tavily.AsyncTavilyClient(api_key="tvly-test", client=custom_client)
|
|
285
|
+
assert "apim.example.com" in str(client._client.base_url)
|
|
286
|
+
|
|
287
|
+
def test_close_does_not_close_external_client(self):
|
|
288
|
+
closed = []
|
|
289
|
+
custom_client = httpx.AsyncClient()
|
|
290
|
+
|
|
291
|
+
async def run():
|
|
292
|
+
client = async_tavily.AsyncTavilyClient(api_key="tvly-test", client=custom_client)
|
|
293
|
+
original_aclose = custom_client.aclose
|
|
294
|
+
|
|
295
|
+
async def track_close():
|
|
296
|
+
closed.append(True)
|
|
297
|
+
await original_aclose()
|
|
298
|
+
|
|
299
|
+
custom_client.aclose = track_close
|
|
300
|
+
await client.close()
|
|
301
|
+
|
|
302
|
+
asyncio.run(run())
|
|
303
|
+
assert len(closed) == 0
|
|
304
|
+
|
|
305
|
+
def test_context_manager_does_not_close_external_client(self):
|
|
306
|
+
closed = []
|
|
307
|
+
custom_client = httpx.AsyncClient()
|
|
308
|
+
|
|
309
|
+
async def run():
|
|
310
|
+
original_aclose = custom_client.aclose
|
|
311
|
+
|
|
312
|
+
async def track_close():
|
|
313
|
+
closed.append(True)
|
|
314
|
+
await original_aclose()
|
|
315
|
+
|
|
316
|
+
custom_client.aclose = track_close
|
|
317
|
+
async with async_tavily.AsyncTavilyClient(api_key="tvly-test", client=custom_client):
|
|
318
|
+
pass
|
|
319
|
+
|
|
320
|
+
asyncio.run(run())
|
|
321
|
+
assert len(closed) == 0
|
|
322
|
+
|
|
323
|
+
def test_custom_client_sends_request(self, async_interceptor):
|
|
324
|
+
async_interceptor.set_response(200, json={"results": []})
|
|
325
|
+
custom_client = httpx.AsyncClient(
|
|
326
|
+
headers={"Authorization": "Bearer apim-token"},
|
|
327
|
+
base_url="https://api.tavily.com",
|
|
328
|
+
)
|
|
329
|
+
|
|
330
|
+
client = async_tavily.AsyncTavilyClient(api_key="tvly-test", client=custom_client)
|
|
331
|
+
asyncio.run(client.search("test query"))
|
|
332
|
+
|
|
333
|
+
req = async_interceptor.get_request()
|
|
334
|
+
assert req is not None
|
|
335
|
+
assert req.headers["Authorization"] == "Bearer apim-token"
|
|
336
|
+
|
|
337
|
+
# --- API key validation edge cases ---
|
|
338
|
+
|
|
339
|
+
def test_no_api_key_no_client_raises(self, monkeypatch):
|
|
340
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
341
|
+
with pytest.raises(MissingAPIKeyError):
|
|
342
|
+
async_tavily.AsyncTavilyClient()
|
|
343
|
+
|
|
344
|
+
def test_no_api_key_with_client_allowed(self, monkeypatch):
|
|
345
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
346
|
+
custom_client = httpx.AsyncClient(
|
|
347
|
+
headers={"Authorization": "Bearer apim-token"},
|
|
348
|
+
)
|
|
349
|
+
client = async_tavily.AsyncTavilyClient(client=custom_client)
|
|
350
|
+
assert client._client.headers["Authorization"] == "Bearer apim-token"
|
|
351
|
+
|
|
352
|
+
def test_no_api_key_with_client_no_auth_header_on_defaults(self, monkeypatch):
|
|
353
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
354
|
+
custom_client = httpx.AsyncClient()
|
|
355
|
+
client = async_tavily.AsyncTavilyClient(client=custom_client)
|
|
356
|
+
# httpx always has headers dict but Authorization shouldn't be added
|
|
357
|
+
assert "authorization" not in [k.lower() for k in client._client.headers.keys()
|
|
358
|
+
if k.lower() == "authorization"
|
|
359
|
+
and client._client.headers[k].startswith("Bearer None")]
|
|
360
|
+
|
|
361
|
+
def test_no_api_key_with_client_sends_request(self, async_interceptor, monkeypatch):
|
|
362
|
+
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
|
|
363
|
+
async_interceptor.set_response(200, json={"results": []})
|
|
364
|
+
custom_client = httpx.AsyncClient(
|
|
365
|
+
headers={"Authorization": "Bearer apim-token"},
|
|
366
|
+
base_url="https://api.tavily.com",
|
|
367
|
+
)
|
|
368
|
+
|
|
369
|
+
client = async_tavily.AsyncTavilyClient(client=custom_client)
|
|
370
|
+
asyncio.run(client.search("test query"))
|
|
371
|
+
|
|
372
|
+
req = async_interceptor.get_request()
|
|
373
|
+
assert req is not None
|
|
374
|
+
assert req.headers["Authorization"] == "Bearer apim-token"
|
|
375
|
+
|
|
376
|
+
def test_custom_client_with_all_endpoints(self, async_interceptor):
|
|
377
|
+
custom_client = httpx.AsyncClient(
|
|
378
|
+
headers={"Authorization": "Bearer apim-token"},
|
|
379
|
+
base_url="https://api.tavily.com",
|
|
380
|
+
)
|
|
381
|
+
client = async_tavily.AsyncTavilyClient(client=custom_client)
|
|
382
|
+
|
|
383
|
+
# search
|
|
384
|
+
async_interceptor.set_response(200, json={"results": []})
|
|
385
|
+
asyncio.run(client.search("test"))
|
|
386
|
+
assert async_interceptor.get_request().headers["Authorization"] == "Bearer apim-token"
|
|
387
|
+
|
|
388
|
+
# extract
|
|
389
|
+
async_interceptor.set_response(200, json={"results": [], "failed_results": []})
|
|
390
|
+
asyncio.run(client.extract(urls=["https://example.com"]))
|
|
391
|
+
assert async_interceptor.get_request().headers["Authorization"] == "Bearer apim-token"
|
|
392
|
+
|
|
393
|
+
# crawl
|
|
394
|
+
async_interceptor.set_response(200, json={"results": []})
|
|
395
|
+
asyncio.run(client.crawl(url="https://example.com"))
|
|
396
|
+
assert async_interceptor.get_request().headers["Authorization"] == "Bearer apim-token"
|
|
397
|
+
|
|
398
|
+
# map
|
|
399
|
+
async_interceptor.set_response(200, json={"results": []})
|
|
400
|
+
asyncio.run(client.map(url="https://example.com"))
|
|
401
|
+
assert async_interceptor.get_request().headers["Authorization"] == "Bearer apim-token"
|
|
402
|
+
|
|
403
|
+
def test_custom_client_project_id_header(self):
|
|
404
|
+
custom_client = httpx.AsyncClient()
|
|
405
|
+
client = async_tavily.AsyncTavilyClient(
|
|
406
|
+
api_key="tvly-test",
|
|
407
|
+
client=custom_client,
|
|
408
|
+
project_id="my-project",
|
|
409
|
+
)
|
|
410
|
+
assert client._client.headers["X-Project-ID"] == "my-project"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|