agents.py

import math

from PIL.Image import Image
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.tools import tool
from pytesseract import pytesseract
from youtube_transcript_api import YouTubeTranscriptApi


@tool
def add_numbers(a: float, b: float) -> float:
    """
    Adds two numbers and returns the result.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a + b

@tool
def subtract_numbers(a: float, b: float) -> float:
    """
    Subtracts the second number from the first and returns the result.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a - b

@tool
def multiply_numbers(a: float, b: float) -> float:
    """
    Multiplies two numbers and returns the result.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a * b

@tool
def divide_numbers(a: float, b: float) -> float:
    """
    Divides the first number by the second. Raises error if dividing by zero.
    Args:
        a (float): the first number
        b (float): the second number
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def modulo_numbers(a: int, b: int) -> int:
    """
    Returns the remainder when the first number is divided by the second.
    Args:
        a (float): the first number
        b (float): the second number
    """
    if b == 0:
        raise ValueError("Cannot perform modulo by zero.")
    return a % b

@tool
def calculate_percentage(a: float, b: float) -> float:
    """
    Returns what percentage the first number is of the second number. Raises error if dividing by zero.
    Args:
        a (float): the first number
        b (float): the second number
    """
    if b == 0:
        raise ValueError("Whole cannot be zero.")
    return (a / b) * 100

@tool
def square_root(value: float) -> float:
    """
    Returns the square root of the given number.
    Args:
        value (float): the given number
    """
    if value < 0:
        raise ValueError("Cannot take the square root of a negative number.")
    return math.sqrt(value)

@tool
def logarithm(value: float, base: float = 10) -> float:
    """
    Returns the logarithm of 'value' with the specified base (default: 10).
    Args:
        value (float): the given number
        base (float): the specified base
    """
    if value <= 0 or base <= 0:
        raise ValueError("Logarithm undefined for non-positive values.")
    return math.log(value, base)


@tool
def search_wikipedia(query: str) -> str:
    """
    Searches Wikipedia for a given query and returns a short summary of the first matching article.
    Args:
        query (str): the given query
    """
    try:
        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
        output = []
        for result in search_docs:
            output.append(
                f"Title: {result.metadata['title']}\n"
                f"Summary: {result.metadata['summary'][:1000]}...\n"
                f"Source: {result.metadata['source']}"
            )
        return "\n\n".join(output)
    except Exception as e:
        return f"Wikipedia error: {e}"

@tool
def read_text_from_image(image_path: str) -> str:
    """Extracts text from an image file using OCR."""
    try:
        img = Image.open(image_path)
        text = pytesseract.image_to_string(img)
        return text.strip()
    except Exception as e:
        return f"Error processing image: {e}"

@tool
def get_youtube_transcript(video_id: str) -> str:
    """
    Retrieves and returns the transcript of a YouTube video by video ID.
    Args:
        video_id (str): the video ID
    """
    try:
        transcript = YouTubeTranscriptApi().fetch(video_id)
        full_text = '\n'.join([f"{entry.text}" for entry in transcript.snippets])
        return full_text
    except Exception as e:
        return f"Transcript error: {e}"

@tool
def search_arxiv(query: str, max_results: int = 3) -> str:
    """
    Searches arXiv and returns titles and summaries of top papers.
    Args:
        :param query: the query to search
        :param max_results: the number of top titles and summaries to return
    """
    try:
        results = ArxivLoader(query=query, load_max_docs=3).load()
        output = []
        for result in results:
            output.append(
                f"Authors: {result.metadata['Authors']}\n"
                f"Title: {result.metadata['Title']}\n"
                f"Published: {result.metadata['Published']}\n"
                f"Summary: {result.metadata['Summary'][:1000]}..."
            )
        return "\n\n".join(output)
    except Exception as e:
        return f"arXiv error: {e}"


def tool_list():
    web_search_tool = DuckDuckGoSearchRun()
    return [
        web_search_tool,
        add_numbers,
        subtract_numbers,
        multiply_numbers,
        divide_numbers,
        modulo_numbers,
        calculate_percentage,
        square_root,
        logarithm,
        search_wikipedia,
        get_youtube_transcript,
        search_arxiv,
        read_text_from_image,
    ]

# if __name__ == "__main__":
#     # print("\nArxiv Search:\n", search_arxiv.invoke({"query": "Reinforcement learning forex", "max_results": 2}))
#     # print("\nWikipedia Search:\n", search_wikipedia.invoke({"query": "Alexander the Great"}))
#     # print("\nYouTube transcript:\n", get_youtube_transcript.invoke({"video_id": "1htKBjuUWec"}))
#     print("\nText from Image:\n", read_text_from_image.invoke({"image_path": "1htKBjuUWec"}))