cutting new verison with exa search tool

This commit is contained in:
João Moura
2024-05-02 23:15:32 -03:00
parent 59d9d9eb1f
commit 037e80c6a3

View File

@@ -5,24 +5,23 @@ from typing import Any
from .exa_base_tool import EXABaseTool from .exa_base_tool import EXABaseTool
class EXASearchTool(EXABaseTool): class EXASearchTool(EXABaseTool):
def _run( def _run(
self, self,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
search_query = kwargs.get('search_query') search_query = kwargs.get('search_query')
if search_query is None: if search_query is None:
search_query = kwargs.get('query') search_query = kwargs.get('query')
payload = { payload = {
"query": search_query, "query": search_query,
} }
headers = self.headers.copy() headers = self.headers.copy()
headers["x-api-key"] = os.environ['EXA_API_KEY'] headers["x-api-key"] = os.environ['EXA_API_KEY']
response = requests.post(self.search_url, json=payload, headers=headers) response = requests.post(self.search_url, json=payload, headers=headers)
results = response.json() results = response.json()
if 'results' in results: if 'results' in results:
results = super()._parse_results(results['results']) results = super()._parse_results(results['results'])
else: return results
return results