from smolagents import PythonInterpreterTool, GoogleSearchTool, VisitWebpageTool, WikipediaSearchTool, SpeechToTextTool
import wikipedia
from wikipedia import WikipediaPage
from smolagents import tool

@tool
def search_wikipedia_page(query: str) -> list[str]:
    """
    Search for Wikipedia pages.

    Args:
        query: The query to search for.

    Returns:
        A list of Wikipedia pages that match the query.
    """
    return wikipedia.search(query)

@tool
def get_wikipedia_page(title: str) -> WikipediaPage:
    """
    Get the content of a specific Wikipedia page.

    Args:
        title: The title of the Wikipedia page to get.

    Returns:
        A WikipediaPage object. The object includes the following attributes: 'content', 'summary', 'images', 'references', 'links', 'sections'
    """
    return wikipedia.page(title)

@tool
def get_wikipedia_summary(title: str) -> str:
    """
    Get a summary of a Wikipedia page.

    Args:
        title: The title of the Wikipedia page to get.

    Returns:
        The plain-text summary of the Wikipedia page.
    """
    return wikipedia.summary(title)

python_wikipedia_tools = [search_wikipedia_page, get_wikipedia_page, get_wikipedia_summary]

smolagents_code_tools = [PythonInterpreterTool()]
smolagents_web_tools = [GoogleSearchTool(), VisitWebpageTool(), WikipediaSearchTool()]
smolagents_speech_tools = [SpeechToTextTool()]