import datasets
import torch
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.retrievers import WikipediaRetriever
from langchain_community.document_loaders import ArxivLoader, WikipediaLoader
from langchain_core.documents import Document
from langchain_core.tools import tool
from langchain_tavily import TavilySearch
from sentence_transformers import SentenceTransformer



search_tool = TavilySearch(
    max_results=5,
    topic="general",
    # include_answer=False,
    # include_raw_content=False,
    # include_images=False,
    # include_image_descriptions=False,
    # search_depth="basic",
    # time_range="day",
    # include_domains=None,
    # exclude_domains=None
)

@tool
def wikipedia_search(query: str) -> str:
    """Retrieves information for a question from wikipedia.

    Args:
        query: The query for which the information is to be retrieved.

    Returns:
        The retrieved information from wikipedia.
    """
    loader = WikipediaLoader(query=query)
    results = loader.load()
    if results:
        return "\n\n".join([doc.page_content for doc in results])
    else:
        return "No relevant information found."


@tool
def arxiv_search(query: str) -> str:
    """Retrieves information for a question from arxiv papers.

    Args:
        query: The query for which the information is to be searched.

    Returns:
        The retrieved information from arxiv.
    """
    loader = ArxivLoader(query=query)
    top_results = loader.load()
    if len(top_results) > 0:
        return "\n\n".join([doc.page_content for doc in top_results])
    else:
        return "No relevant information found."

@tool
def divide(a: float, b:float) -> str:
    """Divides two numbers.
    Args:
        a: The dividend.
        b: The divisor.

    Returns:
        The result of the division.
    """
    if b == 0:
        return "Error: Division by zero is not allowed."
    return str(a / b)

@tool
def multiply(a: float, b:float) -> str:
    """Multiplies two numbers.
    Args:
        a: The first number.
        b: The second number.

    Returns:
        The result of the multiplication.
    """
    return str(a * b)

@tool
def add(a: float, b:float) -> str:
    """Adds two numbers.
    Args:
        a: The first number.
        b: The second number.

    Returns:
        The result of the addition.
    """
    return str(a + b)

@tool
def subtract(a: float, b:float) -> str:
    """Subtracts two numbers.
    Args:
        a: The first number.
        b: The second number.

    Returns:
        The result of the subtraction.
    """
    return str(a - b)

@tool
def modulus(a: int, b:int) -> str:
    """Calculates the modulus of two numbers.
    Args:
        a: The first number.
        b: The second number.

    Returns:
        The result of the modulus operation.
    """
    if b == 0:
        return "Error: Modulus by zero is not allowed."
    return str(a % b)