""" 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)}"}]