files.py

"""Download and manage GAIA task attachments."""

from __future__ import annotations

import os
import shutil
from functools import lru_cache
from pathlib import Path

import requests
from huggingface_hub import hf_hub_download
from huggingface_hub.errors import GatedRepoError, HfHubHTTPError

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
GAIA_REPO = "gaia-benchmark/GAIA"
ATTACHMENTS_DIR = Path(os.getenv("GAIA_ATTACHMENTS_DIR", "task_files"))


def ensure_attachments_dir() -> Path:
    ATTACHMENTS_DIR.mkdir(parents=True, exist_ok=True)
    return ATTACHMENTS_DIR


def _copy_to_workspace(source: str | Path, file_name: str) -> Path:
    dest_dir = ensure_attachments_dir()
    dest_path = dest_dir / file_name
    if Path(source).resolve() == dest_path.resolve():
        return dest_path
    shutil.copy2(source, dest_path)
    return dest_path


def _download_from_scoring_api(task_id: str, file_name: str, api_url: str) -> Path | None:
    url = f"{api_url.rstrip('/')}/files/{task_id}"
    response = requests.get(url, timeout=60)
    if response.status_code == 404:
        return None
    response.raise_for_status()
    dest_path = ensure_attachments_dir() / file_name
    dest_path.write_bytes(response.content)
    return dest_path


@lru_cache(maxsize=256)
def _gaia_relative_candidates(file_name: str) -> tuple[str, ...]:
    """Common GAIA repo-relative paths for a given attachment name."""
    return tuple(
        {
            f"2023/validation/{file_name}",
            f"2023/test/{file_name}",
            f"2023/level1/{file_name}",
            file_name,
        }
    )


def _download_from_gaia(file_name: str) -> Path | None:
    token = os.getenv("HF_TOKEN")
    last_error: Exception | None = None

    for relative_path in _gaia_relative_candidates(file_name):
        try:
            cached = hf_hub_download(
                repo_id=GAIA_REPO,
                filename=relative_path,
                repo_type="dataset",
                token=token,
            )
            return _copy_to_workspace(cached, file_name)
        except GatedRepoError as exc:
            raise RuntimeError(
                "GAIA dataset access required for file attachments. "
                "Accept the terms at https://huggingface.co/datasets/gaia-benchmark/GAIA "
                "and ensure HF_TOKEN is set on your Space."
            ) from exc
        except HfHubHTTPError as exc:
            last_error = exc
            continue

    if last_error:
        print(f"GAIA download failed for {file_name}: {last_error}")
    return None


def download_task_file(
    task_id: str,
    file_name: str,
    api_url: str = DEFAULT_API_URL,
) -> Path:
    """Download a task attachment via scoring API or GAIA dataset fallback."""
    if not file_name:
        raise ValueError("file_name is required")

    dest_path = ensure_attachments_dir() / file_name
    if dest_path.exists() and dest_path.stat().st_size > 0:
        return dest_path

    from_api = _download_from_scoring_api(task_id, file_name, api_url)
    if from_api is not None:
        return from_api

    from_gaia = _download_from_gaia(file_name)
    if from_gaia is not None:
        return from_gaia

    raise FileNotFoundError(
        f"Could not download attachment '{file_name}' for task {task_id}. "
        "Scoring API returned 404; GAIA fallback also failed."
    )


def file_context_block(file_path: Path | None) -> str:
    if file_path is None:
        return ""
    suffix = file_path.suffix.lower()
    hints = {
        ".png": "PNG image — use analyze_image.",
        ".jpg": "JPEG image — use analyze_image.",
        ".jpeg": "JPEG image — use analyze_image.",
        ".webp": "WebP image — use analyze_image.",
        ".mp3": "MP3 audio — use transcribe_audio.",
        ".wav": "WAV audio — use transcribe_audio.",
        ".xlsx": "Excel spreadsheet — use read_spreadsheet.",
        ".xls": "Excel spreadsheet — use read_spreadsheet.",
        ".csv": "CSV file — use read_spreadsheet.",
        ".py": "Python script — use execute_python_file.",
    }
    hint = hints.get(suffix, "Use the appropriate file tool.")
    return (
        f"\n\nAttached file:\n"
        f"- path: {file_path.resolve()}\n"
        f"- name: {file_path.name}\n"
        f"- hint: {hint}\n"
    )