tools/search.py

import requests
from dotenv import load_dotenv
import os
from .tool import Tool

load_dotenv("./.env")

class SearchTool(Tool):
    def __init__(self):
        super().__init__(
            name="search",
            description="Search the web for information",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "The search query"}
                }
            }
        )
        
        self.api_key = os.environ.get("GOOGLE_API_KEY")
        self.search_engine_id = os.environ.get("GOOGLE_CSE_ID")
        
        if not self.api_key:
            raise ValueError("Please set GOOGLE_API_KEY environment variable")
        if not self.search_engine_id:
            raise ValueError("Please set GOOGLE_CSE_ID environment variable")
    
    def __call__(self, query: str):
        try:
            if not query:
                return "Error: Query parameter is required"
                
            params = {
                "q": query,
                "key": self.api_key,
                "cx": self.search_engine_id
            }
            
            resp = requests.get("https://www.googleapis.com/customsearch/v1", params=params)
            resp.raise_for_status()  # Raise an exception for bad status codes
            
            _results = resp.json().get("items", [])
            results = []
            for result in _results[:3]:
                results.append({
                    "title": result.get("title", "No title"),
                    "link": result.get("link", "No link"),
                    "snippet": result.get("snippet", "No snippet")
                })
            
            if not results:
                return "No results found for the given query."
            
            # Format results as a string
            formatted_results = []
            for i, result in enumerate(results, 1):
                formatted_results.append(f"Result {i}:\nTitle: {result['title']}\nLink: {result['link']}\nSnippet: {result['snippet']}\n")
            
            return "\n".join(formatted_results)
            
        except requests.exceptions.RequestException as e:
            return f"Error during search: {str(e)}"
        except Exception as e:
            return f"Unexpected error during search: {str(e)}"