test_agent.py

"""
Test suite for GAIA Agent
Validates imports, tools, model, and agent functionality
"""
import os
import sys

def test_imports():
    """Test 1: Verify all critical imports work"""
    print("\n" + "="*60)
    print("TEST 1: Checking Imports")
    print("="*60)
    
    try:
        import smolagents
        print("✅ smolagents imported successfully")
    except ImportError as e:
        print(f"❌ smolagents import failed: {e}")
        return False
    
    try:
        import litellm
        print("✅ litellm imported successfully")
    except ImportError as e:
        print(f"❌ litellm import failed: {e}")
        return False
    
    try:
        import google.generativeai as genai
        print("✅ google.generativeai imported successfully")
    except ImportError as e:
        print(f"❌ google.generativeai import failed: {e}")
        return False
    
    try:
        from agents import Agent
        print("✅ Agent class imported successfully")
    except ImportError as e:
        print(f"❌ Agent import failed: {e}")
        return False
    
    try:
        from tool import get_tools
        print("✅ get_tools imported successfully")
    except ImportError as e:
        print(f"❌ get_tools import failed: {e}")
        return False
    
    try:
        from model import get_model
        print("✅ get_model imported successfully")
    except ImportError as e:
        print(f"❌ get_model import failed: {e}")
        return False
    
    print("\n✅ All imports successful!")
    return True


def test_tools():
    """Test 2: Verify tools can be loaded"""
    print("\n" + "="*60)
    print("TEST 2: Checking Tools")
    print("="*60)
    
    try:
        from tool import get_tools
        tools = get_tools()
        print(f"✅ Loaded {len(tools)} tools")
        
        # List tools
        for i, tool in enumerate(tools, 1):
            tool_name = getattr(tool, 'name', type(tool).__name__)
            print(f"   {i}. {tool_name}")
        
        return True
    except Exception as e:
        print(f"❌ Tools loading failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_environment():
    """Test 3: Check environment variables"""
    print("\n" + "="*60)
    print("TEST 3: Checking Environment")
    print("="*60)
    
    gemini_key = os.getenv("GEMINI_API_KEY")
    
    if gemini_key:
        # Mask the key for security
        masked_key = gemini_key[:8] + "..." + gemini_key[-4:] if len(gemini_key) > 12 else "***"
        print(f"✅ GEMINI_API_KEY configured: {masked_key}")
        return True
    else:
        print("⚠️  GEMINI_API_KEY not configured")
        print("   Set it with: $env:GEMINI_API_KEY=\"your_key_here\" (PowerShell)")
        print("   Or create a .env file with: GEMINI_API_KEY=your_key_here")
        return False


def test_model():
    """Test 4: Verify model can be instantiated"""
    print("\n" + "="*60)
    print("TEST 4: Checking Model")
    print("="*60)
    
    if not os.getenv("GEMINI_API_KEY"):
        print("⚠️  Skipping model test (GEMINI_API_KEY not configured)")
        return False
    
    try:
        from model import get_model
        
        model_id = "gemini/gemini-2.0-flash-exp"
        print(f"Initializing model: {model_id}")
        
        model = get_model("LiteLLMModel", model_id)
        print("✅ Model instantiated successfully")
        
        # Test simple inference
        print("Testing simple inference...")
        messages = [{"role": "user", "content": "Say 'test successful' and nothing else"}]
        response = model(messages)
        print(f"✅ Model response: {response.content[:100]}")
        
        return True
    except Exception as e:
        print(f"❌ Model test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_agent():
    """Test 5: Verify agent can be created and used"""
    print("\n" + "="*60)
    print("TEST 5: Checking Agent")
    print("="*60)
    
    if not os.getenv("GEMINI_API_KEY"):
        print("⚠️  Skipping agent test (GEMINI_API_KEY not configured)")
        return False
    
    try:
        from agents import Agent
        from tool import get_tools
        from model import get_model
        
        model_id = "gemini/gemini-2.0-flash-exp"
        print(f"Creating agent with model: {model_id}")
        
        agent = Agent(
            model=get_model("LiteLLMModel", model_id),
            tools=get_tools(),
            verbose=True
        )
        print("✅ Agent created successfully")
        
        # Test simple question
        print("\nTesting with simple question...")
        question = "What is 2+2? Answer with only the number."
        print(f"Question: {question}")
        
        answer = agent(question)
        print(f"✅ Agent answer: {answer}")
        
        return True
    except Exception as e:
        print(f"❌ Agent test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """Run all tests"""
    print("\n" + "🔬" * 30)
    print(" GAIA AGENT TEST SUITE ".center(60, "="))
    print("🔬" * 30)
    
    results = {
        "Imports": test_imports(),
        "Tools": test_tools(),
        "Environment": test_environment(),
        "Model": test_model(),
        "Agent": test_agent(),
    }
    
    # Summary
    print("\n" + "="*60)
    print("TEST SUMMARY")
    print("="*60)
    
    passed = sum(results.values())
    total = len(results)
    
    for test_name, passed_test in results.items():
        status = "✅ PASS" if passed_test else "❌ FAIL"
        print(f"{test_name:20s} {status}")
    
    print("="*60)
    print(f"Total: {passed}/{total} tests passed ({passed/total*100:.1f}%)")
    print("="*60)
    
    if passed == total:
        print("\n🎉 ALL TESTS PASSED! Ready for deployment!")
        return 0
    else:
        print(f"\n⚠️  {total - passed} test(s) failed. Please fix before deployment.")
        return 1


if __name__ == "__main__":
    sys.exit(main())