tools/search_tools.py

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