tools/wikipedia.py

from typing import Any
from smolagents.tools import Tool
from langchain_community.document_loaders import WikipediaLoader

class WikiSearchTool(Tool):
    name = "wiki_search"
    description = (
        "Search Wikipedia for a query and return up to 2 results formatted as XML-like Document tags."
    )
    inputs = {
        "query": {"type": "string", "description": "The search query."}
    }
    output_type = "any"

    def __init__(self, *args, **kwargs):
        self.is_initialized = False

    def forward(self, query: str) -> Any:
        # load at most 2 documents from Wikipedia
        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
        # format each as an XML-like fragment and join with separators
        formatted_search_docs = "\n\n---\n\n".join([
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n'
            f'{doc.page_content}\n'
            f'</Document>'
            for doc in search_docs
        ])
        return {"wiki_results": formatted_search_docs}