Spaces:
Sleeping
Sleeping
File size: 5,298 Bytes
ce0ec3b 7539685 ce0ec3b 7539685 ce0ec3b 7539685 ce0ec3b 0e9bb01 ce0ec3b 0e9bb01 ce0ec3b 0e9bb01 ce0ec3b |
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 |
"""
Search tools for finding song lyrics and related information.
"""
import os
import random
import time
from typing import Any, Dict, List
import requests
from loguru import logger
from smolagents import DuckDuckGoSearchTool, Tool
class BraveSearchTool(Tool):
"""
A tool for performing web searches using the Brave Search API.
This tool requires a Brave Search API key to be set in the environment
variable BRAVE_API_KEY or passed directly to the constructor.
Documentation: https://api.search.brave.com/app/documentation
"""
def __init__(self, max_results: int = 10, **kwargs):
"""
Initialize the Brave Search tool.
Args:
max_results: Maximum number of results to return (default: 10)
"""
super().__init__(**kwargs)
self.api_key = os.environ.get("BRAVE_API_KEY")
if not self.api_key:
logger.warning("No Brave API key found. Set BRAVE_API_KEY environment variable or pass api_key parameter.")
self.max_results = max_results
self.name = "brave_search"
self.description = "Search the web using Brave Search API"
self.inputs = {"query": {"type": "string", "description": "The search query string"}}
self.output_type = "string"
logger.info(f"Initialized BraveSearchTool with max_results={max_results}")
def forward(self, query: str) -> List[Dict[str, Any]]:
"""
Execute a search using the Brave Search API.
Args:
query: The search query string
Returns:
List of search results in the format:
[{"title": str, "href": str, "body": str}, ...]
"""
if not self.api_key:
logger.error("Brave Search API key is not set")
return [{"title": "API Key Error", "href": "", "body": "Brave Search API key is not set"}]
url = "https://api.search.brave.com/res/v1/web/search"
headers = {"Accept": "application/json", "X-Subscription-Token": self.api_key}
params = {"q": query, "count": self.max_results}
try:
logger.info(f"Performing Brave search for query: '{query}'")
response = requests.get(url, headers=headers, params=params)
response.raise_for_status()
data = response.json()
results = []
if "web" in data and "results" in data["web"]:
for result in data["web"]["results"]:
results.append({
"title": result.get("title", ""),
"href": result.get("url", ""),
"body": result.get("description", "")
})
logger.info(f"Found {len(results)} results for query: '{query}'")
return results
except Exception as e:
logger.error(f"Error in Brave search: {str(e)}")
return [{"title": "Search error", "href": "", "body": f"Error performing search: {str(e)}"}]
class ThrottledDuckDuckGoSearchTool(DuckDuckGoSearchTool):
"""
A wrapper around DuckDuckGoSearchTool that adds a delay between requests
to avoid rate limiting issues.
This tool implements a delay mechanism to prevent hitting DuckDuckGo's rate limits.
Each search request will be followed by a random delay within the specified range.
"""
def __init__(self, min_delay: float = 7.0, max_delay: float = 15.0, **kwargs):
"""
Initialize the throttled search tool with delay parameters.
Args:
min_delay: Minimum delay in seconds between requests (default: 5.0)
max_delay: Maximum delay in seconds between requests (default: 5.0)
**kwargs: Additional arguments to pass to DuckDuckGoSearchTool
"""
super().__init__(**kwargs)
self.min_delay = min_delay
self.max_delay = max_delay
self.name = "search" # Keep the same name as the parent class
logger.info(f"Initialized ThrottledDuckDuckGoSearchTool with delay range: {min_delay}-{max_delay}s")
def forward(self, query: str) -> List[Dict[str, Any]]:
"""
Execute a search with a delay to avoid rate limiting.
Args:
query: The search query string
Returns:
List of search results
"""
# Add a random delay before the search to avoid rate limiting
delay = random.uniform(self.min_delay, self.max_delay)
logger.info(f"Throttling DuckDuckGo search for {delay:.2f} seconds before query: '{query}'")
time.sleep(delay)
# Call the parent class implementation
try:
results = super().forward(query)
# Add another delay after the search to ensure spacing between requests
time.sleep(random.uniform(self.min_delay / 2, self.max_delay / 2))
return results
except Exception as e:
logger.error(f"Error in DuckDuckGo search: {str(e)}")
# Return empty results on error to allow the agent to continue
return [{"title": "Search error", "href": "", "body": f"Error performing search: {str(e)}"}]
|