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.
Files changed (25) hide show
  1. {tavily_python-0.7.22 → tavily_python-0.7.24}/PKG-INFO +72 -2
  2. {tavily_python-0.7.22 → tavily_python-0.7.24}/README.md +71 -1
  3. {tavily_python-0.7.22 → tavily_python-0.7.24}/setup.py +1 -1
  4. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/async_tavily.py +77 -33
  5. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/hybrid_rag/hybrid_rag.py +6 -3
  6. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/tavily.py +63 -18
  7. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/PKG-INFO +72 -2
  8. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/SOURCES.txt +1 -0
  9. tavily_python-0.7.24/tests/test_custom_session.py +410 -0
  10. {tavily_python-0.7.22 → tavily_python-0.7.24}/LICENSE +0 -0
  11. {tavily_python-0.7.22 → tavily_python-0.7.24}/setup.cfg +0 -0
  12. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/__init__.py +0 -0
  13. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/config.py +0 -0
  14. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/errors.py +0 -0
  15. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/hybrid_rag/__init__.py +0 -0
  16. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily/utils.py +0 -0
  17. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/dependency_links.txt +0 -0
  18. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/requires.txt +0 -0
  19. {tavily_python-0.7.22 → tavily_python-0.7.24}/tavily_python.egg-info/top_level.txt +0 -0
  20. {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_crawl.py +0 -0
  21. {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_errors.py +0 -0
  22. {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_map.py +0 -0
  23. {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_research.py +0 -0
  24. {tavily_python-0.7.22 → tavily_python-0.7.24}/tests/test_search.py +0 -0
  25. {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.22
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 Wrapper
28
+ # Tavily Python SDK
29
29
 
30
30
  [![GitHub stars](https://img.shields.io/github/stars/tavily-ai/tavily-python?style=social)](https://github.com/tavily-ai/tavily-python/stargazers)
31
31
  [![PyPI - Downloads](https://img.shields.io/pypi/dm/tavily-python)](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 Wrapper
1
+ # Tavily Python SDK
2
2
 
3
3
  [![GitHub stars](https://img.shields.io/github/stars/tavily-ai/tavily-python?style=social)](https://github.com/tavily-ai/tavily-python/stargazers)
4
4
  [![PyPI - Downloads](https://img.shields.io/pypi/dm/tavily-python)](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).
@@ -5,7 +5,7 @@ with open('README.md', 'r', encoding='utf-8') as f:
5
5
 
6
6
  setup(
7
7
  name='tavily-python',
8
- version='0.7.22',
8
+ version='0.7.24',
9
9
  url='https://github.com/tavily-ai/tavily-python',
10
10
  author='Tavily AI',
11
11
  author_email='support@tavily.com',
@@ -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
- proxies = proxies or {}
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
- mapped_proxies = {
32
- "http://": proxies.get("http", os.getenv("TAVILY_HTTP_PROXY")),
33
- "https://": proxies.get("https", os.getenv("TAVILY_HTTPS_PROXY")),
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
- mapped_proxies = {key: value for key, value in mapped_proxies.items() if value}
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
- proxy_mounts = (
39
- {scheme: httpx.AsyncHTTPTransport(proxy=proxy) for scheme, proxy in mapped_proxies.items()}
40
- if mapped_proxies
41
- else None
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
- tavily_project = project_id or os.getenv("TAVILY_PROJECT")
67
+ mapped_proxies = {key: value for key, value in mapped_proxies.items() if value}
45
68
 
46
- self._api_base_url = api_base_url or "https://api.tavily.com"
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
- # Create a persistent client for connection pooling
49
- self._client = httpx.AsyncClient(
50
- headers={
51
- "Content-Type": "application/json",
52
- "Authorization": f"Bearer {api_key}",
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
- await self._client.aclose()
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__(self, api_key: Optional[str] = None, proxies: Optional[dict[str, str]] = None, api_base_url: Optional[str] = None, client_source: Optional[str] = None, project_id: Optional[str] = None):
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.session = requests.Session()
41
- self.session.headers.update(self.headers)
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.session.proxies.update(self.proxies)
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.session.close()
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.22
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 Wrapper
28
+ # Tavily Python SDK
29
29
 
30
30
  [![GitHub stars](https://img.shields.io/github/stars/tavily-ai/tavily-python?style=social)](https://github.com/tavily-ai/tavily-python/stargazers)
31
31
  [![PyPI - Downloads](https://img.shields.io/pypi/dm/tavily-python)](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).
@@ -15,6 +15,7 @@ tavily_python.egg-info/dependency_links.txt
15
15
  tavily_python.egg-info/requires.txt
16
16
  tavily_python.egg-info/top_level.txt
17
17
  tests/test_crawl.py
18
+ tests/test_custom_session.py
18
19
  tests/test_errors.py
19
20
  tests/test_map.py
20
21
  tests/test_research.py
@@ -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