from langchain_community.document_loaders import WikipediaLoader
from smolagents import Tool

class WikipediaTool(Tool):
    name = "WikipediaTool"
    description = "This tool fetches documents from Wikipedia"

    inputs = {
        "query": {
            "type": "string",
            "description": "the query to search in Wikipedia",
        }
    }

    output_type = "string"

    def forward(self, query: str) -> str:
        """
        Fetch documents from Wikipedia
        https://python.langchain.com/docs/integrations/document_loaders/wikipedia/
        """
        documents = WikipediaLoader(query=query, load_max_docs=1).load()
        return documents[0].page_content