Spaces:
Sleeping
Sleeping
File size: 5,194 Bytes
d445f2a |
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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
import os
import requests
import time
import hashlib
from typing import Dict, Any, Optional, Type
from crewai.tools import BaseTool
from pydantic import BaseModel, Field
class TavilySearchArgs(BaseModel):
"""Input schema for TavilySearchTool."""
query: str = Field(..., description="The search query to look up")
class TavilySearchTool(BaseTool):
"""
Tool for performing web searches using the Tavily Search API.
This tool sends a search query to Tavily and returns relevant search results.
"""
name: str = Field(
default="Tavily Web Search",
description="Search the internet using Tavily"
)
description: str = Field(
default="Use this tool to search for information on the internet using Tavily Search API.",
description="Description of the Tavily search tool"
)
api_key: Optional[str] = Field(
default=None,
description="Tavily API key. If not provided, will look for TAVILY_API_KEY environment variable"
)
search_depth: str = Field(
default="basic",
description="The depth of the search, 'basic' or 'advanced'"
)
max_results: int = Field(
default=5,
description="Maximum number of search results to return (1-10)"
)
include_answer: bool = Field(
default=False,
description="Whether to include an AI-generated answer in the response"
)
timeout: int = Field(
default=10,
description="Timeout for the API request in seconds"
)
args_schema: Type[BaseModel] = TavilySearchArgs
def __init__(self, **data):
super().__init__(**data)
self.api_key = self.api_key or os.getenv("TAVILY_API_KEY")
if not self.api_key:
print("WARNING: Tavily API key is missing. The tool will return an error message when used.")
self._cache = {} # Simple in-memory cache
def _run(self, query: str) -> str:
"""
Execute a web search using Tavily.
Args:
query: The search query to look up
Returns:
String containing the search results
"""
# Check if API key is missing
if not self.api_key:
return (
"ERROR: Tavily API key is missing. Please set the TAVILY_API_KEY environment variable. "
"Search cannot be performed without a valid API key."
)
# Check cache first
cache_key = self._get_cache_key(query)
if cache_key in self._cache:
timestamp, result = self._cache[cache_key]
# Cache valid for 30 minutes
if time.time() - timestamp < 1800:
return f"{result}\n\n[Cached Tavily result]"
url = "https://api.tavily.com/search"
payload = {
"api_key": self.api_key,
"query": query,
"search_depth": self.search_depth,
"max_results": min(self.max_results, 10), # Ensure we don't exceed API limits
"include_answer": self.include_answer
}
try:
response = requests.post(url, json=payload, timeout=self.timeout)
response.raise_for_status()
result = response.json()
if "results" not in result:
return f"Error in search: {result.get('error', 'Unknown error')}"
# Format the results
formatted_results = self._format_results(result)
# Cache the result
self._cache[cache_key] = (time.time(), formatted_results)
return formatted_results
except requests.exceptions.Timeout:
return "Error: Tavily search request timed out. Please try again later."
except requests.exceptions.RequestException as e:
return f"Error during Tavily search: {str(e)}"
def _format_results(self, result: Dict[str, Any]) -> str:
"""Format the search results into a readable string."""
output = []
# Add the answer if included
if "answer" in result and result["answer"]:
output.append(f"Answer: {result['answer']}\n")
# Add search results
output.append("Search Results:")
for i, r in enumerate(result.get("results", []), 1):
title = r.get("title", "No Title")
url = r.get("url", "No URL")
content = r.get("content", "No Content").strip()
result_text = f"\n{i}. {title}\n URL: {url}\n Content: {content}\n"
output.append(result_text)
return "\n".join(output)
def _get_cache_key(self, query: str) -> str:
"""Generate a cache key for the given query."""
# Include search parameters in the key
params_str = f"{query}|{self.search_depth}|{self.max_results}|{self.include_answer}"
return hashlib.md5(params_str.encode()).hexdigest() |