test_tools_integration.py

#!/usr/bin/env python3
"""
Test script to verify tool integration in the LangGraph agent system
"""

from src.langgraph_system import run_agent_system

def test_retrieval_tools():
    """Test that retrieval tools (Wikipedia, web search, etc.) are working"""
    print("=" * 60)
    print("Testing Retrieval Tools Integration")
    print("=" * 60)
    
    # Test Wikipedia search
    query = "When was Albert Einstein born?"
    print(f"\nTesting query: {query}")
    print("-" * 40)
    
    result = run_agent_system(query, user_id="test_user", session_id="test_session")
    print(f"Result: {result}")
    
    return result

def test_execution_tools():
    """Test that execution tools (Python code execution) are working"""
    print("=" * 60)
    print("Testing Execution Tools Integration")
    print("=" * 60)
    
    # Test code execution
    query = "Calculate the first 10 numbers in the Fibonacci sequence"
    print(f"\nTesting query: {query}")
    print("-" * 40)
    
    result = run_agent_system(query, user_id="test_user", session_id="test_session")
    print(f"Result: {result}")
    
    return result

def test_web_search_tools():
    """Test web search functionality"""
    print("=" * 60)
    print("Testing Web Search Tools Integration")
    print("=" * 60)
    
    # Test web search
    query = "What is the latest news about artificial intelligence?"
    print(f"\nTesting query: {query}")
    print("-" * 40)
    
    result = run_agent_system(query, user_id="test_user", session_id="test_session")
    print(f"Result: {result}")
    
    return result

if __name__ == "__main__":
    print("Starting Tool Integration Tests...")
    
    try:
        # Test retrieval tools
        test_retrieval_tools()
        
        print("\n" + "=" * 60)
        input("Press Enter to continue to execution tools test...")
        
        # Test execution tools
        test_execution_tools()
        
        print("\n" + "=" * 60)
        input("Press Enter to continue to web search tools test...")
        
        # Test web search tools
        test_web_search_tools()
        
        print("\n" + "=" * 60)
        print("Tool integration tests completed!")
        
    except Exception as e:
        print(f"Test failed with error: {e}")
        import traceback
        traceback.print_exc()