tavily-python 0.7.3__tar.gz → 0.7.5__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.3 → tavily_python-0.7.5}/PKG-INFO +6 -1
  2. {tavily_python-0.7.3 → tavily_python-0.7.5}/README.md +5 -0
  3. tavily_python-0.7.5/setup.cfg +4 -0
  4. {tavily_python-0.7.3 → tavily_python-0.7.5}/setup.py +2 -2
  5. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily/async_tavily.py +22 -4
  6. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily/config.py +0 -1
  7. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily/tavily.py +22 -4
  8. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily_python.egg-info/PKG-INFO +6 -1
  9. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily_python.egg-info/SOURCES.txt +4 -3
  10. tavily_python-0.7.5/tests/test_crawl.py +102 -0
  11. tavily_python-0.7.5/tests/test_errors.py +50 -0
  12. tavily_python-0.7.5/tests/test_map.py +98 -0
  13. tavily_python-0.7.5/tests/test_search.py +99 -0
  14. tavily_python-0.7.3/setup.cfg +0 -8
  15. tavily_python-0.7.3/tests/test_async_search.py +0 -219
  16. tavily_python-0.7.3/tests/test_sync_search.py +0 -219
  17. {tavily_python-0.7.3 → tavily_python-0.7.5}/LICENSE +0 -0
  18. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily/__init__.py +0 -0
  19. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily/errors.py +0 -0
  20. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily/hybrid_rag/__init__.py +0 -0
  21. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily/hybrid_rag/hybrid_rag.py +0 -0
  22. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily/utils.py +0 -0
  23. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily_python.egg-info/dependency_links.txt +0 -0
  24. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily_python.egg-info/requires.txt +0 -0
  25. {tavily_python-0.7.3 → tavily_python-0.7.5}/tavily_python.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tavily-python
3
- Version: 0.7.3
3
+ Version: 0.7.5
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
@@ -27,6 +27,11 @@ Dynamic: summary
27
27
 
28
28
  # Tavily Python Wrapper
29
29
 
30
+ [![GitHub stars](https://img.shields.io/github/stars/tavily-ai/tavily-python?style=social)](https://github.com/tavily-ai/tavily-python/stargazers)
31
+ [![PyPI - Downloads](https://img.shields.io/pypi/dm/tavily-python)](https://pypi.org/project/tavily-python/)
32
+ [![License](https://img.shields.io/github/license/tavily-ai/tavily-python)](https://github.com/tavily-ai/tavily-python/blob/main/LICENSE)
33
+ [![CI](https://github.com/tavily-ai/tavily-python/actions/workflows/tests.yml/badge.svg)](https://github.com/tavily-ai/tavily-python/actions)
34
+
30
35
  The Tavily Python wrapper allows for easy interaction with the Tavily API, offering the full range of our search and extract functionalities directly from your Python programs. Easily integrate smart search and content extraction capabilities into your applications, harnessing Tavily's powerful search and extract features.
31
36
 
32
37
  ## Installing
@@ -1,5 +1,10 @@
1
1
  # Tavily Python Wrapper
2
2
 
3
+ [![GitHub stars](https://img.shields.io/github/stars/tavily-ai/tavily-python?style=social)](https://github.com/tavily-ai/tavily-python/stargazers)
4
+ [![PyPI - Downloads](https://img.shields.io/pypi/dm/tavily-python)](https://pypi.org/project/tavily-python/)
5
+ [![License](https://img.shields.io/github/license/tavily-ai/tavily-python)](https://github.com/tavily-ai/tavily-python/blob/main/LICENSE)
6
+ [![CI](https://github.com/tavily-ai/tavily-python/actions/workflows/tests.yml/badge.svg)](https://github.com/tavily-ai/tavily-python/actions)
7
+
3
8
  The Tavily Python wrapper allows for easy interaction with the Tavily API, offering the full range of our search and extract functionalities directly from your Python programs. Easily integrate smart search and content extraction capabilities into your applications, harnessing Tavily's powerful search and extract features.
4
9
 
5
10
  ## Installing
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -5,14 +5,14 @@ with open('README.md', 'r', encoding='utf-8') as f:
5
5
 
6
6
  setup(
7
7
  name='tavily-python',
8
- version='0.7.3',
8
+ version='0.7.5',
9
9
  url='https://github.com/tavily-ai/tavily-python',
10
10
  author='Tavily AI',
11
11
  author_email='support@tavily.com',
12
12
  description='Python wrapper for the Tavily API',
13
13
  long_description=long_description,
14
14
  long_description_content_type='text/markdown',
15
- packages=find_packages(),
15
+ packages=find_packages(exclude=['tests']),
16
16
  install_requires=['requests', 'tiktoken>=0.5.1', 'httpx'],
17
17
  classifiers=[
18
18
  'Programming Language :: Python :: 3',
@@ -60,9 +60,10 @@ class AsyncTavilyClient:
60
60
  include_domains: Sequence[str] = None,
61
61
  exclude_domains: Sequence[str] = None,
62
62
  include_answer: Union[bool, Literal["basic", "advanced"]] = False,
63
- include_raw_content: bool = False,
63
+ include_raw_content: Union[bool, Literal["markdown", "text"]] = False,
64
64
  include_images: bool = False,
65
65
  timeout: int = 60,
66
+ country: str = None,
66
67
  **kwargs,
67
68
  ) -> dict:
68
69
  """
@@ -72,7 +73,7 @@ class AsyncTavilyClient:
72
73
  "query": query,
73
74
  "search_depth": search_depth,
74
75
  "topic": topic,
75
- "time_range'": time_range,
76
+ "time_range": time_range,
76
77
  "days": days,
77
78
  "include_answer": include_answer,
78
79
  "include_raw_content": include_raw_content,
@@ -80,6 +81,7 @@ class AsyncTavilyClient:
80
81
  "include_domains": include_domains,
81
82
  "exclude_domains": exclude_domains,
82
83
  "include_images": include_images,
84
+ "country": country,
83
85
  }
84
86
 
85
87
  if kwargs:
@@ -123,9 +125,10 @@ class AsyncTavilyClient:
123
125
  include_domains: Sequence[str] = None,
124
126
  exclude_domains: Sequence[str] = None,
125
127
  include_answer: Union[bool, Literal["basic", "advanced"]] = False,
126
- include_raw_content: bool = False,
128
+ include_raw_content: Union[bool, Literal["markdown", "text"]] = False,
127
129
  include_images: bool = False,
128
130
  timeout: int = 60,
131
+ country: str = None,
129
132
  **kwargs, # Accept custom arguments
130
133
  ) -> dict:
131
134
  """
@@ -144,6 +147,7 @@ class AsyncTavilyClient:
144
147
  include_raw_content=include_raw_content,
145
148
  include_images=include_images,
146
149
  timeout=timeout,
150
+ country=country,
147
151
  **kwargs,
148
152
  )
149
153
 
@@ -158,6 +162,7 @@ class AsyncTavilyClient:
158
162
  urls: Union[List[str], str],
159
163
  include_images: bool = False,
160
164
  extract_depth: Literal["basic", "advanced"] = "basic",
165
+ format: Literal["markdown", "text"] = "markdown",
161
166
  timeout: int = 60,
162
167
  **kwargs
163
168
  ) -> dict:
@@ -168,6 +173,7 @@ class AsyncTavilyClient:
168
173
  "urls": urls,
169
174
  "include_images": include_images,
170
175
  "extract_depth": extract_depth,
176
+ "format": format,
171
177
  }
172
178
  if kwargs:
173
179
  data.update(kwargs)
@@ -205,6 +211,7 @@ class AsyncTavilyClient:
205
211
  urls: Union[List[str], str], # Accept a list of URLs or a single URL
206
212
  include_images: bool = False,
207
213
  extract_depth: Literal["basic", "advanced"] = "basic",
214
+ format: Literal["markdown", "text"] = "markdown",
208
215
  timeout: int = 60,
209
216
  **kwargs, # Accept custom arguments
210
217
  ) -> dict:
@@ -215,6 +222,7 @@ class AsyncTavilyClient:
215
222
  response_dict = await self._extract(urls,
216
223
  include_images,
217
224
  extract_depth,
225
+ format,
218
226
  timeout,
219
227
  **kwargs,
220
228
  )
@@ -241,6 +249,7 @@ class AsyncTavilyClient:
241
249
  include_images: bool = None,
242
250
  categories: Sequence[AllowedCategory] = None,
243
251
  extract_depth: Literal["basic", "advanced"] = None,
252
+ format: Literal["markdown", "text"] = None,
244
253
  timeout: int = 60,
245
254
  **kwargs
246
255
  ) -> dict:
@@ -261,6 +270,7 @@ class AsyncTavilyClient:
261
270
  "categories": categories,
262
271
  "include_images": include_images,
263
272
  "extract_depth": extract_depth,
273
+ "format": format
264
274
  }
265
275
 
266
276
  if kwargs:
@@ -310,6 +320,7 @@ class AsyncTavilyClient:
310
320
  categories: Sequence[AllowedCategory] = None,
311
321
  extract_depth: Literal["basic", "advanced"] = None,
312
322
  include_images: bool = None,
323
+ format: Literal["markdown", "text"] = None,
313
324
  timeout: int = 60,
314
325
  **kwargs
315
326
  ) -> dict:
@@ -331,6 +342,7 @@ class AsyncTavilyClient:
331
342
  categories=categories,
332
343
  extract_depth=extract_depth,
333
344
  include_images=include_images,
345
+ format=format,
334
346
  timeout=timeout,
335
347
  **kwargs)
336
348
 
@@ -451,6 +463,7 @@ class AsyncTavilyClient:
451
463
  exclude_domains: Sequence[str] = None,
452
464
  max_tokens: int = 4000,
453
465
  timeout: int = 60,
466
+ country: str = None,
454
467
  **kwargs, # Accept custom arguments
455
468
  ) -> str:
456
469
  """
@@ -473,6 +486,7 @@ class AsyncTavilyClient:
473
486
  include_raw_content=False,
474
487
  include_images=False,
475
488
  timeout = timeout,
489
+ country=country,
476
490
  **kwargs,
477
491
  )
478
492
  sources = response_dict.get("results", [])
@@ -488,6 +502,7 @@ class AsyncTavilyClient:
488
502
  include_domains: Sequence[str] = None,
489
503
  exclude_domains: Sequence[str] = None,
490
504
  timeout: int = 60,
505
+ country: str = None,
491
506
  **kwargs, # Accept custom arguments
492
507
  ) -> str:
493
508
  """
@@ -505,6 +520,7 @@ class AsyncTavilyClient:
505
520
  include_images=False,
506
521
  include_answer=True,
507
522
  timeout = timeout,
523
+ country=country,
508
524
  **kwargs,
509
525
  )
510
526
  return response_dict.get("answer", "")
@@ -514,6 +530,7 @@ class AsyncTavilyClient:
514
530
  search_depth: Literal["basic", "advanced"] = "advanced",
515
531
  max_results: int = 5,
516
532
  timeout: int = 60,
533
+ country: str = None,
517
534
  ) -> Sequence[dict]:
518
535
  """ Company information search method. Search depth is advanced by default to get the best answer. """
519
536
  timeout = min(timeout, 120)
@@ -524,7 +541,8 @@ class AsyncTavilyClient:
524
541
  topic=topic,
525
542
  max_results=max_results,
526
543
  include_answer=False,
527
- timeout = timeout)
544
+ timeout = timeout,
545
+ country=country)
528
546
 
529
547
  all_results = []
530
548
  for data in await asyncio.gather(*[_perform_search(topic) for topic in self._company_info_tags]):
@@ -10,4 +10,3 @@ AllowedCategory = Literal[
10
10
  "E-Commerce", "Authentication", "Developer", "Developers", "Solutions",
11
11
  "Partners", "Downloads", "Media", "Events", "People"
12
12
  ]
13
-
@@ -45,9 +45,10 @@ class TavilyClient:
45
45
  include_domains: Sequence[str] = None,
46
46
  exclude_domains: Sequence[str] = None,
47
47
  include_answer: Union[bool, Literal["basic", "advanced"]] = False,
48
- include_raw_content: bool = False,
48
+ include_raw_content: Union[bool, Literal["markdown", "text"]] = False,
49
49
  include_images: bool = False,
50
50
  timeout: int = 60,
51
+ country: str = None,
51
52
  **kwargs
52
53
  ) -> dict:
53
54
  """
@@ -66,6 +67,7 @@ class TavilyClient:
66
67
  "include_domains": include_domains,
67
68
  "exclude_domains": exclude_domains,
68
69
  "include_images": include_images,
70
+ "country": country,
69
71
  }
70
72
 
71
73
  if kwargs:
@@ -110,9 +112,10 @@ class TavilyClient:
110
112
  include_domains: Sequence[str] = None,
111
113
  exclude_domains: Sequence[str] = None,
112
114
  include_answer: Union[bool, Literal["basic", "advanced"]] = False,
113
- include_raw_content: bool = False,
115
+ include_raw_content: Union[bool, Literal["markdown", "text"]] = False,
114
116
  include_images: bool = False,
115
117
  timeout: int = 60,
118
+ country: str = None,
116
119
  **kwargs, # Accept custom arguments
117
120
  ) -> dict:
118
121
  """
@@ -131,6 +134,7 @@ class TavilyClient:
131
134
  include_raw_content=include_raw_content,
132
135
  include_images=include_images,
133
136
  timeout=timeout,
137
+ country=country,
134
138
  **kwargs,
135
139
  )
136
140
 
@@ -144,6 +148,7 @@ class TavilyClient:
144
148
  urls: Union[List[str], str],
145
149
  include_images: bool = False,
146
150
  extract_depth: Literal["basic", "advanced"] = "basic",
151
+ format: Literal["markdown", "text"] = "markdown",
147
152
  timeout: int = 60,
148
153
  **kwargs
149
154
  ) -> dict:
@@ -154,6 +159,7 @@ class TavilyClient:
154
159
  "urls": urls,
155
160
  "include_images": include_images,
156
161
  "extract_depth": extract_depth,
162
+ "format": format,
157
163
  }
158
164
  if kwargs:
159
165
  data.update(kwargs)
@@ -189,6 +195,7 @@ class TavilyClient:
189
195
  urls: Union[List[str], str], # Accept a list of URLs or a single URL
190
196
  include_images: bool = False,
191
197
  extract_depth: Literal["basic", "advanced"] = "basic",
198
+ format: Literal["markdown", "text"] = "markdown",
192
199
  timeout: int = 60,
193
200
  **kwargs, # Accept custom arguments
194
201
  ) -> dict:
@@ -199,6 +206,7 @@ class TavilyClient:
199
206
  response_dict = self._extract(urls,
200
207
  include_images,
201
208
  extract_depth,
209
+ format,
202
210
  timeout,
203
211
  **kwargs)
204
212
 
@@ -224,6 +232,7 @@ class TavilyClient:
224
232
  include_images: bool = None,
225
233
  categories: Sequence[AllowedCategory] = None,
226
234
  extract_depth: Literal["basic", "advanced"] = None,
235
+ format: Literal["markdown", "text"] = None,
227
236
  timeout: int = 60,
228
237
  **kwargs
229
238
  ) -> dict:
@@ -244,6 +253,7 @@ class TavilyClient:
244
253
  "include_images": include_images,
245
254
  "categories": categories,
246
255
  "extract_depth": extract_depth,
256
+ "format": format
247
257
  }
248
258
 
249
259
  if kwargs:
@@ -293,6 +303,7 @@ class TavilyClient:
293
303
  include_images: bool = None,
294
304
  categories: Sequence[AllowedCategory] = None,
295
305
  extract_depth: Literal["basic", "advanced"] = None,
306
+ format: Literal["markdown", "text"] = None,
296
307
  timeout: int = 60,
297
308
  **kwargs
298
309
  ) -> dict:
@@ -311,9 +322,10 @@ class TavilyClient:
311
322
  exclude_paths=exclude_paths,
312
323
  exclude_domains=exclude_domains,
313
324
  allow_external=allow_external,
325
+ include_images=include_images,
314
326
  categories=categories,
315
327
  extract_depth=extract_depth,
316
- include_images=include_images,
328
+ format=format,
317
329
  timeout=timeout,
318
330
  **kwargs)
319
331
 
@@ -434,6 +446,7 @@ class TavilyClient:
434
446
  exclude_domains: Sequence[str] = None,
435
447
  max_tokens: int = 4000,
436
448
  timeout: int = 60,
449
+ country: str = None,
437
450
  **kwargs, # Accept custom arguments
438
451
  ) -> str:
439
452
  """
@@ -456,6 +469,7 @@ class TavilyClient:
456
469
  include_raw_content=False,
457
470
  include_images=False,
458
471
  timeout=timeout,
472
+ country=country,
459
473
  **kwargs,
460
474
  )
461
475
  sources = response_dict.get("results", [])
@@ -472,6 +486,7 @@ class TavilyClient:
472
486
  include_domains: Sequence[str] = None,
473
487
  exclude_domains: Sequence[str] = None,
474
488
  timeout: int = 60,
489
+ country: str = None,
475
490
  **kwargs, # Accept custom arguments
476
491
  ) -> str:
477
492
  """
@@ -489,6 +504,7 @@ class TavilyClient:
489
504
  include_images=False,
490
505
  include_answer=True,
491
506
  timeout=timeout,
507
+ country=country,
492
508
  **kwargs,
493
509
  )
494
510
  return response_dict.get("answer", "")
@@ -499,6 +515,7 @@ class TavilyClient:
499
515
  "advanced"] = "advanced",
500
516
  max_results: int = 5,
501
517
  timeout: int = 60,
518
+ country: str = None,
502
519
  ) -> Sequence[dict]:
503
520
  """ Company information search method. Search depth is advanced by default to get the best answer. """
504
521
  timeout = min(timeout, 120)
@@ -508,7 +525,8 @@ class TavilyClient:
508
525
  topic=topic,
509
526
  max_results=max_results,
510
527
  include_answer=False,
511
- timeout=timeout)
528
+ timeout=timeout,
529
+ country=country)
512
530
 
513
531
  with ThreadPoolExecutor() as executor:
514
532
  # Initiate the search for each topic in parallel
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tavily-python
3
- Version: 0.7.3
3
+ Version: 0.7.5
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
@@ -27,6 +27,11 @@ Dynamic: summary
27
27
 
28
28
  # Tavily Python Wrapper
29
29
 
30
+ [![GitHub stars](https://img.shields.io/github/stars/tavily-ai/tavily-python?style=social)](https://github.com/tavily-ai/tavily-python/stargazers)
31
+ [![PyPI - Downloads](https://img.shields.io/pypi/dm/tavily-python)](https://pypi.org/project/tavily-python/)
32
+ [![License](https://img.shields.io/github/license/tavily-ai/tavily-python)](https://github.com/tavily-ai/tavily-python/blob/main/LICENSE)
33
+ [![CI](https://github.com/tavily-ai/tavily-python/actions/workflows/tests.yml/badge.svg)](https://github.com/tavily-ai/tavily-python/actions)
34
+
30
35
  The Tavily Python wrapper allows for easy interaction with the Tavily API, offering the full range of our search and extract functionalities directly from your Python programs. Easily integrate smart search and content extraction capabilities into your applications, harnessing Tavily's powerful search and extract features.
31
36
 
32
37
  ## Installing
@@ -1,6 +1,5 @@
1
1
  LICENSE
2
2
  README.md
3
- setup.cfg
4
3
  setup.py
5
4
  tavily/__init__.py
6
5
  tavily/async_tavily.py
@@ -15,5 +14,7 @@ tavily_python.egg-info/SOURCES.txt
15
14
  tavily_python.egg-info/dependency_links.txt
16
15
  tavily_python.egg-info/requires.txt
17
16
  tavily_python.egg-info/top_level.txt
18
- tests/test_async_search.py
19
- tests/test_sync_search.py
17
+ tests/test_crawl.py
18
+ tests/test_errors.py
19
+ tests/test_map.py
20
+ tests/test_search.py
@@ -0,0 +1,102 @@
1
+ import asyncio
2
+
3
+ dummy_response = {
4
+ "base_url": "tavily.com",
5
+ "results": [
6
+ {
7
+ "url": "https://www.tavily.com/",
8
+ "raw_content": "Tavily page content",
9
+ "images": []
10
+ },
11
+ ],
12
+ "response_time": 1.5
13
+ }
14
+
15
+ def validate_default(request, response):
16
+ assert request.method == "POST"
17
+ assert request.url == "https://api.tavily.com/crawl"
18
+ assert request.headers["Authorization"] == "Bearer tvly-test"
19
+ assert request.json().get('url') == "https://tavily.com"
20
+ assert response == dummy_response
21
+
22
+ def validate_specific(request, response):
23
+ assert request.method == "POST"
24
+ assert request.url == "https://api.tavily.com/crawl"
25
+ assert request.headers["Authorization"] == "Bearer tvly-test"
26
+ assert request.timeout == 10
27
+
28
+ request_json = request.json()
29
+ for key, value in {
30
+ "url": "https://tavily.com",
31
+ "max_depth": 2,
32
+ "max_breadth": 3,
33
+ "limit": 5,
34
+ "instructions": "Focus on pricing pages",
35
+ "select_paths": ["/privacy", "/terms"],
36
+ "select_domains": ["tavily.com"],
37
+ "exclude_paths": ["/contact"],
38
+ "exclude_domains": ["example.com"],
39
+ "allow_external": False,
40
+ "include_images": True,
41
+ "categories": ["pricing", "documentation"],
42
+ "extract_depth": "advanced"
43
+ }.items():
44
+ assert request_json.get(key) == value
45
+
46
+ assert response == dummy_response
47
+
48
+ def test_sync_crawl_defaults(sync_interceptor, sync_client):
49
+ sync_interceptor.set_response(200, json=dummy_response)
50
+ response = sync_client.crawl("https://tavily.com")
51
+ request = sync_interceptor.get_request()
52
+ validate_default(request, response)
53
+
54
+ def test_sync_crawl_specific(sync_interceptor, sync_client):
55
+ sync_interceptor.set_response(200, json=dummy_response)
56
+ response = sync_client.crawl(
57
+ url="https://tavily.com",
58
+ max_depth=2,
59
+ max_breadth=3,
60
+ limit=5,
61
+ instructions="Focus on pricing pages",
62
+ select_paths=["/privacy", "/terms"],
63
+ select_domains=["tavily.com"],
64
+ exclude_paths=["/contact"],
65
+ exclude_domains=["example.com"],
66
+ allow_external=False,
67
+ include_images=True,
68
+ categories=["pricing", "documentation"],
69
+ extract_depth="advanced",
70
+ timeout=10
71
+ )
72
+
73
+ request = sync_interceptor.get_request()
74
+ validate_specific(request, response)
75
+
76
+ def test_async_crawl_defaults(async_interceptor, async_client):
77
+ async_interceptor.set_response(200, json=dummy_response)
78
+ response = asyncio.run(async_client.crawl("https://tavily.com"))
79
+ request = async_interceptor.get_request()
80
+ validate_default(request, response)
81
+
82
+ def test_async_crawl_specific(async_interceptor, async_client):
83
+ async_interceptor.set_response(200, json=dummy_response)
84
+ response = asyncio.run(async_client.crawl(
85
+ url="https://tavily.com",
86
+ max_depth=2,
87
+ max_breadth=3,
88
+ limit=5,
89
+ instructions="Focus on pricing pages",
90
+ select_paths=["/privacy", "/terms"],
91
+ select_domains=["tavily.com"],
92
+ exclude_paths=["/contact"],
93
+ exclude_domains=["example.com"],
94
+ allow_external=False,
95
+ include_images=True,
96
+ categories=["pricing", "documentation"],
97
+ extract_depth="advanced",
98
+ timeout=10
99
+ ))
100
+
101
+ request = async_interceptor.get_request()
102
+ validate_specific(request, response)
@@ -0,0 +1,50 @@
1
+ import pytest
2
+ import os
3
+ import asyncio
4
+
5
+ import tavily.tavily as sync_tavily
6
+ import tavily.async_tavily as async_tavily
7
+ from tavily.errors import MissingAPIKeyError, InvalidAPIKeyError
8
+
9
+ @pytest.fixture
10
+ def set_api_key():
11
+ old_key = os.getenv("TAVILY_API_KEY")
12
+ os.environ["TAVILY_API_KEY"] = "test_api_key"
13
+ yield
14
+ os.environ["TAVILY_API_KEY"] = old_key
15
+
16
+ @pytest.fixture
17
+ def clear_api_key():
18
+ old_key = os.getenv("TAVILY_API_KEY")
19
+ if old_key:
20
+ del os.environ["TAVILY_API_KEY"]
21
+ yield
22
+ if old_key:
23
+ os.environ["TAVILY_API_KEY"] = old_key
24
+
25
+ def test_load_key_from_env(set_api_key):
26
+ sync_tavily.TavilyClient()
27
+ async_tavily.AsyncTavilyClient()
28
+
29
+ # No error should be raised
30
+
31
+ def test_missing_api_key(clear_api_key):
32
+ with pytest.raises(MissingAPIKeyError):
33
+ sync_tavily.TavilyClient(api_key='')
34
+
35
+ with pytest.raises(MissingAPIKeyError):
36
+ async_tavily.AsyncTavilyClient(api_key='')
37
+
38
+ with pytest.raises(MissingAPIKeyError):
39
+ sync_tavily.TavilyClient()
40
+
41
+ with pytest.raises(MissingAPIKeyError):
42
+ async_tavily.AsyncTavilyClient()
43
+
44
+ def test_invalid_api_key():
45
+ with pytest.raises(InvalidAPIKeyError):
46
+ sync_tavily.TavilyClient(api_key="invalid_api_key").search("What is Tavily?")
47
+
48
+ with pytest.raises(InvalidAPIKeyError):
49
+ print(async_tavily.httpx.AsyncClient.post)
50
+ print(asyncio.run(async_tavily.AsyncTavilyClient(api_key="invalid_api_key").search("What is Tavily?")))
@@ -0,0 +1,98 @@
1
+ import asyncio
2
+
3
+ dummy_response = {
4
+ "base_url": "tavily.com",
5
+ "results": [
6
+ "https://www.tavily.com/",
7
+ "https://www.tavily.com/enterprise",
8
+ "https://www.tavily.com/terms",
9
+ "https://www.tavily.com/contact",
10
+ ],
11
+ "response_time": 0.5,
12
+ }
13
+
14
+ def validate_default(request, response):
15
+ assert request.method == "POST"
16
+ assert request.url == "https://api.tavily.com/map"
17
+ assert request.headers["Authorization"] == "Bearer tvly-test"
18
+ assert request.json().get('url') == "https://tavily.com"
19
+ assert response == dummy_response
20
+
21
+ def validate_specific(request, response):
22
+ assert request.method == "POST"
23
+ assert request.url == "https://api.tavily.com/map"
24
+ assert request.headers["Authorization"] == "Bearer tvly-test"
25
+ assert request.timeout == 10
26
+
27
+ request_json = request.json()
28
+ for key, value in {
29
+ "url": "https://tavily.com",
30
+ "max_depth": 2,
31
+ "max_breadth": 3,
32
+ "limit": 5,
33
+ "instructions": "Focus on navigation structure",
34
+ "select_paths": ["/pricing", "/docs"],
35
+ "select_domains": ["tavily.com"],
36
+ "exclude_paths": ["/blog"],
37
+ "exclude_domains": ["example.com"],
38
+ "allow_external": False,
39
+ "include_images": True,
40
+ "categories": ["pricing", "documentation"]
41
+ }.items():
42
+ assert request_json.get(key) == value
43
+
44
+ assert response == dummy_response
45
+
46
+ def test_sync_map_defaults(sync_interceptor, sync_client):
47
+ sync_interceptor.set_response(200, json=dummy_response)
48
+ response = sync_client.map("https://tavily.com")
49
+ request = sync_interceptor.get_request()
50
+ validate_default(request, response)
51
+
52
+ def test_sync_map_specific(sync_interceptor, sync_client):
53
+ sync_interceptor.set_response(200, json=dummy_response)
54
+ response = sync_client.map(
55
+ url="https://tavily.com",
56
+ max_depth=2,
57
+ max_breadth=3,
58
+ limit=5,
59
+ instructions="Focus on navigation structure",
60
+ select_paths=["/pricing", "/docs"],
61
+ select_domains=["tavily.com"],
62
+ exclude_paths=["/blog"],
63
+ exclude_domains=["example.com"],
64
+ allow_external=False,
65
+ include_images=True,
66
+ categories=["pricing", "documentation"],
67
+ timeout=10
68
+ )
69
+
70
+ request = sync_interceptor.get_request()
71
+ validate_specific(request, response)
72
+
73
+ def test_async_map_defaults(async_interceptor, async_client):
74
+ async_interceptor.set_response(200, json=dummy_response)
75
+ response = asyncio.run(async_client.map("https://tavily.com"))
76
+ request = async_interceptor.get_request()
77
+ validate_default(request, response)
78
+
79
+ def test_async_map_specific(async_interceptor, async_client):
80
+ async_interceptor.set_response(200, json=dummy_response)
81
+ response = asyncio.run(async_client.map(
82
+ url="https://tavily.com",
83
+ max_depth=2,
84
+ max_breadth=3,
85
+ limit=5,
86
+ instructions="Focus on navigation structure",
87
+ select_paths=["/pricing", "/docs"],
88
+ select_domains=["tavily.com"],
89
+ exclude_paths=["/blog"],
90
+ exclude_domains=["example.com"],
91
+ allow_external=False,
92
+ include_images=True,
93
+ categories=["pricing", "documentation"],
94
+ timeout=10
95
+ ))
96
+
97
+ request = async_interceptor.get_request()
98
+ validate_specific(request, response)
@@ -0,0 +1,99 @@
1
+ import asyncio
2
+
3
+ dummy_response = {
4
+ "query": "What is Tavily?",
5
+ "follow_up_questions": None,
6
+ "answer": None,
7
+ "images": [],
8
+ "results": [
9
+ {
10
+ "title": "Tavily",
11
+ "url": "https://tavily.com",
12
+ "content": "Connect Your LLM to the Web Empowering your AI applications with " \
13
+ "real-time, accurate search results tailored for LLMs and RAG.",
14
+ "score": 0.99,
15
+ "raw_content": None
16
+ }
17
+ ],
18
+ "response_time": 1.5
19
+ }
20
+
21
+ def validate_default(request, response):
22
+ assert request.method == "POST"
23
+ assert request.url == "https://api.tavily.com/search"
24
+ assert request.headers["Authorization"] == "Bearer tvly-test"
25
+ assert request.json().get('query') == "What is Tavily?"
26
+ assert response == dummy_response
27
+
28
+ def validate_specific(request, response):
29
+ assert request.method == "POST"
30
+ assert request.url == "https://api.tavily.com/search"
31
+ assert request.headers["Authorization"] == "Bearer tvly-test"
32
+ assert request.timeout == 10
33
+
34
+ request_json = request.json()
35
+ for key, value in {
36
+ "query": "What is Tavily?",
37
+ "search_depth": "advanced",
38
+ "topic": "news",
39
+ "days": 5,
40
+ "max_results": 10,
41
+ "include_domains": ["tavily.com"],
42
+ "exclude_domains": ["example.com"],
43
+ "include_answer": "advanced",
44
+ "include_raw_content": True,
45
+ "include_images": True
46
+ }.items():
47
+ assert request_json.get(key) == value
48
+
49
+ assert response == dummy_response
50
+
51
+ def test_sync_search_defaults(sync_interceptor, sync_client):
52
+ sync_interceptor.set_response(200, json=dummy_response)
53
+ response = sync_client.search("What is Tavily?")
54
+ request = sync_interceptor.get_request()
55
+ validate_default(request, response)
56
+
57
+ def test_sync_search_specific(sync_interceptor, sync_client):
58
+ sync_interceptor.set_response(200, json=dummy_response)
59
+ response = sync_client.search(
60
+ "What is Tavily?",
61
+ search_depth="advanced",
62
+ topic="news",
63
+ days=5,
64
+ max_results=10,
65
+ include_domains=["tavily.com"],
66
+ exclude_domains=["example.com"],
67
+ include_answer="advanced",
68
+ include_raw_content=True,
69
+ include_images=True,
70
+ timeout=10
71
+ )
72
+
73
+ request = sync_interceptor.get_request()
74
+ validate_specific(request, response)
75
+
76
+ def test_async_search_defaults(async_interceptor, async_client):
77
+ async_interceptor.set_response(200, json=dummy_response)
78
+ response = asyncio.run(async_client.search("What is Tavily?"))
79
+ request = async_interceptor.get_request()
80
+ validate_default(request, response)
81
+
82
+ def test_async_search_specific(async_interceptor, async_client):
83
+ async_interceptor.set_response(200, json=dummy_response)
84
+ response = asyncio.run(async_client.search(
85
+ "What is Tavily?",
86
+ search_depth="advanced",
87
+ topic="news",
88
+ days=5,
89
+ max_results=10,
90
+ include_domains=["tavily.com"],
91
+ exclude_domains=["example.com"],
92
+ include_answer="advanced",
93
+ include_raw_content=True,
94
+ include_images=True,
95
+ timeout=10
96
+ ))
97
+
98
+ request = async_interceptor.get_request()
99
+ validate_specific(request, response)
@@ -1,8 +0,0 @@
1
- [metadata]
2
- description_file = README.rst
3
- license_files = LICENSE
4
-
5
- [egg_info]
6
- tag_build =
7
- tag_date = 0
8
-
@@ -1,219 +0,0 @@
1
- import unittest
2
- import os
3
- from tavily import AsyncTavilyClient, MissingAPIKeyError, InvalidAPIKeyError
4
- from urllib.parse import urlparse
5
- import asyncio
6
-
7
- from unit_tests import cases
8
- class SearchTest(unittest.TestCase):
9
-
10
- def setUp(self) -> None:
11
- self.tavily_client = AsyncTavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
12
- return super().setUp()
13
-
14
- def tearDown(self) -> None:
15
- return super().tearDown()
16
-
17
- # Every single search result should have these properties
18
- def common_search_result_properties(self, result) -> None:
19
- self.assertIsInstance(result, dict)
20
- self.assertIn("title", result)
21
- self.assertIn("url", result)
22
- self.assertIn("content", result)
23
- self.assertIn("score", result)
24
-
25
- # General search results should have these properties
26
- def general_search_result_properties(self, result) -> None:
27
- self.common_search_result_properties(result)
28
- self.assertIn("raw_content", result)
29
-
30
- # News search results should have these properties
31
- def news_search_result_properties(self, result) -> None:
32
- self.common_search_result_properties(result)
33
- self.assertIn("published_date", result)
34
-
35
- # Topic-specific properties
36
- def topic_specific_properties(self, result, **params) -> None:
37
- if params.get("topic", "general") == "general":
38
- self.assertIn("raw_content", result)
39
- elif params.get("topic", "general") == "news":
40
- self.assertIn("published_date", result)
41
-
42
- # Domain inclusion/exclusion-dependent properties
43
- def domain_dependent_properties(self, response, **params) -> None:
44
- if params.get("topic", "general") != "general":
45
- return
46
-
47
- if params.get("include_domains", False) and len(params["include_domains"]) > 0:
48
- for result in response["results"]:
49
- self.assertTrue(any(domain in urlparse(result["url"]).netloc for domain in params["include_domains"]))
50
-
51
- if params.get("exclude_domains", False) and len(params["exclude_domains"]) > 0:
52
- for result in response["results"]:
53
- self.assertFalse(any(domain in urlparse(result["url"]).netloc for domain in params["exclude_domains"]))
54
-
55
- # Image-dependent properties
56
- def image_properties(self, response, **params) -> None:
57
- if params.get("topic", "general") == "general" and params.get("include_images", False):
58
- self.assertIsNotNone(response["images"])
59
- self.assertIsInstance(response["images"], list)
60
- for image in response["images"]:
61
- self.assertIsInstance(image, str)
62
-
63
- # Answer-dependent properties
64
- def answer_properties(self, response, **params) -> None:
65
- if params.get("include_answer", False):
66
- self.assertIn("answer", response)
67
- self.assertIsInstance(response["answer"], str)
68
-
69
-
70
- # Every single search response should have these properties
71
- def common_response_properties(self, response) -> None:
72
- self.assertIsNotNone(response)
73
- self.assertIsInstance(response, dict)
74
- self.assertIn("answer", response)
75
- self.assertIn("query", response)
76
- self.assertIn("results", response)
77
- self.assertIn("images", response)
78
- self.assertIn("response_time", response)
79
- self.assertIn("follow_up_questions", response)
80
-
81
- self.assertIsNotNone(response["query"])
82
- self.assertIsNotNone(response["results"])
83
- self.assertIsNotNone(response["response_time"])
84
- self.assertIsNotNone(response["images"])
85
-
86
- self.assertIsInstance(response["query"], str)
87
- self.assertIsInstance(response["results"], list)
88
- self.assertIsInstance(response["response_time"], float)
89
- self.assertIsInstance(response["images"], list)
90
-
91
- # Search responses also have properties that depend on the request params
92
- def custom_response_properties(self, response, **params) -> None:
93
- self.domain_dependent_properties(response, **params)
94
- self.image_properties(response, **params)
95
- self.answer_properties(response, **params)
96
- for result in response["results"]:
97
- self.topic_specific_properties(result, **params)
98
-
99
- def test_internal_search(self) -> None:
100
- for test_case in cases:
101
- with self.subTest(msg=test_case["name"]):
102
- response = asyncio.run(self.tavily_client._search(**test_case["params"]))
103
- self.common_response_properties(response)
104
- if test_case["params"].get("topic", "general") == "general":
105
- for search_result in response["results"]:
106
- self.general_search_result_properties(search_result)
107
- elif test_case["params"].get("topic", "general") == "news":
108
- for search_result in response["results"]:
109
- self.news_search_result_properties(search_result)
110
- self.custom_response_properties(response, **test_case["params"])
111
-
112
- def test_external_search(self) -> None:
113
- for test_case in cases:
114
- with self.subTest(msg=test_case["name"]):
115
- response = asyncio.run(self.tavily_client.search(**test_case["params"]))
116
- self.common_response_properties(response)
117
- if test_case["params"].get("topic", "general") == "general":
118
- for search_result in response["results"]:
119
- self.general_search_result_properties(search_result)
120
- elif test_case["params"].get("topic", "general") == "news":
121
- for search_result in response["results"]:
122
- self.news_search_result_properties(search_result)
123
- self.custom_response_properties(response, **test_case["params"])
124
-
125
- class QNASearchTest(unittest.TestCase):
126
-
127
- def setUp(self) -> None:
128
- self.tavily_client = AsyncTavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
129
- return super().setUp()
130
-
131
- def tearDown(self) -> None:
132
- return super().tearDown()
133
-
134
- def test_qna_search(self) -> None:
135
- for test_case in cases:
136
- if "include_answer" in test_case["params"]:
137
- del test_case["params"]["include_answer"]
138
- if "include_raw_content" in test_case["params"]:
139
- del test_case["params"]["include_raw_content"]
140
- if "include_images" in test_case["params"]:
141
- del test_case["params"]["include_images"]
142
- with self.subTest(msg=test_case["name"]):
143
- response = asyncio.run(self.tavily_client.qna_search(**test_case["params"]))
144
- self.assertIsNotNone(response)
145
- self.assertIsInstance(response, str)
146
- self.assertTrue(len(response) > 0)
147
-
148
- class CompanyInfoSearchTest(unittest.TestCase):
149
-
150
- def setUp(self) -> None:
151
- self.tavily_client = AsyncTavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
152
- return super().setUp()
153
-
154
- def tearDown(self) -> None:
155
- return super().tearDown()
156
-
157
- # Every single search result should have these properties
158
- def common_search_result_properties(self, result) -> None:
159
- self.assertIsInstance(result, dict)
160
- self.assertIn("title", result)
161
- self.assertIn("url", result)
162
- self.assertIn("content", result)
163
- self.assertIn("score", result)
164
-
165
- def test_company_info_search(self) -> None:
166
- for test_case in cases:
167
- if "topic" in test_case["params"]:
168
- del test_case["params"]["topic"]
169
- if "include_domains" in test_case["params"]:
170
- del test_case["params"]["include_domains"]
171
- if "exclude_domains" in test_case["params"]:
172
- del test_case["params"]["exclude_domains"]
173
- if "include_raw_content" in test_case["params"]:
174
- del test_case["params"]["include_raw_content"]
175
- if "include_images" in test_case["params"]:
176
- del test_case["params"]["include_images"]
177
- if "include_answer" in test_case["params"]:
178
- del test_case["params"]["include_answer"]
179
- if "use_cache" in test_case["params"]:
180
- del test_case["params"]["use_cache"]
181
- with self.subTest(msg=test_case["name"]):
182
- response = asyncio.run(self.tavily_client.get_company_info(**test_case["params"]))
183
- self.assertIsNotNone(response)
184
- self.assertIsInstance(response, list)
185
- self.assertTrue(len(response) > 0)
186
- for search_result in response:
187
- self.common_search_result_properties(search_result)
188
-
189
- class ErrorTest(unittest.TestCase):
190
-
191
- def setUp(self) -> None:
192
- return super().setUp()
193
-
194
- def tearDown(self) -> None:
195
- return super().tearDown()
196
-
197
- # This test is here to ensure that no MissingAPIKeyError is raised when the API key is in the environment
198
- def test_load_key_from_env(self) -> None:
199
- self.assertIn('results', asyncio.run(AsyncTavilyClient().search("Why is Tavily the best search API?")))
200
-
201
- def test_missing_api_key(self) -> None:
202
- with self.assertRaises(MissingAPIKeyError):
203
- AsyncTavilyClient(api_key='')
204
-
205
- old_key = os.getenv("TAVILY_API_KEY")
206
- del os.environ["TAVILY_API_KEY"]
207
- with self.assertRaises(MissingAPIKeyError):
208
- AsyncTavilyClient()
209
-
210
- os.environ["TAVILY_API_KEY"] = old_key
211
-
212
-
213
- def test_invalid_api_key(self) -> None:
214
- with self.assertRaises(InvalidAPIKeyError):
215
- asyncio.run(AsyncTavilyClient(api_key="invalid_api_key").search("Why is Tavily the best search API?"))
216
-
217
- if __name__ == "__main__":
218
-
219
- unittest.main()
@@ -1,219 +0,0 @@
1
- import unittest
2
- import os
3
- from tavily import TavilyClient, InvalidAPIKeyError, MissingAPIKeyError
4
- from urllib.parse import urlparse
5
-
6
- from unit_tests import cases
7
-
8
- class SearchTest(unittest.TestCase):
9
-
10
- def setUp(self) -> None:
11
- self.tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
12
- return super().setUp()
13
-
14
- def tearDown(self) -> None:
15
- return super().tearDown()
16
-
17
- # Every single search result should have these properties
18
- def common_search_result_properties(self, result) -> None:
19
- self.assertIsInstance(result, dict)
20
- self.assertIn("title", result)
21
- self.assertIn("url", result)
22
- self.assertIn("content", result)
23
- self.assertIn("score", result)
24
-
25
- # General search results should have these properties
26
- def general_search_result_properties(self, result) -> None:
27
- self.common_search_result_properties(result)
28
- self.assertIn("raw_content", result)
29
-
30
- # News search results should have these properties
31
- def news_search_result_properties(self, result) -> None:
32
- self.common_search_result_properties(result)
33
- self.assertIn("published_date", result)
34
-
35
- # Topic-specific properties
36
- def topic_specific_properties(self, result, **params) -> None:
37
- if params.get("topic", "general") == "general":
38
- self.assertIn("raw_content", result)
39
- elif params.get("topic", "general") == "news":
40
- self.assertIn("published_date", result)
41
-
42
- # Domain inclusion/exclusion-dependent properties
43
- def domain_dependent_properties(self, response, **params) -> None:
44
- if params.get("topic", "general") != "general":
45
- return
46
-
47
- if params.get("include_domains", False) and len(params["include_domains"]) > 0:
48
- for result in response["results"]:
49
- self.assertTrue(any(domain in urlparse(result["url"]).netloc for domain in params["include_domains"]))
50
-
51
- if params.get("exclude_domains", False) and len(params["exclude_domains"]) > 0:
52
- for result in response["results"]:
53
- self.assertFalse(any(domain in urlparse(result["url"]).netloc for domain in params["exclude_domains"]))
54
-
55
- # Image-dependent properties
56
- def image_properties(self, response, **params) -> None:
57
- if params.get("topic", "general") == "general" and params.get("include_images", False):
58
- self.assertIsNotNone(response["images"])
59
- self.assertIsInstance(response["images"], list)
60
- for image in response["images"]:
61
- self.assertIsInstance(image, str)
62
-
63
- # Answer-dependent properties
64
- def answer_properties(self, response, **params) -> None:
65
- if params.get("include_answer", False):
66
- self.assertIn("answer", response)
67
- self.assertIsInstance(response["answer"], str)
68
-
69
-
70
- # Every single search response should have these properties
71
- def common_response_properties(self, response) -> None:
72
- self.assertIsNotNone(response)
73
- self.assertIsInstance(response, dict)
74
- self.assertIn("answer", response)
75
- self.assertIn("query", response)
76
- self.assertIn("results", response)
77
- self.assertIn("images", response)
78
- self.assertIn("response_time", response)
79
- self.assertIn("follow_up_questions", response)
80
-
81
- self.assertIsNotNone(response["query"])
82
- self.assertIsNotNone(response["results"])
83
- self.assertIsNotNone(response["response_time"])
84
- self.assertIsNotNone(response["images"])
85
-
86
- self.assertIsInstance(response["query"], str)
87
- self.assertIsInstance(response["results"], list)
88
- self.assertIsInstance(response["response_time"], float)
89
- self.assertIsInstance(response["images"], list)
90
-
91
- # Search responses also have properties that depend on the request params
92
- def custom_response_properties(self, response, **params) -> None:
93
- self.domain_dependent_properties(response, **params)
94
- self.image_properties(response, **params)
95
- self.answer_properties(response, **params)
96
- for result in response["results"]:
97
- self.topic_specific_properties(result, **params)
98
-
99
- def test_internal_search(self) -> None:
100
- for test_case in cases:
101
- with self.subTest(msg=test_case["name"]):
102
- result = self.tavily_client._search(**test_case["params"])
103
- self.common_response_properties(result)
104
- if test_case["params"].get("topic", "general") == "general":
105
- for search_result in result["results"]:
106
- self.general_search_result_properties(search_result)
107
- elif test_case["params"].get("topic", "general") == "news":
108
- for search_result in result["results"]:
109
- self.news_search_result_properties(search_result)
110
- self.custom_response_properties(result, **test_case["params"])
111
-
112
- def test_external_search(self) -> None:
113
- for test_case in cases:
114
- with self.subTest(msg=test_case["name"]):
115
- response = self.tavily_client._search(**test_case["params"])
116
- self.common_response_properties(response)
117
- if test_case["params"].get("topic", "general") == "general":
118
- for search_result in response["results"]:
119
- self.general_search_result_properties(search_result)
120
- elif test_case["params"].get("topic", "general") == "news":
121
- for search_result in response["results"]:
122
- self.news_search_result_properties(search_result)
123
- self.custom_response_properties(response, **test_case["params"])
124
-
125
- class QNASearchTest(unittest.TestCase):
126
-
127
- def setUp(self) -> None:
128
- self.tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
129
- return super().setUp()
130
-
131
- def tearDown(self) -> None:
132
- return super().tearDown()
133
-
134
- def test_qna_search(self) -> None:
135
- for test_case in cases:
136
- if "include_answer" in test_case["params"]:
137
- del test_case["params"]["include_answer"]
138
- if "include_raw_content" in test_case["params"]:
139
- del test_case["params"]["include_raw_content"]
140
- if "include_images" in test_case["params"]:
141
- del test_case["params"]["include_images"]
142
- with self.subTest(msg=test_case["name"]):
143
- response = self.tavily_client.qna_search(**test_case["params"])
144
- self.assertIsNotNone(response)
145
- self.assertIsInstance(response, str)
146
- self.assertTrue(len(response) > 0)
147
-
148
- class CompanyInfoSearchTest(unittest.TestCase):
149
-
150
- def setUp(self) -> None:
151
- self.tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
152
- return super().setUp()
153
-
154
- def tearDown(self) -> None:
155
- return super().tearDown()
156
-
157
- # Every single search result should have these properties
158
- def common_search_result_properties(self, result) -> None:
159
- self.assertIsInstance(result, dict)
160
- self.assertIn("title", result)
161
- self.assertIn("url", result)
162
- self.assertIn("content", result)
163
- self.assertIn("score", result)
164
-
165
- def test_company_info_search(self) -> None:
166
- for test_case in cases:
167
- if "topic" in test_case["params"]:
168
- del test_case["params"]["topic"]
169
- if "include_domains" in test_case["params"]:
170
- del test_case["params"]["include_domains"]
171
- if "exclude_domains" in test_case["params"]:
172
- del test_case["params"]["exclude_domains"]
173
- if "include_raw_content" in test_case["params"]:
174
- del test_case["params"]["include_raw_content"]
175
- if "include_images" in test_case["params"]:
176
- del test_case["params"]["include_images"]
177
- if "include_answer" in test_case["params"]:
178
- del test_case["params"]["include_answer"]
179
- if "use_cache" in test_case["params"]:
180
- del test_case["params"]["use_cache"]
181
- with self.subTest(msg=test_case["name"]):
182
- response = self.tavily_client.get_company_info(**test_case["params"])
183
- self.assertIsNotNone(response)
184
- self.assertIsInstance(response, list)
185
- self.assertTrue(len(response) > 0)
186
- for search_result in response:
187
- self.common_search_result_properties(search_result)
188
-
189
- class ErrorTest(unittest.TestCase):
190
-
191
- def setUp(self) -> None:
192
- return super().setUp()
193
-
194
- def tearDown(self) -> None:
195
- return super().tearDown()
196
-
197
- # This test is here to ensure that no MissingAPIKeyError is raised when the API key is in the environment
198
- def test_load_key_from_env(self) -> None:
199
- self.assertIn('results', TavilyClient().search("Why is Tavily the best search API?"))
200
-
201
- def test_missing_api_key(self) -> None:
202
- with self.assertRaises(MissingAPIKeyError):
203
- TavilyClient(api_key='')
204
-
205
- old_key = os.getenv("TAVILY_API_KEY")
206
- del os.environ["TAVILY_API_KEY"]
207
- with self.assertRaises(MissingAPIKeyError):
208
- TavilyClient()
209
-
210
- os.environ["TAVILY_API_KEY"] = old_key
211
-
212
-
213
- def test_invalid_api_key(self) -> None:
214
- with self.assertRaises(InvalidAPIKeyError):
215
- TavilyClient(api_key="invalid_api_key").search("Why is Tavily the best search API?")
216
-
217
- if __name__ == "__main__":
218
-
219
- unittest.main()
File without changes