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.
- {tavily_python-0.7.10 → tavily_python-0.7.12}/PKG-INFO +1 -1
- {tavily_python-0.7.10 → tavily_python-0.7.12}/setup.py +1 -1
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/async_tavily.py +3 -11
- tavily_python-0.7.12/tavily/config.py +2 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/tavily.py +3 -11
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/PKG-INFO +1 -1
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/SOURCES.txt +1 -3
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tests/test_crawl.py +0 -3
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tests/test_map.py +1 -4
- tavily_python-0.7.10/tavily/config.py +0 -12
- tavily_python-0.7.10/tests/test_async_search.py +0 -219
- tavily_python-0.7.10/tests/test_sync_search.py +0 -219
- {tavily_python-0.7.10 → tavily_python-0.7.12}/LICENSE +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/README.md +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/setup.cfg +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/__init__.py +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/errors.py +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/hybrid_rag/__init__.py +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/hybrid_rag/hybrid_rag.py +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily/utils.py +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/dependency_links.txt +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/requires.txt +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tavily_python.egg-info/top_level.txt +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tests/test_errors.py +0 -0
- {tavily_python-0.7.10 → tavily_python-0.7.12}/tests/test_search.py +0 -0
|
@@ -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 =
|
|
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 =
|
|
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
|
|
|
@@ -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 =
|
|
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 =
|
|
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
|
|
|
@@ -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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|