test_enhanced_agent.py

#!/usr/bin/env python3
"""
Test script for the enhanced GAIA agent
"""
import os
import sys
from dotenv import load_dotenv

# Add current directory to path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

try:
    from langgraph_new import graph, answer_gaia_question, get_random_gaia_question
    print("✅ Successfully imported enhanced GAIA agent")
except ImportError as e:
    print(f"❌ Import error: {e}")
    sys.exit(1)

def test_basic_functionality():
    """Test basic agent functionality"""
    print("\n🔧 Testing basic functionality...")
    
    test_cases = [
        ("What is 2 + 2?", "4"),
        ("What is the capital of France?", "Paris"),
        ("List these items alphabetically: zebra, apple, banana", "apple, banana, zebra"),
    ]
    
    for question, expected in test_cases:
        try:
            answer = answer_gaia_question(question)
            print(f"Q: {question}")
            print(f"A: {answer}")
            print(f"Expected: {expected}")
            print(f"Match: {'✅' if expected.lower() in answer.lower() else '❌'}")
            print("-" * 50)
        except Exception as e:
            print(f"❌ Error answering '{question}': {e}")

def test_file_analysis():
    """Test file analysis capabilities"""
    print("\n📊 Testing file analysis...")
    
    # Test Excel file if it exists
    if os.path.exists("test_sales.xlsx"):
        try:
            question = "Given the Excel file at test_sales.xlsx, what is the structure of the data?"
            answer = answer_gaia_question(question)
            print(f"Q: {question}")
            print(f"A: {answer}")
        except Exception as e:
            print(f"❌ Excel test error: {e}")
    else:
        print("⚠️ test_sales.xlsx not found, skipping Excel test")
    
    # Test audio file if it exists
    if os.path.exists("test.wav"):
        try:
            question = "What does the speaker say in the audio file test.wav?"
            answer = answer_gaia_question(question)
            print(f"Q: {question}")
            print(f"A: {answer}")
        except Exception as e:
            print(f"❌ Audio test error: {e}")
    else:
        print("⚠️ test.wav not found, skipping audio test")

def test_youtube_capability():
    """Test YouTube transcript capability"""
    print("\n🎥 Testing YouTube capability...")
    
    try:
        # Test with a known working video
        question = """Examine the video at https://www.youtube.com/watch?v=1htKBjuUWec. What does Teal'c say in response to the question "Isn't that hot?" """
        answer = answer_gaia_question(question)
        print(f"Q: {question}")
        print(f"A: {answer}")
    except Exception as e:
        print(f"❌ YouTube test error: {e}")

def test_web_search():
    """Test web search capabilities"""
    print("\n🌐 Testing web search...")
    
    try:
        question = "Who is the current president of France in 2025?"
        answer = answer_gaia_question(question)
        print(f"Q: {question}")
        print(f"A: {answer}")
    except Exception as e:
        print(f"❌ Web search test error: {e}")

def test_real_gaia_question():
    """Test with a real GAIA question from the API"""
    print("\n🎯 Testing with real GAIA question...")
    
    try:
        question_data = get_random_gaia_question()
        if question_data:
            question = question_data.get('question', '')
            task_id = question_data.get('task_id', 'Unknown')
            
            print(f"Task ID: {task_id}")
            print(f"Question: {question}")
            
            answer = answer_gaia_question(question)
            print(f"Agent Answer: {answer}")
            
            return {"task_id": task_id, "question": question, "answer": answer}
        else:
            print("⚠️ Could not fetch random GAIA question")
            return None
    except Exception as e:
        print(f"❌ Real GAIA question test error: {e}")
        return None

def main():
    """Main test runner"""
    load_dotenv()
    
    print("🚀 Starting GAIA Agent Tests")
    print("=" * 60)
    
    # Check environment variables
    required_vars = ["OPENAI_API_KEY", "TAVILY_API_KEY"]
    missing_vars = [var for var in required_vars if not os.getenv(var)]
    
    if missing_vars:
        print(f"❌ Missing environment variables: {missing_vars}")
        print("Please set these in your .env file")
        return
    
    # Run tests
    test_basic_functionality()
    test_file_analysis()
    test_web_search()
    test_youtube_capability()
    
    # Test with real GAIA question
    gaia_result = test_real_gaia_question()
    
    print("\n" + "=" * 60)
    print("🎉 Test suite completed!")
    
    if gaia_result:
        print("\n📋 Sample GAIA Result:")
        print(f"Task ID: {gaia_result['task_id']}")
        print(f"Answer: {gaia_result['answer']}")

if __name__ == "__main__":
    main()