tools.py

"""Tools for the GAIA Level-1 evaluation agent."""

from __future__ import annotations

import os
import re
import subprocess
import tempfile
from pathlib import Path

import requests
from langchain_core.tools import tool

API_URL = os.getenv("SCORING_API_URL", "https://agents-course-unit4-scoring.hf.space")
GAIA_REPO = "gaia-benchmark/GAIA"
FILES_DIR = Path(tempfile.gettempdir()) / "gaia_task_files"
FILES_DIR.mkdir(parents=True, exist_ok=True)

_GAIA_FILES: list[str] | None = None


def _truncate(text: str, limit: int = 1800) -> str:
    text = text.strip()
    if len(text) <= limit:
        return text
    return text[:limit] + "\n...[truncated]"


@tool
def wikipedia_search(query: str) -> str:
    """Search English Wikipedia and return a short page summary."""
    try:
        import wikipedia

        wikipedia.set_lang("en")
        results = wikipedia.search(query, results=3)
        if not results:
            return f"No Wikipedia results for: {query}"
        title = results[0]
        page = wikipedia.page(title, auto_suggest=False)
        return _truncate(f"Title: {page.title}\nURL: {page.url}\n\n{page.summary}")
    except Exception as e:  # noqa: BLE001
        return f"Wikipedia error: {e}"


@tool
def web_search(query: str) -> str:
    """Search the public web and return top result snippets."""
    try:
        try:
            from ddgs import DDGS
        except ImportError:
            from duckduckgo_search import DDGS

        rows = []
        with DDGS() as ddgs:
            for i, item in enumerate(ddgs.text(query, max_results=5), start=1):
                rows.append(
                    f"{i}. {item.get('title')}\n"
                    f"URL: {item.get('href')}\n"
                    f"{item.get('body')}"
                )
        return _truncate("\n\n".join(rows) if rows else f"No web results for: {query}")
    except Exception as e:  # noqa: BLE001
        return f"Web search error: {e}"


@tool
def youtube_transcript(url: str) -> str:
    """Fetch the transcript/captions text for a YouTube video URL."""
    try:
        from youtube_transcript_api import YouTubeTranscriptApi

        match = re.search(r"(?:v=|youtu\.be/)([A-Za-z0-9_-]{6,})", url)
        if not match:
            return "Could not parse YouTube video id from URL."
        video_id = match.group(1)
        api = YouTubeTranscriptApi()
        parts = api.fetch(video_id)
        text = " ".join(getattr(p, "text", str(p)) for p in parts)
        return _truncate(text, 3000)
    except Exception as e:  # noqa: BLE001
        return f"YouTube transcript error: {e}"


def _fetch_from_api(task_id: str) -> Path | None:
    resp = requests.get(f"{API_URL}/files/{task_id}", timeout=60)
    if resp.status_code != 200:
        return None
    filename = task_id
    match = re.search(r'filename="?([^";]+)"?', resp.headers.get("content-disposition", ""))
    if match:
        filename = match.group(1)
    path = FILES_DIR / filename
    path.write_bytes(resp.content)
    return path


def _fetch_from_gaia(task_id: str) -> Path | None:
    """The scoring API often has no file path; GAIA stores attachments as <task_id>.<ext>."""
    global _GAIA_FILES
    from huggingface_hub import hf_hub_download, list_repo_files

    token = os.getenv("HF_TOKEN")
    if _GAIA_FILES is None:
        _GAIA_FILES = list_repo_files(GAIA_REPO, repo_type="dataset", token=token)
    remote = next((f for f in _GAIA_FILES if Path(f).stem == task_id), None)
    if not remote:
        return None
    return Path(hf_hub_download(GAIA_REPO, remote, repo_type="dataset", token=token))


def _preview(path: Path) -> str:
    suffix = path.suffix.lower()
    if suffix in {".txt", ".py", ".csv", ".md", ".json", ".jsonld"}:
        return path.read_text(errors="ignore")[:1500]
    if suffix in {".xlsx", ".xls"}:
        return "Excel file saved. Use analyze_excel to compute values."
    if suffix in {".mp3", ".wav", ".m4a"}:
        return "Audio file saved. Use transcribe_audio to listen."
    if suffix in {".png", ".jpg", ".jpeg", ".webp"}:
        return "Image file saved. Use analyze_image to inspect it."
    if suffix == ".pdf":
        return "PDF file saved."
    return f"Binary file saved ({path.stat().st_size} bytes)."


@tool
def download_task_file(task_id: str) -> str:
    """Download the file attached to a GAIA task_id.

    Tries the scoring API first, then the GAIA dataset on the Hugging Face Hub.
    Returns the saved path plus a short content preview.
    """
    try:
        path = _fetch_from_api(task_id)
        source = "scoring API"
        if path is None:
            path = _fetch_from_gaia(task_id)
            source = "GAIA dataset"
        if path is None:
            return f"No file found for task_id {task_id}."
        return f"Saved to: {path} (via {source})\nPreview:\n{_preview(path)}"
    except Exception as e:  # noqa: BLE001
        if "gated" in str(e).lower() or "403" in str(e):
            return (
                f"The file for {task_id} lives in the gated GAIA dataset. Accept the terms "
                f"at https://huggingface.co/datasets/{GAIA_REPO} to enable downloads."
            )
        return f"download_task_file error: {e}"


@tool
def run_python_file(path: str) -> str:
    """Execute a local Python file and return stdout/stderr (for attached .py tasks)."""
    try:
        proc = subprocess.run(
            ["python", path],
            capture_output=True,
            text=True,
            timeout=30,
            cwd=str(Path(path).parent),
        )
        out = (proc.stdout or "") + (("\n" + proc.stderr) if proc.stderr else "")
        return _truncate(out.strip() or f"(no output, exit={proc.returncode})")
    except Exception as e:  # noqa: BLE001
        return f"run_python_file error: {e}"


@tool
def analyze_excel(path: str, question: str) -> str:
    """Read an Excel file and return sheet names plus a compact table preview to answer sales questions."""
    try:
        import pandas as pd

        xls = pd.ExcelFile(path)
        chunks = [f"Sheets: {xls.sheet_names}"]
        for sheet in xls.sheet_names:
            df = pd.read_excel(xls, sheet_name=sheet)
            chunks.append(f"\nSheet={sheet} columns={list(df.columns)}")
            chunks.append(df.head(30).to_csv(index=False))
            # helpful totals if numeric columns exist
            num = df.select_dtypes(include="number")
            if not num.empty:
                chunks.append("Numeric column sums:\n" + num.sum().to_string())
        chunks.append(f"\nQuestion reminder: {question}")
        return _truncate("\n".join(chunks), 3000)
    except Exception as e:  # noqa: BLE001
        return f"analyze_excel error: {e}"


@tool
def transcribe_audio(path: str) -> str:
    """Transcribe an audio file (mp3/wav) using Groq Whisper."""
    try:
        from groq import Groq

        client = Groq(api_key=os.getenv("GROQ_API_KEY"))
        with open(path, "rb") as f:
            result = client.audio.transcriptions.create(
                file=f,
                model="whisper-large-v3",
            )
        text = getattr(result, "text", None) or str(result)
        return _truncate(text, 3000)
    except Exception as e:  # noqa: BLE001
        return f"transcribe_audio error: {e}"


@tool
def analyze_image(path: str, question: str) -> str:
    """Answer a question about a local image file (chess positions, charts, photos).

    Needs GROQ_VISION_MODEL set to a vision-capable Groq model.
    """
    model = os.getenv("GROQ_VISION_MODEL")
    if not model:
        return (
            "No vision model is configured, so the image cannot be read. "
            "Answer from the question text alone."
        )
    try:
        import base64

        from groq import Groq

        image = Path(path)
        mime = "image/png" if image.suffix.lower() == ".png" else "image/jpeg"
        encoded = base64.b64encode(image.read_bytes()).decode()
        client = Groq(api_key=os.getenv("GROQ_API_KEY"))
        resp = client.chat.completions.create(
            model=model,
            temperature=0,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": question},
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:{mime};base64,{encoded}"},
                        },
                    ],
                }
            ],
        )
        return _truncate(resp.choices[0].message.content or "")
    except Exception as e:  # noqa: BLE001
        return f"analyze_image error: {e}"


@tool
def reverse_text(text: str) -> str:
    """Reverse a string. Useful when a question is written backwards."""
    return text[::-1]


TOOLS = [
    wikipedia_search,
    web_search,
    youtube_transcript,
    download_task_file,
    run_python_file,
    analyze_excel,
    transcribe_audio,
    analyze_image,
    reverse_text,
]