/
tavily.py
47 lines (33 loc) · 1.46 KB
/
tavily.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import os
from typing import Any, Dict, List
from langchain_community.tools.tavily_search import TavilySearchResults
from tavily import TavilyClient
from backend.tools.retrieval.base import BaseRetrieval
class TavilyInternetSearch(BaseRetrieval):
tavily_api_key = os.environ.get("TAVILY_API_KEY")
def __init__(self):
self.client = TavilyClient(api_key=self.tavily_api_key)
@classmethod
def is_available(cls) -> bool:
return cls.tavily_api_key is not None
def retrieve_documents(self, query: str, **kwargs: Any) -> List[Dict[str, Any]]:
content = self.client.search(query=query, search_depth="advanced")
if "results" not in content:
return []
return [
{
"url": result["url"],
"text": result["content"],
}
for result in content["results"]
]
def to_langchain_tool(self) -> TavilySearchResults:
internet_search = TavilySearchResults()
internet_search.name = "internet_search"
internet_search.description = "Returns a list of relevant document snippets for a textual query retrieved from the internet."
# pydantic v1 base model
from langchain_core.pydantic_v1 import BaseModel, Field
class TavilySearchInput(BaseModel):
query: str = Field(description="Query to search the internet with")
internet_search.args_schema = TavilySearchInput
return internet_search