gaia_workbench.py

from __future__ import annotations

import argparse
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import requests


API_URL = "https://agents-course-unit4-scoring.hf.space"
ROOT = Path(__file__).resolve().parent
DATA_DIR = ROOT / "data"
QUESTIONS_PATH = DATA_DIR / "questions.json"
FILES_DIR = DATA_DIR / "files"
ANSWERS_PATH = ROOT / "answers.local.json"


@dataclass(frozen=True)
class Question:
    task_id: str
    question: str
    file_name: str = ""


def _request_json(method: str, url: str, **kwargs: Any) -> Any:
    response = requests.request(method, url, timeout=60, **kwargs)
    response.raise_for_status()
    return response.json()


def fetch_questions() -> list[Question]:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    raw_questions = _request_json("GET", f"{API_URL}/questions")
    QUESTIONS_PATH.write_text(json.dumps(raw_questions, indent=2, ensure_ascii=False), encoding="utf-8")
    return [
        Question(
            task_id=str(item["task_id"]),
            question=str(item["question"]),
            file_name=str(item.get("file_name") or ""),
        )
        for item in raw_questions
    ]


def load_questions() -> list[Question]:
    if not QUESTIONS_PATH.exists():
        return fetch_questions()
    raw_questions = json.loads(QUESTIONS_PATH.read_text(encoding="utf-8"))
    return [
        Question(
            task_id=str(item["task_id"]),
            question=str(item["question"]),
            file_name=str(item.get("file_name") or ""),
        )
        for item in raw_questions
    ]


def load_answers() -> dict[str, str]:
    if not ANSWERS_PATH.exists():
        return {}
    raw_answers = json.loads(ANSWERS_PATH.read_text(encoding="utf-8"))
    return {str(task_id): str(answer).strip() for task_id, answer in raw_answers.items() if str(answer).strip()}


def write_answer(task_id: str, answer: str) -> None:
    answers = load_answers()
    answers[task_id] = answer.strip()
    ANSWERS_PATH.write_text(json.dumps(answers, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def download_task_file(task_id: str, file_name: str = "") -> Path:
    FILES_DIR.mkdir(parents=True, exist_ok=True)
    response = requests.get(f"{API_URL}/files/{task_id}", timeout=60)
    response.raise_for_status()
    target_name = file_name or task_id
    target = FILES_DIR / target_name
    target.write_bytes(response.content)
    return target


def download_files() -> tuple[list[Path], list[str]]:
    downloaded: list[Path] = []
    failed: list[str] = []
    for question in load_questions():
        if question.file_name:
            try:
                downloaded.append(download_task_file(question.task_id, question.file_name))
            except requests.HTTPError as exc:
                failed.append(f"{question.task_id} ({question.file_name}): {exc}")
    return downloaded, failed


def answer_payload() -> list[dict[str, str]]:
    answers = load_answers()
    questions = load_questions()
    return [
        {"task_id": question.task_id, "submitted_answer": answers[question.task_id]}
        for question in questions
        if question.task_id in answers
    ]


def submit_answers(username: str, agent_code: str) -> dict[str, Any]:
    payload = {
        "username": username.strip(),
        "agent_code": agent_code.strip(),
        "answers": answer_payload(),
    }
    if not payload["username"]:
        raise ValueError("Missing username")
    if not payload["agent_code"]:
        raise ValueError("Missing agent_code")
    if not payload["answers"]:
        raise ValueError("No answers found in answers.local.json")
    return _request_json("POST", f"{API_URL}/submit", json=payload)


def status_rows() -> list[dict[str, str]]:
    answers = load_answers()
    rows = []
    for index, question in enumerate(load_questions(), start=1):
        rows.append(
            {
                "#": str(index),
                "task_id": question.task_id,
                "file": question.file_name,
                "answered": "yes" if question.task_id in answers else "no",
                "answer": answers.get(question.task_id, ""),
                "question": question.question,
            }
        )
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description="Hugging Face GAIA Unit 4 workbench")
    subparsers = parser.add_subparsers(dest="command", required=True)

    subparsers.add_parser("fetch", help="Fetch and cache questions")
    subparsers.add_parser("download-files", help="Download all attached task files")
    subparsers.add_parser("status", help="Print answer coverage")

    set_answer_parser = subparsers.add_parser("set-answer", help="Set one exact-match answer")
    set_answer_parser.add_argument("task_id")
    set_answer_parser.add_argument("answer")

    submit_parser = subparsers.add_parser("submit", help="Submit filled local answers")
    submit_parser.add_argument("--username", default=os.getenv("HF_USERNAME", ""))
    submit_parser.add_argument("--agent-code", default=os.getenv("AGENT_CODE", ""))

    args = parser.parse_args()

    if args.command == "fetch":
        questions = fetch_questions()
        print(f"Fetched {len(questions)} questions -> {QUESTIONS_PATH}")
    elif args.command == "download-files":
        downloaded, failed = download_files()
        print(f"Downloaded {len(downloaded)} files")
        for path in downloaded:
            print(path)
        if failed:
            print(f"Failed {len(failed)} files")
            for item in failed:
                print(item)
    elif args.command == "status":
        rows = status_rows()
        answered_count = sum(1 for row in rows if row["answered"] == "yes")
        print(f"Answered {answered_count}/{len(rows)}")
        for row in rows:
            marker = "[x]" if row["answered"] == "yes" else "[ ]"
            file_part = f" file={row['file']}" if row["file"] else ""
            print(f"{marker} {row['#']}. {row['task_id']}{file_part} - {row['question']}")
    elif args.command == "set-answer":
        write_answer(args.task_id, args.answer)
        print(f"Saved answer for {args.task_id}")
    elif args.command == "submit":
        result = submit_answers(args.username, args.agent_code)
        print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()