tools.py

"""Custom smolagents tools for GAIA multimodal and file tasks."""

from __future__ import annotations

import base64
import os
import re
import subprocess
import sys
from pathlib import Path

import pandas as pd
from huggingface_hub import InferenceClient
from smolagents import tool
from youtube_transcript_api import YouTubeTranscriptApi


def _hf_client() -> InferenceClient:
    token = os.getenv("HF_TOKEN")
    return InferenceClient(token=token)


def _vision_model() -> str:
    return os.getenv("HF_VISION_MODEL", "Qwen/Qwen2-VL-7B-Instruct")


def _asr_model() -> str:
    return os.getenv("HF_ASR_MODEL", "openai/whisper-large-v3")


@tool
def read_spreadsheet(file_path: str) -> str:
    """Read an Excel or CSV file and return its contents as text for analysis.

    Args:
        file_path: Absolute or relative path to .xlsx, .xls, or .csv file.
    """
    path = Path(file_path)
    if not path.exists():
        return f"File not found: {file_path}"

    suffix = path.suffix.lower()
    if suffix == ".csv":
        df = pd.read_csv(path)
    elif suffix in {".xlsx", ".xls"}:
        df = pd.read_excel(path)
    else:
        return f"Unsupported spreadsheet type: {suffix}"

    buffer = []
    buffer.append(f"Shape: {df.shape[0]} rows x {df.shape[1]} columns")
    buffer.append(f"Columns: {', '.join(str(c) for c in df.columns)}")
    buffer.append("\n--- data ---")
    buffer.append(df.to_string(index=False))
    text = "\n".join(buffer)
    return text[:50000]


@tool
def execute_python_file(file_path: str) -> str:
    """Execute a Python file in a subprocess and return stdout/stderr.

    Args:
        file_path: Path to a .py file to run.
    """
    path = Path(file_path)
    if not path.exists():
        return f"File not found: {file_path}"
    if path.suffix.lower() != ".py":
        return f"Not a Python file: {file_path}"

    try:
        completed = subprocess.run(
            [sys.executable, str(path.resolve())],
            capture_output=True,
            text=True,
            timeout=45,
            cwd=str(path.parent.resolve()),
        )
    except subprocess.TimeoutExpired:
        return "Execution timed out after 45 seconds."

    parts = []
    if completed.stdout:
        parts.append(f"STDOUT:\n{completed.stdout}")
    if completed.stderr:
        parts.append(f"STDERR:\n{completed.stderr}")
    parts.append(f"Exit code: {completed.returncode}")
    return "\n".join(parts)[:20000]


@tool
def transcribe_audio(file_path: str) -> str:
    """Transcribe speech from an audio file (mp3/wav) to text.

    Args:
        file_path: Path to the audio file.
    """
    path = Path(file_path)
    if not path.exists():
        return f"File not found: {file_path}"

    client = _hf_client()
    with path.open("rb") as audio_file:
        result = client.automatic_speech_recognition(
            audio=audio_file.read(),
            model=_asr_model(),
        )

    if isinstance(result, dict):
        return str(result.get("text", result))[:20000]
    return str(getattr(result, "text", result))[:20000]


@tool
def analyze_image(file_path: str, question: str) -> str:
    """Analyze an image file to answer a specific question about it.

    Args:
        file_path: Path to png/jpg/jpeg/webp image.
        question: What to determine from the image.
    """
    path = Path(file_path)
    if not path.exists():
        return f"File not found: {file_path}"

    mime = {
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".webp": "image/webp",
    }.get(path.suffix.lower(), "image/png")

    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    client = _hf_client()
    response = client.chat_completion(
        model=_vision_model(),
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{encoded}"}},
                    {"type": "text", "text": question},
                ],
            }
        ],
        max_tokens=1024,
    )
    return response.choices[0].message.content.strip()[:10000]


@tool
def get_youtube_transcript(video_url: str) -> str:
    """Fetch the transcript/captions of a YouTube video.

    Args:
        video_url: Full YouTube URL or 11-character video ID.
    """
    match = re.search(
        r"(?:youtube\.com/watch\?v=|youtu\.be/|youtube\.com/embed/)([A-Za-z0-9_-]{11})",
        video_url,
    )
    video_id = match.group(1) if match else video_url.strip()

    try:
        api = YouTubeTranscriptApi()
        fetched = api.fetch(video_id)
        lines = [snippet.text for snippet in fetched.snippets]
    except Exception as exc:
        return f"Could not fetch transcript: {exc}"

    return " ".join(lines)[:30000]


def build_custom_tools() -> list:
    return [
        read_spreadsheet,
        execute_python_file,
        transcribe_audio,
        analyze_image,
        get_youtube_transcript,
    ]