"""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" )