from smolagents import Tool
import wikipedia
from bs4 import BeautifulSoup

class WikipediaSearch(Tool):
    name = "wikipedia_search"
    description = "Fetches wikipedia pages."
    inputs = {
        "query": {
            "type": "string",
            "description": "Query to be searched on wikipedia"
        }
        }
    output_type = "string"

    def forward(self, query:str)->str:
        try:
            res = wikipedia.page(query)
            bs = BeautifulSoup(res.html(), 'html.parser')
            text_only = bs.get_text()
            return text_only
        except Exception as e:
            return f'Error in wikipedia search: {str(e)}'