# Wikipedia search tool
import requests
from smolagents import Tool

class WikipediaTool(Tool):
    name = "wikipedia_search"
    description = "Search Wikipedia for information about a topic."
    inputs = {
        "query": {
            "type": "string",
            "description": "The search query"
        }
    }
    output_type = "string"
    
    def forward(self, query: str) -> str:
        assert isinstance(query, str), "Query must be a string"
        try:
            search_url = f"https://en.wikipedia.org/w/api.php?action=query&list=search&srsearch={query}&format=json"
            search_response = requests.get(search_url, timeout=10)
            search_data = search_response.json()
            
            if "query" not in search_data or "search" not in search_data["query"] or not search_data["query"]["search"]:
                return f"No Wikipedia results found for {query}"
                
            # Get the first result
            first_result = search_data["query"]["search"][0]
            page_id = first_result["pageid"]
            
            # Get the page content
            content_url = f"https://en.wikipedia.org/w/api.php?action=query&prop=extracts&exintro&explaintext&pageids={page_id}&format=json"
            content_response = requests.get(content_url, timeout=10)
            content_data = content_response.json()
            
            extract = content_data["query"]["pages"][str(page_id)]["extract"]
            title = content_data["query"]["pages"][str(page_id)]["title"]
            
            return f"""Wikipedia: {title}

{extract[:1500]}... [content truncated]

Source: https://en.wikipedia.org/wiki/{title.replace(' ', '_')}
"""
        except Exception as e:
            print(f"Error searching Wikipedia: {str(e)}")
            return f"Error searching Wikipedia for {query}: {str(e)}"