Spaces:
Runtime error
Runtime error
File size: 3,888 Bytes
105b369 |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import json
from os import getenv
from typing import Optional, Literal, Dict, Any
from phi.tools import Toolkit
from phi.utils.log import logger
try:
from tavily import TavilyClient
except ImportError:
raise ImportError("`tavily-python` not installed. Please install using `pip install tavily-python`")
class TavilyTools(Toolkit):
def __init__(
self,
api_key: Optional[str] = None,
search: bool = True,
max_tokens: int = 6000,
include_answer: bool = True,
search_depth: Literal["basic", "advanced"] = "advanced",
format: Literal["json", "markdown"] = "markdown",
use_search_context: bool = False,
):
super().__init__(name="tavily_tools")
self.api_key = api_key or getenv("TAVILY_API_KEY")
if not self.api_key:
logger.error("TAVILY_API_KEY not provided")
self.client: TavilyClient = TavilyClient(api_key=self.api_key)
self.search_depth: Literal["basic", "advanced"] = search_depth
self.max_tokens: int = max_tokens
self.include_answer: bool = include_answer
self.format: Literal["json", "markdown"] = format
if search:
if use_search_context:
self.register(self.web_search_with_tavily)
else:
self.register(self.web_search_using_tavily)
def web_search_using_tavily(self, query: str, max_results: int = 5) -> str:
"""Use this function to search the web for a given query.
This function uses the Tavily API to provide realtime online information about the query.
Args:
query (str): Query to search for.
max_results (int): Maximum number of results to return. Defaults to 5.
Returns:
str: JSON string of results related to the query.
"""
response = self.client.search(
query=query, search_depth=self.search_depth, include_answer=self.include_answer, max_results=max_results
)
clean_response: Dict[str, Any] = {"query": query}
if "answer" in response:
clean_response["answer"] = response["answer"]
clean_results = []
current_token_count = len(json.dumps(clean_response))
for result in response.get("results", []):
_result = {
"title": result["title"],
"url": result["url"],
"content": result["content"],
"score": result["score"],
}
current_token_count += len(json.dumps(_result))
if current_token_count > self.max_tokens:
break
clean_results.append(_result)
clean_response["results"] = clean_results
if self.format == "json":
return json.dumps(clean_response) if clean_response else "No results found."
elif self.format == "markdown":
_markdown = ""
_markdown += f"# {query}\n\n"
if "answer" in clean_response:
_markdown += "### Summary\n"
_markdown += f"{clean_response.get('answer')}\n\n"
for result in clean_response["results"]:
_markdown += f"### [{result['title']}]({result['url']})\n"
_markdown += f"{result['content']}\n\n"
return _markdown
def web_search_with_tavily(self, query: str) -> str:
"""Use this function to search the web for a given query.
This function uses the Tavily API to provide realtime online information about the query.
Args:
query (str): Query to search for.
Returns:
str: JSON string of results related to the query.
"""
return self.client.get_search_context(
query=query, search_depth=self.search_depth, max_tokens=self.max_tokens, include_answer=self.include_answer
)
|