test_agent.py

#!/usr/bin/env python3
"""
Test script for the GAIA agent to verify it works correctly.
"""

import os
from dotenv import load_dotenv
from agent import build_graph
from langchain_core.messages import HumanMessage

# Load environment variables
load_dotenv()

def test_agent():
    """Test the agent with a simple question."""
    print("Testing GAIA Agent with DeepSeek...")
    
    # Build the graph with DeepSeek provider
    try:
        graph = build_graph(provider="deepseek")
        print("✅ Agent graph built successfully")
    except Exception as e:
        print(f"❌ Error building agent graph: {e}")
        return
    
    # Test question
    test_question = "What is 15 + 27?"
    print(f"\nTest question: {test_question}")
    
    try:
        # Run the agent
        messages = [HumanMessage(content=test_question)]
        result = graph.invoke({"messages": messages})
        
        if result and "messages" in result and result["messages"]:
            answer = result["messages"][-1].content
            print(f"\nAgent response: {answer}")
            
            # Extract final answer (similar to app.py logic)
            lines = answer.strip().split('\n')
            for line in reversed(lines):
                line = line.strip()
                if line and not line.startswith('Based on') and not line.startswith('Therefore') and not line.startswith('In conclusion'):
                    print(f"\nExtracted final answer: {line}")
                    break
            
            print("\n✅ Agent test completed successfully")
        else:
            print("❌ No response from agent")
            
    except Exception as e:
        print(f"❌ Error running agent: {e}")

if __name__ == "__main__":
    test_agent()