test_agent.py

"""
Script to test the GAIA agent with a specific question.
This is useful for testing the agent's response to a specific question
without having to run the full Gradio interface.
"""

import sys
import json
from pathlib import Path
import argparse
import logging
import os

# Configure logging
logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Import the agent class from app2.py
try:
    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
    from app2 import GAIAAgent
except ImportError:
    logger.error("Failed to import GAIAAgent from app2.py")
    sys.exit(1)

def load_questions(file_path):
    """Load questions from a JSON file."""
    try:
        with open(file_path, 'r') as f:
            return json.load(f)
    except Exception as e:
        logger.error(f"Error loading questions from {file_path}: {e}")
        return []

def find_question_by_id(questions, task_id):
    """Find a question by its task_id."""
    for q in questions:
        if q.get("task_id") == task_id:
            return q
    return None

def main():
    parser = argparse.ArgumentParser(description='Test the GAIA agent with a specific question')
    parser.add_argument('--question', '-q', type=str, help='The question to ask the agent')
    parser.add_argument('--task-id', '-t', type=str, help='Task ID to look up in common_questions.json')
    parser.add_argument('--file', '-f', type=str, default='question_set/common_questions.json',
                      help='Path to questions file (default: question_set/common_questions.json)')
    
    args = parser.parse_args()
    
    # Initialize the agent
    logger.info("Initializing GAIA Agent...")
    agent = GAIAAgent()
    logger.info("Agent initialized")
    
    question = args.question
    
    # If task_id is provided, look up the question in the file
    if not question and args.task_id:
        questions = load_questions(args.file)
        question_obj = find_question_by_id(questions, args.task_id)
        
        if question_obj:
            question = question_obj.get("Question")
            expected_answer = question_obj.get("Final answer", "Not provided")
            logger.info(f"Found question for task_id {args.task_id}")
            logger.info(f"Expected answer: {expected_answer}")
        else:
            logger.error(f"Could not find question with task_id {args.task_id}")
            sys.exit(1)
    
    # Check if we have a question to answer
    if not question:
        logger.error("No question provided. Use --question or --task-id")
        sys.exit(1)
    
    logger.info(f"Question: {question}")
    
    # Get the agent's answer
    logger.info("Asking agent...")
    try:
        answer = agent(question)
        logger.info(f"Agent's answer: {answer}")
    except Exception as e:
        logger.error(f"Error getting answer from agent: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()