tools.py

"""Tools available to the GAIA agent.

Each tool is a LangChain `@tool`-decorated function so it can be bound to a
LangGraph ReAct agent. Tools are intentionally small, deterministic, and return
plain strings the LLM can reason over.
"""

from __future__ import annotations

import io
import os
import re

import requests
from langchain_core.tools import tool

# Base URL of the course scoring API; GAIA task files are served from here.
DEFAULT_API_URL = os.getenv(
    "GAIA_API_URL", "https://agents-course-unit4-scoring.hf.space"
)

# Max characters returned by any single tool. Tool output is re-sent on every
# turn, so large outputs blow rate limits (e.g. Groq free tier = 8k tokens/min).
# ~3500 chars is roughly 900 tokens. Override with GAIA_MAX_TOOL_CHARS.
MAX_TOOL_CHARS = int(os.getenv("GAIA_MAX_TOOL_CHARS", "3500"))


def _truncate(text: str) -> str:
    """Clip tool output to MAX_TOOL_CHARS so it fits provider rate limits."""
    if len(text) > MAX_TOOL_CHARS:
        return text[:MAX_TOOL_CHARS] + "\n\n... [truncated]"
    return text


# A polite, real-looking UA stops a lot of sites returning 403 to the agent.
_HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
        "(KHTML, like Gecko) Chrome/124.0 Safari/537.36"
    )
}


@tool
def web_search(query: str) -> str:
    """Search the web and return the top results (title, snippet, URL).

    Use this for current events, facts, or to discover pages to read with
    `visit_webpage`. Keep the query short and keyword-focused.
    """
    try:
        try:
            from ddgs import DDGS  # current package name
        except ImportError:
            from duckduckgo_search import DDGS  # legacy fallback

        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=6))
    except Exception as exc:  # noqa: BLE001 - surface error to the LLM
        return f"web_search error: {exc}"

    if not results:
        return "No results found."

    lines = []
    for r in results:
        title = r.get("title", "")
        body = r.get("body", "")
        href = r.get("href", "")
        lines.append(f"- {title}\n  {body}\n  {href}")
    return "\n".join(lines)


@tool
def wikipedia_search(query: str) -> str:
    """Look up a topic on Wikipedia and return a summary plus the page URL.

    Best for stable, encyclopedic facts (people, places, dates, definitions).
    """
    try:
        import wikipedia

        page = wikipedia.page(query, auto_suggest=True, redirect=True)
        summary = wikipedia.summary(query, sentences=8, auto_suggest=True)
        return f"Title: {page.title}\nURL: {page.url}\n\n{summary}"
    except Exception:
        # Fall back to disambiguation / search results.
        try:
            import wikipedia

            options = wikipedia.search(query, results=5)
            if options:
                return "Did not find an exact page. Closest matches: " + ", ".join(
                    options
                )
        except Exception as exc:  # noqa: BLE001
            return f"wikipedia_search error: {exc}"
        return "No Wikipedia results found."


@tool
def visit_webpage(url: str) -> str:
    """Fetch a web page and return its readable text content as markdown.

    Use after `web_search` to read a specific page in detail. Long pages are
    truncated to keep the context manageable.
    """
    try:
        resp = requests.get(url, headers=_HEADERS, timeout=30)
        resp.raise_for_status()
    except Exception as exc:  # noqa: BLE001
        return f"visit_webpage error: {exc}"

    try:
        from markdownify import markdownify

        text = markdownify(resp.text, heading_style="ATX")
    except Exception:
        text = resp.text

    # Collapse excessive blank lines and trim.
    text = re.sub(r"\n{3,}", "\n\n", text).strip()
    return _truncate(text)


@tool
def read_task_file(task_id: str) -> str:
    """Download the file attached to a GAIA task and return its contents.

    Some questions reference an attached file (CSV/XLSX/TXT/PY/JSON/image).
    Pass the question's `task_id`. Text and spreadsheet files are parsed to
    text; for other binary files a short description is returned.
    """
    url = f"{DEFAULT_API_URL}/files/{task_id}"
    try:
        resp = requests.get(url, headers=_HEADERS, timeout=60)
        resp.raise_for_status()
    except Exception as exc:  # noqa: BLE001
        return f"read_task_file error: {exc}"

    content = resp.content
    # Try to derive a filename from the Content-Disposition header.
    cd = resp.headers.get("content-disposition", "")
    fname = ""
    m = re.search(r'filename="?([^"]+)"?', cd)
    if m:
        fname = m.group(1)
    ext = os.path.splitext(fname)[1].lower()

    # Spreadsheets -> readable table text.
    if ext in {".xlsx", ".xls"}:
        try:
            import pandas as pd

            sheets = pd.read_excel(io.BytesIO(content), sheet_name=None)
            out = []
            for name, df in sheets.items():
                out.append(f"# Sheet: {name}\n{df.to_string()}")
            return _truncate("\n\n".join(out))
        except Exception as exc:  # noqa: BLE001
            return f"read_task_file (excel) error: {exc}"

    if ext == ".csv":
        try:
            import pandas as pd

            df = pd.read_csv(io.BytesIO(content))
            return _truncate(df.to_string())
        except Exception:
            pass  # fall through to text decode

    # Text-like files.
    if ext in {".txt", ".csv", ".py", ".json", ".md", ".tsv", ".xml", ".html", ""}:
        try:
            return _truncate(content.decode("utf-8", errors="replace"))
        except Exception as exc:  # noqa: BLE001
            return f"read_task_file (text) error: {exc}"

    return (
        f"Downloaded binary file '{fname or task_id}' "
        f"({len(content)} bytes, type={ext or 'unknown'}). "
        "This tool cannot interpret this file type directly."
    )


@tool
def reverse_text(text: str) -> str:
    """Reverse a string character-by-character.

    Use when the question text appears written backwards / mirrored, so you can
    read what it actually asks. Returns the reversed string.
    """
    return (text or "")[::-1]


@tool
def calculator(expression: str) -> str:
    """Evaluate a basic arithmetic expression and return the numeric result.

    Supports + - * / // % ** and parentheses. Use for exact arithmetic so the
    model does not make manual calculation mistakes. Example: '(12*7)+3'.
    """
    # Restrict to a safe arithmetic character set.
    if not re.fullmatch(r"[0-9\.\s\+\-\*/%\(\)]+", expression or ""):
        return "calculator error: only numbers and + - * / % ** ( ) are allowed."
    try:
        result = eval(expression, {"__builtins__": {}}, {})  # noqa: S307 - sandboxed
    except Exception as exc:  # noqa: BLE001
        return f"calculator error: {exc}"
    return str(result)


# Exported list consumed by the agent.
TOOLS = [
    web_search,
    wikipedia_search,
    visit_webpage,
    read_task_file,
    reverse_text,
    calculator,
]