tools.py

import base64
import math
import os
import re
import subprocess
from pathlib import Path
from tempfile import NamedTemporaryFile

import arxiv
import pandas as pd
import requests
import wikipedia
from bs4 import BeautifulSoup
from huggingface_hub import InferenceClient
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langchain_tavily import TavilySearch
from markdownify import markdownify as md
from PyPDF2 import PdfReader
from youtube_transcript_api import YouTubeTranscriptApi

# Global configration for the tools
wikipedia.set_user_agent("GaiaAgent (alexinhocora@gmail.com)")
tavily_search = TavilySearch(max_results=5, topic="general")
arxiv_client = arxiv.Client()
ytt_api = YouTubeTranscriptApi()


@tool
def calculator(expression: str) -> str:
    """Evaluate a mathematical expression.

    Args:
        expression: A string containing the mathematical expression to evaluate.
        examples:
            "2 + 2 * (3 - 1)"
            "sqrt(16) + 5"
            "sin(pi / 2) * 10"
    """
    try:
        # Source - https://stackoverflow.com/q/3513292
        # Posted by flybywire, modified by community. See post 'Timeline' for change history
        # Retrieved 2026-07-12, License - CC BY-SA 3.0
        safe_names = {
            name: getattr(math, name) for name in dir(math) if not name.startswith("_")
        }

        safe_names.update(
            {
                "abs": abs,
                "round": round,
                "min": min,
                "max": max,
                "int": int,
                "float": float,
                "complex": complex,
            }
        )

        result = eval(
            expression,
            {"__builtins__": {}},
            safe_names,
        )

        return str(result)

    except Exception as e:
        return f"Error: {str(e)}"


@tool
def wikipedia_search(query: str) -> str:
    """Retrieve a summary of the top 3 Wikipedia pages for a given search query.

    Args:
        query: Search key term to look up in wikipedia.
    """
    documents = []
    suggestions = wikipedia.search(query, results=3)
    for page in suggestions:
        try:
            page = wikipedia.page(page, auto_suggest=False)
            sections = page.sections
            if not sections:
                # Sometime fails to get the proper sections, backup plan is to parse the html and get the sections from the table of contents
                soup = BeautifulSoup(page.html(), "html.parser")
                sections = [
                    a.attrs.get("href", "").strip("#")
                    for a in soup.select("#toc li > a")
                ]
            documents.append(f"""# {page.title}
**URL:** {page.url}\\
**Sections:** {",".join(sections)}

{page.summary.strip()}""")

        except Exception as e:
            print(f"Error retrieving page for {page}: {e}")
            continue

    formatted_docs = "\n\n---\n\n".join(documents)

    return formatted_docs


@tool
def wikipedia_section(
    title: str, section: str | None = None, max_chars: int | None = None
) -> str:
    """Retrieve current Wikipedia content.
    You must provide the exact title of the page and a valid section name to retrieve.
    Always use this tool instead of `fetch_webpage` for Wikipedia pages.
    This tool retrieves the current article revision. It must not be used as
    historical evidence for questions containing "as of" or a past date.

    Example usage:
        # Retrieve the "Discography" section of the "Mercedes Sosa" Wikipedia page
        wikipedia_section.invoke({
            "title": "Mercedes Sosa",
            "section": "Discography"
        })
        # Retrive the entire content of the "Mercedes Sosa" Wikipedia page
        wikipedia_section.invoke({
            "title": "Mercedes Sosa",
        })
        # Return the first 500 characters of the Python (programming language) Wikipedia page
        wikipedia_section.invoke({
            "title": "Python (programming language)",
            "max_chars": 500
        })

    Args:
        title: The exact title of the wikipedia page to retrieve. For example, "Mercedes Sosa" or "Python (programming language)".
        section: The exact section title to retrieve from the page. For example, "Discography" or "Early life". Use None to retrieve the beginning of the entire article.
        max_chars: The maximum number of characters to retrieve from the section content. None to retrieves from the start of the section to the end of the content.
    """
    # page.section fails to retrieve the section content when the section includes tables or other complex HTML elments
    # So we will retrieve the page HTML and parse the content
    page = wikipedia.page(title, auto_suggest=False)
    soup = BeautifulSoup(page.html(), "html.parser")

    # Clean the HTML content and convert it to Markdown
    for tag in [
        "script",
        "style",
        "noscript",
        "nav",
        "footer",
        "header",
        "svg",
        "img",
        "figure",
        "#toc",
        "span.mw-editsection",
    ]:
        for t in soup.select(tag):
            t.decompose()

        page_md = md(
            soup.prettify(),
            heading_style="ATX",
            bullets="-",
            strip=["a", "img"],
        )

    # Avoid llm stress with unnecesary token consume data
    if (i := page_md.find("## References")) > 0:
        page_md = page_md[:i]

    if section is None:
        if max_chars is not None:
            content = page_md[:max_chars]
        else:
            content = page_md

        return f"""# {title}
**URL:** {page.url}

{content.strip()}"""

    # Look for the section in the markdown content
    if (
        m := re.search(r"#{2,4} " + re.escape(section), page_md, re.IGNORECASE)
    ) is None:
        return f"Section '{section}' not found in the page '{title}', you must provide a valid section name."

    # Returns the max_chars characters of the section content
    if max_chars is not None:
        content = page_md[m.start() : m.end() + max_chars]
    else:
        content = page_md[m.start() :]

    return f"""# {title}
**URL:** {page.url}
**Section:** {section}

{content.strip()}"""


@tool
def arxiv_search(query: str) -> str:
    """Retrieve the top 5 relevant arxiv papers for a given search query.
    Use the exact paper title whenever it is known.
    Return the canonical title, authors, arXiv ID,
    abstract URL, and direct PDF URL.

    Args:
        query: Search key term to look up in arxiv.
    """
    search = arxiv.Search(
        query=query, max_results=5, sort_by=arxiv.SortCriterion.Relevance
    )

    documents = [
        f"""# {r.title}

**URL:** {r.entry_id}\\
**PDF:** {r.pdf_url}\\
**Authors:** {",".join(a.name for a in r.authors)}\\
**Published:** {r.published.date()}

{r.summary.strip()}"""
        for r in arxiv_client.results(search)
    ]

    formatted_docs = "\n\n---\n\n".join(documents)
    return formatted_docs


@tool
def web_search(query: str) -> str:
    """Retrieve the top relevant web search results for a given query.

    Args:
        query: Search key term to look for
    """

    data = tavily_search.invoke({"query": query})
    search_docs = data.get("results", [])

    formatted_docs = "\n\n---\n\n".join(
        [
            f"""# {doc["title"]}
**URL:** {doc["url"]}\\
**SCORE:** {doc["score"]:.2f} (Relevance score from 0 to 1, higher is better)

{doc["content"].strip()}"""
            for doc in search_docs
        ]
    )

    return formatted_docs


@tool
def fetch_webpage(url: str, max_chars: int | None = None) -> str:
    """Fetch a webpage content.
    DON'T USE THIS TOOL to try to fetch Wikipedia or arxiv content, use the dedicated tools for those sources `wikipedia_section` and `arxiv_search`.


    Args:
        url: Webpage URL to fetch
        max_chars: Maximum number of characters to return from the webpage content. By default, returns the entire content.
    """
    try:
        response = requests.get(url)
        response.raise_for_status()
        soup = BeautifulSoup(response.text, "html.parser")
        for tag in [
            "script",
            "style",
            "noscript",
            "nav",
            "link",
            "footer",
            "header",
            "iframe",
            "svg",
            "img",
            "figure",
        ]:
            for t in soup.select(tag):
                t.decompose()

        page_md = md(
            soup.prettify(),
            heading_style="ATX",
            bullets="-",
            strip=["a", "img"],
        )

        if max_chars is not None:
            return page_md[:max_chars]
        return page_md

    except requests.exceptions.RequestException as e:
        return f"Failed to fetch the webpage. Error: {str(e)}."


@tool
def read_excel_file(file_path: str, sheet_name: str = "Sheet0") -> str:
    """Read an Excel file and return its content as a string.

    Args:
        file_path: Path to the Excel file.
        sheet_name: Name of the sheet to read from the Excel file.
    """
    try:
        excel_file = pd.ExcelFile(file_path)
        sheet_name = (
            sheet_name
            if sheet_name in excel_file.sheet_names
            else excel_file.sheet_names[0]
        )
        df = pd.read_excel(excel_file, sheet_name=sheet_name)
        return df.to_string(index=False)

    except Exception as e:
        return f"Failed to read the Excel file. Error: {str(e)}"


@tool
def get_youtube_transcript(video_url: str) -> str:
    """Retrieve the transcription of a YouTube video.

    Args:
        video_url: The URL of the YouTube video.
    """
    try:
        # Extract the video ID from the URL
        video_id = re.search(
            r"https?://www.youtube.com/watch\?v=(?P<video_id>[a-zA-Z0-9_-]+)", video_url
        )
        if not video_id:
            return "Invalid YouTube URL. Please provide a valid URL."
        video_id = video_id.group("video_id")

        transcript = ytt_api.fetch(video_id)
        return f"# YouTube video transcript:\n{'=' * 50}\n" + "\n".join(
            [
                f" - [{f.start:05.2f}:{(f.start + f.duration):05.2f}] {f.text}"
                for f in transcript
            ]
        )

    except Exception as e:
        return f"An error occurred: {e}"


@tool
def execute_python_file(file_path: str) -> str:
    """Execute a Python file and return the output.

    Args:
        file_path: The path to the Python file to execute.
    """

    try:
        result = subprocess.run(["python", file_path], check=True, capture_output=True)
        return result.stdout.decode()
    except subprocess.CalledProcessError as e:
        return f"Error executing subprocess: {e}"


@tool
def execute_python_code(code: str) -> str:
    """Execute a Python code snippet and return the output.

    Args:
        code: The Python code to execute.
    """
    try:
        result = subprocess.run(["python", "-c", code], check=True, capture_output=True)
        return result.stdout.decode()
    except subprocess.CalledProcessError as e:
        return f"Error executing subprocess: {e}"


@tool
def transcript_audio_file(file_path: str) -> str:
    """Transcribe an audio file using OpenAI's Whisper model.

    Args:
        file_path: The path to the audio file to transcribe.
    """
    try:
        client = InferenceClient(
            provider="hf-inference",
            api_key=os.environ["HF_TOKEN"],
        )

        output = client.automatic_speech_recognition(
            file_path, model="openai/whisper-large-v3-turbo"
        )
        return output.text

    except Exception as e:
        return str(e)


@tool
def describe_image_file(file_path: str) -> str:
    """Describe an image file using OpenAI's CLIP model.

    Args:
        file_path: The path to the image file to describe.
    """
    try:
        extension = Path(file_path).suffix.lower()

        mime_types = {
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".png": "image/png",
            ".gif": "image/gif",
            ".bmp": "image/bmp",
            ".tiff": "image/tiff",
        }

        if extension not in mime_types:
            return f"Unsupported file type: {extension}. Supported types are: {', '.join(mime_types.keys())}"

        mime_type = mime_types[extension]
        with open(file_path, "rb") as f:
            image_bytes = f.read()
            image_base64 = base64.b64encode(image_bytes).decode("utf-8")

        llm = ChatOpenAI(model_name="gpt-4o", temperature=0)

        message = {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": """Describe the content of this image:
                            - In case of an image of chess, describe the position of the pieces and the state of the game.
                            - In case there is text in the image, transcribe it and describe it.
                            - In case of a graph, describe the axes, the data and the trends.
                            - In case of a map, describe the location and the features.
                            - In case of a diagram, describe the components and their relationships.
                            - In case of a photo, describe the scene, the objects and the people.""",
                },
                {
                    "type": "image",
                    "base64": f"{image_base64}",
                    "mime_type": mime_type,
                },
            ],
        }

        return llm.invoke([message])

    except Exception as e:
        return str(e)


@tool
def download_and_read_pdf(url: str) -> str:
    """Download and extract text from a PDF.

    Use this tool to inspect the complete contents of academic papers,
    especially acknowledgments, funding statements, grant numbers,
    award numbers, references, and appendices.

    The URL must be a verified PDF URL, not a guessed URL.

    Args:
        url: The URL of the PDF file to download.
    """
    try:
        response = requests.get(url)
        response.raise_for_status()

        with NamedTemporaryFile(suffix=".pdf") as fp:
            fp.write(response.content)
            text_content = ""
            pdf_reader = PdfReader(fp)
            for page in pdf_reader.pages:
                text_content += page.extract_text() + "\n"

        pdf_resume = f"""# PDF Content from {url}
{text_content[:6000]}
...
{text_content[-6000:]}
"""
        return pdf_resume

    except Exception as e:
        return f"Failed to download or read the PDF. Error: {str(e)}"