ArthyP/technical-rag-assistant

1#!/usr/bin/env python32"""3HuggingFace Spaces startup script with Ollama support.4Starts Ollama server in background, then launches Streamlit.5"""6 7import os8import sys9import subprocess10import time11import signal12from datetime import datetime13 14 15def log(message):16    """Log message to stderr for visibility in HF Spaces."""17    print(f"[{datetime.now().isoformat()}] {message}", file=sys.stderr, flush=True)18 19 20def start_ollama():21    """Start Ollama server in background."""22    log("🦙 Starting Ollama server...")23    24    try:25        # Start Ollama in background26        ollama_process = subprocess.Popen(27            ["ollama", "serve"],28            stdout=subprocess.PIPE,29            stderr=subprocess.PIPE,30            preexec_fn=os.setsid  # Create new process group31        )32        33        log(f"🦙 Ollama server started with PID {ollama_process.pid}")34        35        # Wait for Ollama to be ready36        log("⏳ Waiting for Ollama to be ready...")37        max_retries = 12  # 60 seconds total38        for attempt in range(max_retries):39            try:40                # Test if Ollama is responding41                result = subprocess.run(42                    ["curl", "-s", "http://localhost:11434/api/tags"],43                    capture_output=True,44                    timeout=545                )46                if result.returncode == 0:47                    log("✅ Ollama server is ready!")48                    49                    # Try to pull the model (optimized for HF Spaces)50                    log("📥 Pulling llama3.2:1b model (optimized for container deployment)...")51                    pull_result = subprocess.run(52                        ["ollama", "pull", "llama3.2:1b"],53                        capture_output=True,54                        timeout=300  # 5 minutes for model download55                    )56                    if pull_result.returncode == 0:57                        log("✅ Model llama3.2:1b ready!")58                    else:59                        log("⚠️ Model pull failed, will download on first use")60                    61                    return ollama_process62                    63            except (subprocess.TimeoutExpired, Exception) as e:64                log(f"🔄 Ollama not ready yet (attempt {attempt + 1}/{max_retries}): {e}")65                time.sleep(5)66                continue67        68        log("❌ Ollama failed to start after 60 seconds")69        ollama_process.terminate()70        return None71        72    except Exception as e:73        log(f"❌ Failed to start Ollama: {e}")74        return None75 76 77def main():78    """Start services and Streamlit based on configuration."""79    log("🚀 Starting Technical RAG Assistant in HuggingFace Spaces...")80 81    # Check which inference method to use82    use_ollama = os.getenv("USE_OLLAMA", "false").lower() == "true"83    use_inference_providers = os.getenv("USE_INFERENCE_PROVIDERS", "false").lower() == "true"84    85    ollama_process = None86    87    # Configure environment variables based on selected inference method88    if use_inference_providers:89        os.environ["USE_INFERENCE_PROVIDERS"] = "true"90        os.environ["USE_OLLAMA"] = "false"91        log("🚀 Using Inference Providers API")92    elif use_ollama:93        os.environ["USE_OLLAMA"] = "true"94        os.environ["USE_INFERENCE_PROVIDERS"] = "false"95        log("🦙 Ollama enabled - starting server...")96        ollama_process = start_ollama()97        98        if ollama_process is None:99            log("🔄 Ollama failed to start, falling back to HuggingFace API")100            os.environ["USE_OLLAMA"] = "false"101            os.environ["USE_INFERENCE_PROVIDERS"] = "false"102    else:103        os.environ["USE_OLLAMA"] = "false"104        os.environ["USE_INFERENCE_PROVIDERS"] = "false"105        log("🤗 Using classic HuggingFace API")106 107    # Start Streamlit108    log("🎯 Starting Streamlit application...")109    110    def signal_handler(signum, frame):111        """Handle shutdown signals."""112        log("🛑 Received shutdown signal, cleaning up...")113        if ollama_process:114            log("🦙 Stopping Ollama server...")115            try:116                os.killpg(os.getpgid(ollama_process.pid), signal.SIGTERM)117            except:118                pass119        sys.exit(0)120    121    # Register signal handlers122    signal.signal(signal.SIGTERM, signal_handler)123    signal.signal(signal.SIGINT, signal_handler)124    125    try:126        subprocess.run(127            [128                "streamlit",129                "run",130                "streamlit_app.py",131                "--server.port=8501",132                "--server.address=0.0.0.0",133                "--server.headless=true",134                "--server.enableCORS=false",135                "--server.enableXsrfProtection=false",136            ],137            check=True,138        )139    except Exception as e:140        log(f"❌ Failed to start Streamlit: {e}")141        sys.exit(1)142    finally:143        # Clean up Ollama if it was started144        if ollama_process:145            log("🦙 Cleaning up Ollama server...")146            try:147                os.killpg(os.getpgid(ollama_process.pid), signal.SIGTERM)148            except:149                pass150 151 152if __name__ == "__main__":153    main()154