tools.py

import os
import random

from typing import Dict
from pathlib import Path

from langchain_core.tools import tool

from langchain_core.messages import ToolMessage
from langchain_tavily import TavilySearch
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader

@tool
def web_search(query: str) -> ToolMessage:
    """Search in the web with Tavily for a query and return maximum 5 results.
    Args:
        query: The search query.
    Returns:
        Tavily output, and snippet for the top 5 results
    """
    return TavilySearch(max_results=5, include_images=False).invoke({"query": query})


@tool
def wikipedia_search(query: str) -> Dict[str, list]:
    """Search Wikipedia for a given query and return the first 5 results.
    Args:
        query: The search term or topic.
    Returns:
        A dictionary containing the formatted Wikipedia results.
    """
    search_docs = WikipediaLoader(query=query, load_max_docs=5).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 {"wiki_results": formatted_search_docs}


#Mathematical tools
@tool
def multiply(a: float, b: float) -> float:
    """Multiply two numbers.
    Args:
        a: first number
        b: second number
    Returns: 
        Multiplication result
    """
    return a * b


@tool
def add(a: float, b: float) -> float:
    """Add two numbers.
    Args:
        a: first number
        b: second number
    Returns: 
        Addition result
    """
    return a + b


@tool
def subtract(a: float, b: float) -> float:
    """Subtract two numbers.
    Args:
        a: first number
        b: second number
    Returns: 
        Subtraction result
    """
    return a - b


@tool
def divide(a: float, b: float) -> float:
    """Divide two numbers.
    Args:
        a: first number
        b: second number
    Returns: 
        Division result
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b


@tool
def modulus(a: int, b: int) -> int:
    """Get the modulus of two numbers.
    Args:
        a: first number
        b: second number
    Returns:
        Modulus result
    """
    return a % b

from langchain_core.tools import tool


@tool
def convert_units(value: float, from_unit: str, to_unit: str) -> float:
    """
    Converts a value from one unit to another.

    Args:
        value: The numerical value to convert.
        from_unit: The original unit (e.g. 'miles', 'kg', 'celsius').
        to_unit: The target unit (e.g. 'kilometers', 'lb', 'fahrenheit').

    Supported conversions:
        - miles <-> kilometers
        - kilograms <-> pounds
        - celsius <-> fahrenheit

    Returns:
        The converted value result.
    """
    conversions = {
        ("miles", "kilometers"): lambda v: v * 1.60934,
        ("kilometers", "miles"): lambda v: v / 1.60934,
        ("kilograms", "pounds"): lambda v: v * 2.20462,
        ("pounds", "kilograms"): lambda v: v / 2.20462,
        ("celsius", "fahrenheit"): lambda v: (v * 9/5) + 32,
        ("fahrenheit", "celsius"): lambda v: (v - 32) * 5/9,
    }

    key = (from_unit.lower(), to_unit.lower())
    if key not in conversions:
        raise ValueError(f"Conversion from {from_unit} to {to_unit} not supported.")

    return conversions[key](value)

@tool
def query_table_data(file_path: str, query: str, sheet_name: str = None) -> str:
    """
    Loads a table from CSV or Excel and filters it using a pandas query.

    Args:
        file_path: Path to the table file (.xlsx, .xls).
        query: A pandas-compatible query string, e.g., "Age > 30 and Country == 'USA'".
        sheet_name: Optional sheet name if the file is Excel.

    Returns:
        A string representation (markdown) of the filtered table (max 10 rows).
    """
    try:
        import pandas as pd
        path = Path(file_path)
        if not path.exists():
           raise FileNotFoundError(f"File not found: {file_path}")
        ext = path.suffix.lower()
        if ext == ".csv":
           df = pd.read_csv(path)
        elif ext in [".xlsx", ".xls"]:
           df = pd.read_excel(path, sheet_name=sheet_name)
        else:
           raise ValueError(f"Unsupported file extension: {ext}")
        try:
           filtered_df = df.query(query)
           return filtered_df.head(10).to_markdown(index=False)
        except Exception as e:
           raise ValueError(f"Invalid query: {query}. Error: {e}")
    except ImportError:
        return "Error: pandas and openpyxl are not installed. Please install them with 'pip install pandas openpyxl'."
    

@tool
def arvix_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 5 result.
    Args:
     query: The search query.
    Returns:
       A dictionary containing the formatted Arvix results, and snippet for the top 5 results.
    """
    search_docs = ArxivLoader(query=query, load_max_docs=5).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ])
    return {"arvix_results": formatted_search_docs}


level1_tools = [
    multiply,
    add,
    subtract,
    divide,
    modulus,
    wikipedia_search,
    web_search,
    arvix_search,
    convert_units,
    query_table_data
]