pushing initial exa search tool

This commit is contained in:
João Moura
2024-05-02 02:48:21 -03:00
parent 5fb9ddfa2a
commit ffd5942b31
2 changed files with 61 additions and 0 deletions

View File

@@ -0,0 +1,37 @@
import os
from typing import Type
from pydantic.v1 import BaseModel, Field
from crewai_tools.tools.base_tool import BaseTool
class EXABaseToolToolSchema(BaseModel):
"""Input for EXABaseTool."""
search_query: str = Field(..., description="Mandatory search query you want to use to search the internet")
class EXABaseTool(BaseTool):
name: str = "Search the internet"
description: str = "A tool that can be used to search the internet from a search_query"
args_schema: Type[BaseModel] = EXABaseToolToolSchema
search_url: str = "https://api.exa.ai/search"
n_results: int = None
headers: dict = {
"accept": "application/json",
"content-type": "application/json",
"x-api-key": os.environ['EXA_API_KEY'],
}
def _parse_results(self, results):
stirng = []
for result in results:
try:
stirng.append('\n'.join([
f"Title: {result['title']}",
f"Score: {result['score']}",
f"Url: {result['url']}",
f"ID: {result['id']}",
"---"
]))
except KeyError:
next
content = '\n'.join(stirng)
return f"\nSearch results: {content}\n"

View File

@@ -0,0 +1,24 @@
import requests
from typing import Any
from .exa_base_tool import EXABaseTool
class EXAURLTool(EXABaseTool):
def _run(
self,
**kwargs: Any,
) -> Any:
search_query = kwargs.get('search_query')
if search_query is None:
search_query = kwargs.get('query')
payload = {
"query": search_query,
}
response = requests.post(self.search_url, json=payload, headers=self.headers)
results = response.json()
if 'results' in results:
results = super()._parse_results(results['results'])
else:
return results