test_local.py

"""
Local test script for the GAIA benchmark agent.

Fetches questions from the scoring API and runs the agent on 2-3 questions
WITHOUT submitting answers. This lets you validate reasoning and answer
formatting before burning submission attempts.

Usage:
    set MODEL_ID=gpt-4o
    set OPENAI_API_KEY=sk-...
    python test_local.py
"""

import os
import sys
import requests

from agent import create_agent, _clean_answer, build_question_prompt


API_URL = "https://agents-course-unit4-scoring.hf.space"
NUM_QUESTIONS = 3


def main():
    print("=" * 60)
    print("GAIA Agent — Local Test")
    print("=" * 60)

    # Check that the agent can be created
    print("\n[1] Creating agent...")
    try:
        agent = create_agent()
        print("    Agent created successfully.")
    except Exception as e:
        print(f"    Failed to create agent: {e}")
        sys.exit(1)

    # Fetch questions
    print(f"\n[2] Fetching questions from {API_URL}/questions ...")
    try:
        resp = requests.get(f"{API_URL}/questions", timeout=15)
        resp.raise_for_status()
        questions = resp.json()
    except Exception as e:
        print(f"    Failed to fetch questions: {e}")
        sys.exit(1)

    print(f"    Got {len(questions)} total questions.")

    # Prefer text-only questions first (no file attachment)
    no_file = [q for q in questions if not q.get("file_name")]
    with_file = [q for q in questions if q.get("file_name")]
    selected = no_file[:NUM_QUESTIONS]
    if len(selected) < NUM_QUESTIONS:
        selected += with_file[: NUM_QUESTIONS - len(selected)]

    print(f"    Selected {len(selected)} questions to test.\n")

    # Run agent on each
    for i, item in enumerate(selected):
        task_id = item.get("task_id", "???")
        question = item.get("question", "")
        file_name = item.get("file_name", "")
        prompt = build_question_prompt(item)

        print("=" * 60)
        print(f"Question {i + 1}/{len(selected)}")
        print(f"  Task ID:     {task_id}")
        print(f"  File:        {file_name or '(none)'}")
        print(f"  Question:    {question[:200]}{'...' if len(question) > 200 else ''}")
        print("-" * 60)

        try:
            raw_output = agent.run(prompt)
        except Exception as e:
            raw_output = f"ERROR: {e}"

        cleaned = _clean_answer(raw_output)

        print(f"  Raw output:   {raw_output!r}")
        print(f"  Clean answer: {cleaned!r}")
        print()

    print("=" * 60)
    print("Test complete.")
    print("=" * 60)


if __name__ == "__main__":
    main()