tavily-python 0.7.10__tar.gz → 0.7.12__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.10 → tavily_python-0.7.12}/PKG-INFO +1 -1
  2. {tavily_python-0.7.10 → tavily_python-0.7.12}/setup.py +1 -1
  3. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/async_tavily.py +3 -11
  4. tavily_python-0.7.12/tavily/config.py +2 -0
  5. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/tavily.py +3 -11
  6. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/PKG-INFO +1 -1
  7. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/SOURCES.txt +1 -3
  8. {tavily_python-0.7.10 → tavily_python-0.7.12}/tests/test_crawl.py +0 -3
  9. {tavily_python-0.7.10 → tavily_python-0.7.12}/tests/test_map.py +1 -4
  10. tavily_python-0.7.10/tavily/config.py +0 -12
  11. tavily_python-0.7.10/tests/test_async_search.py +0 -219
  12. tavily_python-0.7.10/tests/test_sync_search.py +0 -219
  13. {tavily_python-0.7.10 → tavily_python-0.7.12}/LICENSE +0 -0
  14. {tavily_python-0.7.10 → tavily_python-0.7.12}/README.md +0 -0
  15. {tavily_python-0.7.10 → tavily_python-0.7.12}/setup.cfg +0 -0
  16. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/__init__.py +0 -0
  17. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/errors.py +0 -0
  18. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/hybrid_rag/__init__.py +0 -0
  19. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/hybrid_rag/hybrid_rag.py +0 -0
  20. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/utils.py +0 -0
  21. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/dependency_links.txt +0 -0
  22. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/requires.txt +0 -0
  23. {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/top_level.txt +0 -0
  24. {tavily_python-0.7.10 → tavily_python-0.7.12}/tests/test_errors.py +0 -0
  25. {tavily_python-0.7.10 → tavily_python-0.7.12}/tests/test_search.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tavily-python
3
- Version: 0.7.10
3
+ Version: 0.7.12
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
@@ -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.10',
8
+ version='0.7.12',
9
9
  url='https://github.com/tavily-ai/tavily-python',
10
10
  author='Tavily AI',
11
11
  author_email='support@tavily.com',
@@ -7,7 +7,6 @@ import httpx
7
7
 
8
8
  from .utils import get_max_items_from_list
9
9
  from .errors import UsageLimitExceededError, InvalidAPIKeyError, MissingAPIKeyError, BadRequestError, ForbiddenError, TimeoutError
10
- from .config import AllowedCategory
11
10
 
12
11
 
13
12
  class AsyncTavilyClient:
@@ -184,7 +183,7 @@ class AsyncTavilyClient:
184
183
  include_images: bool = None,
185
184
  extract_depth: Literal["basic", "advanced"] = None,
186
185
  format: Literal["markdown", "text"] = None,
187
- timeout: int = 60,
186
+ timeout: int = 30,
188
187
  include_favicon: bool = None,
189
188
  **kwargs
190
189
  ) -> dict:
@@ -197,6 +196,7 @@ class AsyncTavilyClient:
197
196
  "include_images": include_images,
198
197
  "extract_depth": extract_depth,
199
198
  "format": format,
199
+ "timeout": timeout,
200
200
  "include_favicon": include_favicon,
201
201
  }
202
202
 
@@ -239,7 +239,7 @@ class AsyncTavilyClient:
239
239
  include_images: bool = None,
240
240
  extract_depth: Literal["basic", "advanced"] = None,
241
241
  format: Literal["markdown", "text"] = None,
242
- timeout: int = 60,
242
+ timeout: int = 30,
243
243
  include_favicon: bool = None,
244
244
  **kwargs, # Accept custom arguments
245
245
  ) -> dict:
@@ -277,7 +277,6 @@ class AsyncTavilyClient:
277
277
  exclude_domains: Sequence[str] = None,
278
278
  allow_external: bool = None,
279
279
  include_images: bool = None,
280
- categories: Sequence[AllowedCategory] = None,
281
280
  extract_depth: Literal["basic", "advanced"] = None,
282
281
  format: Literal["markdown", "text"] = None,
283
282
  timeout: int = 60,
@@ -298,7 +297,6 @@ class AsyncTavilyClient:
298
297
  "exclude_paths": exclude_paths,
299
298
  "exclude_domains": exclude_domains,
300
299
  "allow_external": allow_external,
301
- "categories": categories,
302
300
  "include_images": include_images,
303
301
  "extract_depth": extract_depth,
304
302
  "format": format,
@@ -349,7 +347,6 @@ class AsyncTavilyClient:
349
347
  exclude_paths: Sequence[str] = None,
350
348
  exclude_domains: Sequence[str] = None,
351
349
  allow_external: bool = None,
352
- categories: Sequence[AllowedCategory] = None,
353
350
  extract_depth: Literal["basic", "advanced"] = None,
354
351
  include_images: bool = None,
355
352
  format: Literal["markdown", "text"] = None,
@@ -372,7 +369,6 @@ class AsyncTavilyClient:
372
369
  exclude_paths=exclude_paths,
373
370
  exclude_domains=exclude_domains,
374
371
  allow_external=allow_external,
375
- categories=categories,
376
372
  extract_depth=extract_depth,
377
373
  include_images=include_images,
378
374
  format=format,
@@ -394,7 +390,6 @@ class AsyncTavilyClient:
394
390
  exclude_domains: Sequence[str] = None,
395
391
  allow_external: bool = None,
396
392
  include_images: bool = None,
397
- categories: Sequence[AllowedCategory] = None,
398
393
  timeout: int = 60,
399
394
  **kwargs
400
395
  ) -> dict:
@@ -413,7 +408,6 @@ class AsyncTavilyClient:
413
408
  "exclude_domains": exclude_domains,
414
409
  "allow_external": allow_external,
415
410
  "include_images": include_images,
416
- "categories": categories,
417
411
  }
418
412
 
419
413
  if kwargs:
@@ -461,7 +455,6 @@ class AsyncTavilyClient:
461
455
  exclude_domains: Sequence[str] = None,
462
456
  allow_external: bool = None,
463
457
  include_images: bool = None,
464
- categories: Sequence[AllowedCategory] = None,
465
458
  timeout: int = 60,
466
459
  **kwargs
467
460
  ) -> dict:
@@ -481,7 +474,6 @@ class AsyncTavilyClient:
481
474
  exclude_domains=exclude_domains,
482
475
  allow_external=allow_external,
483
476
  include_images=include_images,
484
- categories=categories,
485
477
  timeout=timeout,
486
478
  **kwargs)
487
479
 
@@ -0,0 +1,2 @@
1
+ DEFAULT_MODEL_ENCODING = "gpt-3.5-turbo"
2
+ DEFAULT_MAX_TOKENS = 4000
@@ -6,7 +6,6 @@ from typing import Literal, Sequence, Optional, List, Union
6
6
  from concurrent.futures import ThreadPoolExecutor, as_completed
7
7
  from .utils import get_max_items_from_list
8
8
  from .errors import UsageLimitExceededError, InvalidAPIKeyError, MissingAPIKeyError, BadRequestError, ForbiddenError, TimeoutError
9
- from .config import AllowedCategory
10
9
 
11
10
  class TavilyClient:
12
11
  """
@@ -168,7 +167,7 @@ class TavilyClient:
168
167
  include_images: bool = None,
169
168
  extract_depth: Literal["basic", "advanced"] = None,
170
169
  format: Literal["markdown", "text"] = None,
171
- timeout: int = 60,
170
+ timeout: int = 30,
172
171
  include_favicon: bool = None,
173
172
  **kwargs
174
173
  ) -> dict:
@@ -180,6 +179,7 @@ class TavilyClient:
180
179
  "include_images": include_images,
181
180
  "extract_depth": extract_depth,
182
181
  "format": format,
182
+ "timeout": timeout,
183
183
  "include_favicon": include_favicon,
184
184
  }
185
185
 
@@ -220,7 +220,7 @@ class TavilyClient:
220
220
  include_images: bool = None,
221
221
  extract_depth: Literal["basic", "advanced"] = None,
222
222
  format: Literal["markdown", "text"] = None,
223
- timeout: int = 60,
223
+ timeout: int = 30,
224
224
  include_favicon: bool = None,
225
225
  **kwargs, # Accept custom arguments
226
226
  ) -> dict:
@@ -256,7 +256,6 @@ class TavilyClient:
256
256
  exclude_domains: Sequence[str] = None,
257
257
  allow_external: bool = None,
258
258
  include_images: bool = None,
259
- categories: Sequence[AllowedCategory] = None,
260
259
  extract_depth: Literal["basic", "advanced"] = None,
261
260
  format: Literal["markdown", "text"] = None,
262
261
  timeout: int = 60,
@@ -279,7 +278,6 @@ class TavilyClient:
279
278
  "exclude_domains": exclude_domains,
280
279
  "allow_external": allow_external,
281
280
  "include_images": include_images,
282
- "categories": categories,
283
281
  "extract_depth": extract_depth,
284
282
  "format": format,
285
283
  "include_favicon": include_favicon,
@@ -330,7 +328,6 @@ class TavilyClient:
330
328
  exclude_domains: Sequence[str] = None,
331
329
  allow_external: bool = None,
332
330
  include_images: bool = None,
333
- categories: Sequence[AllowedCategory] = None,
334
331
  extract_depth: Literal["basic", "advanced"] = None,
335
332
  format: Literal["markdown", "text"] = None,
336
333
  timeout: int = 60,
@@ -353,7 +350,6 @@ class TavilyClient:
353
350
  exclude_domains=exclude_domains,
354
351
  allow_external=allow_external,
355
352
  include_images=include_images,
356
- categories=categories,
357
353
  extract_depth=extract_depth,
358
354
  format=format,
359
355
  timeout=timeout,
@@ -374,7 +370,6 @@ class TavilyClient:
374
370
  exclude_domains: Sequence[str] = None,
375
371
  allow_external: bool = None,
376
372
  include_images: bool = None,
377
- categories: Sequence[AllowedCategory] = None,
378
373
  timeout: int = 60,
379
374
  **kwargs
380
375
  ) -> dict:
@@ -393,7 +388,6 @@ class TavilyClient:
393
388
  "exclude_domains": exclude_domains,
394
389
  "allow_external": allow_external,
395
390
  "include_images": include_images,
396
- "categories": categories,
397
391
  }
398
392
 
399
393
  if kwargs:
@@ -441,7 +435,6 @@ class TavilyClient:
441
435
  exclude_domains: Sequence[str] = None,
442
436
  allow_external: bool = None,
443
437
  include_images: bool = None,
444
- categories: Sequence[AllowedCategory] = None,
445
438
  timeout: int = 60,
446
439
  **kwargs
447
440
  ) -> dict:
@@ -461,7 +454,6 @@ class TavilyClient:
461
454
  exclude_domains=exclude_domains,
462
455
  allow_external=allow_external,
463
456
  include_images=include_images,
464
- categories=categories,
465
457
  timeout=timeout,
466
458
  **kwargs)
467
459
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tavily-python
3
- Version: 0.7.10
3
+ Version: 0.7.12
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
@@ -14,9 +14,7 @@ tavily_python.egg-info/SOURCES.txt
14
14
  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
- tests/test_async_search.py
18
17
  tests/test_crawl.py
19
18
  tests/test_errors.py
20
19
  tests/test_map.py
21
- tests/test_search.py
22
- tests/test_sync_search.py
20
+ tests/test_search.py
@@ -40,7 +40,6 @@ def validate_specific(request, response):
40
40
  "exclude_domains": ["example.com"],
41
41
  "allow_external": False,
42
42
  "include_images": True,
43
- "categories": ["pricing", "documentation"],
44
43
  "extract_depth": "advanced"
45
44
  }.items():
46
45
  assert request_json.get(key) == value
@@ -67,7 +66,6 @@ def test_sync_crawl_specific(sync_interceptor, sync_client):
67
66
  exclude_domains=["example.com"],
68
67
  allow_external=False,
69
68
  include_images=True,
70
- categories=["pricing", "documentation"],
71
69
  extract_depth="advanced",
72
70
  timeout=10
73
71
  )
@@ -95,7 +93,6 @@ def test_async_crawl_specific(async_interceptor, async_client):
95
93
  exclude_domains=["example.com"],
96
94
  allow_external=False,
97
95
  include_images=True,
98
- categories=["pricing", "documentation"],
99
96
  extract_depth="advanced",
100
97
  timeout=10
101
98
  ))
@@ -38,8 +38,7 @@ def validate_specific(request, response):
38
38
  "exclude_paths": ["/blog"],
39
39
  "exclude_domains": ["example.com"],
40
40
  "allow_external": False,
41
- "include_images": True,
42
- "categories": ["pricing", "documentation"]
41
+ "include_images": True
43
42
  }.items():
44
43
  assert request_json.get(key) == value
45
44
 
@@ -65,7 +64,6 @@ def test_sync_map_specific(sync_interceptor, sync_client):
65
64
  exclude_domains=["example.com"],
66
65
  allow_external=False,
67
66
  include_images=True,
68
- categories=["pricing", "documentation"],
69
67
  timeout=10
70
68
  )
71
69
 
@@ -92,7 +90,6 @@ def test_async_map_specific(async_interceptor, async_client):
92
90
  exclude_domains=["example.com"],
93
91
  allow_external=False,
94
92
  include_images=True,
95
- categories=["pricing", "documentation"],
96
93
  timeout=10
97
94
  ))
98
95
 
@@ -1,12 +0,0 @@
1
- from typing import Literal
2
-
3
- DEFAULT_MODEL_ENCODING = "gpt-3.5-turbo"
4
- DEFAULT_MAX_TOKENS = 4000
5
-
6
- # Create a type that represents all allowed categories
7
- AllowedCategory = Literal[
8
- "Documentation", "Blog", "Blogs", "Community", "About", "Contact",
9
- "Privacy", "Terms", "Status", "Pricing", "Enterprise", "Careers",
10
- "E-Commerce", "Authentication", "Developer", "Developers", "Solutions",
11
- "Partners", "Downloads", "Media", "Events", "People"
12
- ]
@@ -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
File without changes
File without changes