# this file contains the tools that the agent will use
from smolagents import tool
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import ArxivLoader

# wikipedia page loader
@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 results.
    
    This tool retrieves information from Wikipedia articles related to your query.
    It returns up to 2 full Wikipedia articles with their content and metadata.
    Use this tool when you need factual information, definitions, historical context,
    or general knowledge that would be found in an encyclopedia.
    
    Args:
        query: The search query."""
    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return formatted_search_docs

@tool
def web_search(query: str) -> str:
    """Search the web using Tavily search engine and return up to 3 relevant results.
    This tool provides access to current information from the internet, including news, articles, and websites.
    Each result includes the title, URL, and content snippet from the webpage.
    Use this tool when you need to find up-to-date information that might not be available in other sources.
    
    Args:
        query: The search query."""
    
    
    search_docs = TavilySearchResults(max_results=3).invoke(input=query)
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document title="{doc["title"]}" \n url="{doc["url"]}"/>\n{doc["content"]}\n</Document>'
            for doc in search_docs
        ])
    return formatted_search_docs

@tool
def arvix_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 results.
    
    This tool retrieves scientific papers and research articles from Arxiv based on your query.
    It returns up to 3 papers with their publishing date, title, authors, summary, and the first 1000 characters of content.
    Use this tool when you need academic research, scientific information, technical papers, or the latest findings in fields like physics, mathematics, computer science, and other scientific disciplines. You should use this tool when looking for files on the arxiv website. Arxiv URLs typically look like "https://arxiv.org/abs/XXXX.XXXXX" where X represents numbers.
    Args:
        query: The search query."""
    
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document Publishing Date="{doc.metadata["Published"]}" \n Title="{doc.metadata["Title"]}" \n Authors="{doc.metadata["Authors"]}" \n Summary="{doc.metadata["Summary"]}" \n />\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ])
    return formatted_search_docs