tools.py

import math

from langchain.tools import tool
from langchain_community.document_loaders import WikipediaLoader
from langchain_tavily import TavilySearch

@tool
def wikipedia_search(query: str) -> str:
    """Search Wikipedia for a given query and return results.
    
    Args:
        query (str): The search query, should be the main subject to search (e.g. person name).
    """
    max_docs: int=1; max_text_len: int=99999
    print(f"wikipedia_search: {query}, {max_docs}, {max_text_len}")
    loader = WikipediaLoader(query, load_max_docs=max_docs, doc_content_chars_max=max_text_len)
    documents = loader.load()
    return {"wikipedia_content": (documents[0].page_content if len(documents) > 0 else "")}

@tool
def web_search(query: str) -> str:
    """Search the web for a given query and return results.
    
    Args:
        query (str): The search query.
    """
    print(f"web_search: {query}")
    search_results = TavilySearch(max_results=3, include_images=False).invoke({"query": query})
    return {"web_results": search_results}

@tool
def add(a: float, b: float) -> float:
    """
    Adds two numbers.

    Args:
        a (float): the first number
        b (float): the second number
    """
    return a + b

@tool
def subtract(a: float, b: float) -> int:
    """
    Subtracts two numbers.

    Args:
        a (float): the first number
        b (float): the second number
    """
    return a - b

@tool
def multiply(a: float, b: float) -> float:
    """
    Multiplies two numbers.

    Args:
        a (float): the first number
        b (float): the second number
    """
    return a * b

@tool
def divide(a: float, b: float) -> float:
    """
    Divides two numbers.

    Args:
        a (float): the first float number
        b (float): the second float number
    """
    if b == 0:
        raise ValueError("Cannot divided by zero.")
    return a / b


@tool
def modulus(a: int, b: int) -> int:
    """
    Get the modulus of two numbers.

    Args:
        a (int): the first number
        b (int): the second number
    """
    return a % b


@tool
def power(a: float, b: float) -> float:
    """
    Get the power of two numbers.

    Args:
        a (float): the first number
        b (float): the second number
    """
    return a**b


@tool
def square_root(a: float) -> float | complex:
    """
    Get the square root of a number.

    Args:
        a (float): the number to get the square root of
    """
    if a >= 0:
        return a**0.5
    return math.sqrt(a)

my_tools = [
    wikipedia_search,
    web_search,
    add,
    subtract,
    multiply,
    divide,
    modulus,
    power,
    square_root
]

if __name__ == "__main__":
    result = web_search.invoke("Mercedes Sosa discography")
    print(result)

    #from main import main as main_fn
    #main_fn()