tavily-python 0.7.5__tar.gz → 0.7.6__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.5 → tavily_python-0.7.6}/PKG-INFO +1 -1
- {tavily_python-0.7.5 → tavily_python-0.7.6}/setup.py +1 -1
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily/async_tavily.py +2 -1
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily/tavily.py +2 -1
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily_python.egg-info/PKG-INFO +1 -1
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily_python.egg-info/SOURCES.txt +3 -1
- tavily_python-0.7.6/tests/test_async_search.py +219 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tests/test_crawl.py +2 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tests/test_map.py +2 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tests/test_search.py +2 -0
- tavily_python-0.7.6/tests/test_sync_search.py +219 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/LICENSE +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/README.md +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/setup.cfg +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily/__init__.py +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily/config.py +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily/errors.py +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily/hybrid_rag/__init__.py +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily/hybrid_rag/hybrid_rag.py +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily/utils.py +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily_python.egg-info/dependency_links.txt +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily_python.egg-info/requires.txt +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tavily_python.egg-info/top_level.txt +0 -0
- {tavily_python-0.7.5 → tavily_python-0.7.6}/tests/test_errors.py +0 -0
|
@@ -42,7 +42,8 @@ class AsyncTavilyClient:
|
|
|
42
42
|
self._client_creator = lambda: httpx.AsyncClient(
|
|
43
43
|
headers={
|
|
44
44
|
"Content-Type": "application/json",
|
|
45
|
-
"Authorization": f"Bearer {api_key}"
|
|
45
|
+
"Authorization": f"Bearer {api_key}",
|
|
46
|
+
"X-Client-Source": "tavily-python"
|
|
46
47
|
},
|
|
47
48
|
base_url="https://api.tavily.com",
|
|
48
49
|
mounts=proxy_mounts
|
|
@@ -32,7 +32,8 @@ class TavilyClient:
|
|
|
32
32
|
self.proxies = resolved_proxies
|
|
33
33
|
self.headers = {
|
|
34
34
|
"Content-Type": "application/json",
|
|
35
|
-
"Authorization": f"Bearer {self.api_key}"
|
|
35
|
+
"Authorization": f"Bearer {self.api_key}",
|
|
36
|
+
"X-Client-Source": "tavily-python"
|
|
36
37
|
}
|
|
37
38
|
|
|
38
39
|
def _search(self,
|
|
@@ -14,7 +14,9 @@ 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
|
|
17
18
|
tests/test_crawl.py
|
|
18
19
|
tests/test_errors.py
|
|
19
20
|
tests/test_map.py
|
|
20
|
-
tests/test_search.py
|
|
21
|
+
tests/test_search.py
|
|
22
|
+
tests/test_sync_search.py
|
|
@@ -0,0 +1,219 @@
|
|
|
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()
|
|
@@ -16,6 +16,7 @@ def validate_default(request, response):
|
|
|
16
16
|
assert request.method == "POST"
|
|
17
17
|
assert request.url == "https://api.tavily.com/crawl"
|
|
18
18
|
assert request.headers["Authorization"] == "Bearer tvly-test"
|
|
19
|
+
assert request.headers["X-Client-Source"] == "tavily-python"
|
|
19
20
|
assert request.json().get('url') == "https://tavily.com"
|
|
20
21
|
assert response == dummy_response
|
|
21
22
|
|
|
@@ -23,6 +24,7 @@ def validate_specific(request, response):
|
|
|
23
24
|
assert request.method == "POST"
|
|
24
25
|
assert request.url == "https://api.tavily.com/crawl"
|
|
25
26
|
assert request.headers["Authorization"] == "Bearer tvly-test"
|
|
27
|
+
assert request.headers["X-Client-Source"] == "tavily-python"
|
|
26
28
|
assert request.timeout == 10
|
|
27
29
|
|
|
28
30
|
request_json = request.json()
|
|
@@ -15,6 +15,7 @@ def validate_default(request, response):
|
|
|
15
15
|
assert request.method == "POST"
|
|
16
16
|
assert request.url == "https://api.tavily.com/map"
|
|
17
17
|
assert request.headers["Authorization"] == "Bearer tvly-test"
|
|
18
|
+
assert request.headers["X-Client-Source"] == "tavily-python"
|
|
18
19
|
assert request.json().get('url') == "https://tavily.com"
|
|
19
20
|
assert response == dummy_response
|
|
20
21
|
|
|
@@ -22,6 +23,7 @@ def validate_specific(request, response):
|
|
|
22
23
|
assert request.method == "POST"
|
|
23
24
|
assert request.url == "https://api.tavily.com/map"
|
|
24
25
|
assert request.headers["Authorization"] == "Bearer tvly-test"
|
|
26
|
+
assert request.headers["X-Client-Source"] == "tavily-python"
|
|
25
27
|
assert request.timeout == 10
|
|
26
28
|
|
|
27
29
|
request_json = request.json()
|
|
@@ -22,6 +22,7 @@ def validate_default(request, response):
|
|
|
22
22
|
assert request.method == "POST"
|
|
23
23
|
assert request.url == "https://api.tavily.com/search"
|
|
24
24
|
assert request.headers["Authorization"] == "Bearer tvly-test"
|
|
25
|
+
assert request.headers["X-Client-Source"] == "tavily-python"
|
|
25
26
|
assert request.json().get('query') == "What is Tavily?"
|
|
26
27
|
assert response == dummy_response
|
|
27
28
|
|
|
@@ -29,6 +30,7 @@ def validate_specific(request, response):
|
|
|
29
30
|
assert request.method == "POST"
|
|
30
31
|
assert request.url == "https://api.tavily.com/search"
|
|
31
32
|
assert request.headers["Authorization"] == "Bearer tvly-test"
|
|
33
|
+
assert request.headers["X-Client-Source"] == "tavily-python"
|
|
32
34
|
assert request.timeout == 10
|
|
33
35
|
|
|
34
36
|
request_json = request.json()
|
|
@@ -0,0 +1,219 @@
|
|
|
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
|