test_system.py

"""
Test script to verify RAG system is working correctly
"""
import requests
import time
import sys

API_BASE = "http://localhost:8000"

def print_status(message, success=True):
    """Print colored status message"""
    color = "\033[92m" if success else "\033[91m"
    reset = "\033[0m"
    symbol = "✓" if success else "✗"
    print(f"{color}{symbol} {message}{reset}")

def test_api_connection():
    """Test if API is running"""
    try:
        response = requests.get(f"{API_BASE}/", timeout=2)
        if response.status_code == 200:
            print_status("API is running")
            return True
        else:
            print_status("API returned non-200 status", False)
            return False
    except requests.exceptions.ConnectionError:
        print_status("API is not running. Start it with: python api.py", False)
        return False
    except Exception as e:
        print_status(f"Error connecting to API: {e}", False)
        return False

def test_status_endpoint():
    """Test status endpoint"""
    try:
        response = requests.get(f"{API_BASE}/status")
        if response.status_code == 200:
            data = response.json()
            print_status("Status endpoint working")
            print(f"  - RAG initialized: {data['details']['rag_initialized']}")
            print(f"  - Index loaded: {data['details']['index_loaded']}")
            return True
        else:
            print_status("Status endpoint failed", False)
            return False
    except Exception as e:
        print_status(f"Error testing status: {e}", False)
        return False

def test_chunks_endpoint():
    """Test chunks endpoint"""
    try:
        response = requests.get(f"{API_BASE}/chunks")
        if response.status_code == 200:
            data = response.json()
            chunk_count = data['total']
            print_status(f"Chunks endpoint working - {chunk_count} chunks available")
            
            if chunk_count == 0:
                print("  ℹ️  No chunks found. Upload a PDF to test further.")
            else:
                # Test viewing specific chunk
                chunk_response = requests.get(f"{API_BASE}/chunks/0")
                if chunk_response.status_code == 200:
                    print_status("Can view individual chunks")
                
            return True
        else:
            print_status("Chunks endpoint failed", False)
            return False
    except Exception as e:
        print_status(f"Error testing chunks: {e}", False)
        return False

def test_statistics_endpoint():
    """Test statistics endpoint"""
    try:
        response = requests.get(f"{API_BASE}/chunks/statistics/summary")
        if response.status_code == 200:
            stats = response.json()
            print_status("Statistics endpoint working")
            if 'total_chunks' in stats:
                print(f"  - Total chunks: {stats['total_chunks']}")
                print(f"  - Avg length: {stats['avg_chunk_length']:.0f}")
            return True
        elif response.status_code == 404:
            print_status("Statistics endpoint working (no data yet)")
            return True
        else:
            print_status("Statistics endpoint failed", False)
            return False
    except Exception as e:
        print_status(f"Error testing statistics: {e}", False)
        return False

def test_search_endpoint():
    """Test search endpoint"""
    try:
        response = requests.get(f"{API_BASE}/chunks/search/text", 
                               params={'query': 'test'})
        if response.status_code == 200:
            data = response.json()
            print_status(f"Search endpoint working - {data['total_results']} results")
            return True
        else:
            print_status("Search endpoint failed", False)
            return False
    except Exception as e:
        print_status(f"Error testing search: {e}", False)
        return False

def test_embeddings_endpoint():
    """Test embeddings endpoint"""
    try:
        response = requests.get(f"{API_BASE}/embeddings", params={'limit': 1})
        if response.status_code == 200:
            data = response.json()
            print_status("Embeddings endpoint working")
            if data.get('total_embeddings', 0) > 0:
                print(f"  - Total embeddings: {data['total_embeddings']}")
                print(f"  - Model: {data['embedding_model']}")
                print(f"  - Dimension: {data['embedding_dimension']}")
            return True
        elif response.status_code == 404:
            print_status("Embeddings endpoint working (no data yet)")
            return True
        else:
            print_status("Embeddings endpoint failed", False)
            return False
    except Exception as e:
        print_status(f"Error testing embeddings: {e}", False)
        return False

def test_query_endpoint():
    """Test query endpoint"""
    try:
        # First check if index is loaded
        status_response = requests.get(f"{API_BASE}/status")
        status = status_response.json()
        
        if not status['details']['index_loaded']:
            print_status("Query endpoint skipped (no PDF ingested yet)")
            print("  ℹ️  Upload a PDF first to test querying")
            return True
        
        # Try a simple query
        response = requests.post(f"{API_BASE}/query", json={
            'question': 'What is this document about?',
            'top_k': 3
        })
        
        if response.status_code == 200:
            result = response.json()
            print_status("Query endpoint working")
            print(f"  - Answer length: {len(result['answer'])} chars")
            print(f"  - Sources: {len(result['sources'])} chunks")
            return True
        else:
            print_status("Query endpoint failed", False)
            return False
    except Exception as e:
        print_status(f"Error testing query: {e}", False)
        return False

def test_chunk_viewer():
    """Test ChunkViewer class directly"""
    try:
        from chunk_viewer import ChunkViewer
        
        viewer = ChunkViewer()
        
        # Test loading chunks
        chunks_data = viewer.load_chunks()
        if 'error' not in chunks_data:
            print_status("ChunkViewer can load chunks")
            print(f"  - Total chunks: {chunks_data['total_chunks']}")
        else:
            print_status("ChunkViewer working (no data yet)")
        
        # Test loading embeddings
        embeddings_data = viewer.load_embeddings()
        if 'error' not in embeddings_data:
            print_status("ChunkViewer can load embeddings")
            print(f"  - Total embeddings: {embeddings_data['total_embeddings']}")
        else:
            print_status("ChunkViewer working (no embeddings yet)")
        
        return True
    except Exception as e:
        print_status(f"Error testing ChunkViewer: {e}", False)
        return False

def main():
    print("\n" + "="*60)
    print("RAG SYSTEM TEST SUITE")
    print("="*60 + "\n")
    
    print("Testing API Endpoints:")
    print("-" * 60)
    
    # Test API connection first
    if not test_api_connection():
        print("\n❌ API is not running. Please start it first:")
        print("   python api.py")
        print("\nOr use the startup scripts:")
        print("   ./start_both.sh  (Linux/Mac)")
        print("   start_both.bat   (Windows)")
        sys.exit(1)
    
    # Run all tests
    tests = [
        ("Status Endpoint", test_status_endpoint),
        ("Chunks Endpoint", test_chunks_endpoint),
        ("Statistics Endpoint", test_statistics_endpoint),
        ("Search Endpoint", test_search_endpoint),
        ("Embeddings Endpoint", test_embeddings_endpoint),
        ("Query Endpoint", test_query_endpoint),
    ]
    
    results = []
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append(result)
        except Exception as e:
            print_status(f"{test_name} crashed: {e}", False)
            results.append(False)
        print()
    
    print("-" * 60)
    print("\nTesting Direct Components:")
    print("-" * 60)
    
    chunk_viewer_result = test_chunk_viewer()
    results.append(chunk_viewer_result)
    print()
    
    # Summary
    print("="*60)
    passed = sum(results)
    total = len(results)
    
    if passed == total:
        print(f"✅ ALL TESTS PASSED ({passed}/{total})")
        print("\nYour RAG system is working correctly! 🎉")
        print("\nNext steps:")
        print("  1. Open Streamlit UI: http://localhost:8501")
        print("  2. Upload a PDF and test the full workflow")
        print("  3. Check API docs: http://localhost:8000/docs")
    else:
        print(f"⚠️  SOME TESTS FAILED ({passed}/{total} passed)")
        print("\nPlease check the errors above and:")
        print("  1. Ensure the API is running: python api.py")
        print("  2. Check your .env file has OPENAI_API_KEY")
        print("  3. Verify all dependencies are installed")
    
    print("="*60 + "\n")
    
    return passed == total

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)