setup.py

#!/usr/bin/env python
"""
Setup script for RAG Latency Demo
Run this to initialize the project after cloning.
"""
import subprocess
import sys
import os
from pathlib import Path

def run_command(cmd, description):
    """Run a command and print status."""
    print(f"\n🔧 {description}...")
    try:
        result = subprocess.run(cmd, shell=True, check=True, capture_output=True, text=True)
        print(f"   ✅ Success")
        return True
    except subprocess.CalledProcessError as e:
        print(f"   ❌ Failed: {e.stderr}")
        return False

def main():
    print("=" * 60)
    print("🚀 RAG Latency Demo - Setup Script")
    print("=" * 60)
    
    # Check Python version
    print(f"\n🐍 Python {sys.version}")
    if sys.version_info < (3, 10):
        print("❌ Python 3.10+ required")
        return False
    
    # Create virtual environment
    if not os.path.exists(".venv"):
        if run_command("python -m venv .venv", "Creating virtual environment"):
            print("   📁 Virtual environment created at .venv/")
    
    # Install requirements
    if run_command(".venv\\Scripts\\pip install -r requirements.txt", "Installing dependencies"):
        print("   📦 Dependencies installed")
    
    # Download sample data
    if run_command(".venv\\Scripts\\python scripts/download_sample_data.py", "Downloading sample data"):
        print("   📄 Sample data downloaded to data/")
    
    # Initialize RAG system
    if run_command(".venv\\Scripts\\python scripts/initialize_rag.py", "Initializing RAG system"):
        print("   🗃️  Databases and indexes created")
    
    # Run a quick test
    print("\n🧪 Running quick test...")
    test_cmd = '.venv\\Scripts\\python -c "from app.rag_naive import NaiveRAG; r=NaiveRAG(); r.initialize(); a,c=r.query(\"Test\"); print(f\"✅ System works: {c} chunks\"); r.close()"'
    if run_command(test_cmd, "Testing system"):
        print("   ✅ System is functional")
    
    print("\n" + "=" * 60)
    print("🎉 Setup complete!")
    print("\nNext steps:")
    print("1. Run benchmarks: python working_benchmark.py")
    print("2. Start API server: uvicorn app.main:app --reload")
    print("3. View docs: http://localhost:8000/docs")
    print("\nFor production deployment, see DEPLOYMENT.md")
    print("=" * 60)
    
    return True

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