Fixed ngrok issues

UtkarshTheDev · UtkarshTheDev · commit 0c1bc4668734 · 2025-03-15T13:28:23.000+05:30
diff --git a/locallab/cli/interactive.py b/locallab/cli/interactive.py
@@ -43,7 +43,7 @@ def get_missing_required_env_vars() -> List[str]:
     
     return missing
 
-def prompt_for_config(use_ngrok: bool = None, port: int = None, ngrok_auth_token: str = None, force_reconfigure: bool = False) -> Dict[str, Any]:
+def prompt_for_config(use_ngrok: bool = None, port: int = None, ngrok_auth_token: Optional[str] = None, force_reconfigure: bool = False) -> Dict[str, Any]:
     """
     Interactive prompt for configuration
     """
@@ -59,105 +59,153 @@ def prompt_for_config(use_ngrok: bool = None, port: int = None, ngrok_auth_token
     # Override with provided parameters
     if use_ngrok is not None:
         config["use_ngrok"] = use_ngrok
+        # Set environment variable for use_ngrok
+        os.environ["LOCALLAB_USE_NGROK"] = str(use_ngrok).lower()
+        
     if port is not None:
         config["port"] = port
+        os.environ["LOCALLAB_PORT"] = str(port)
+        
     if ngrok_auth_token is not None:
         config["ngrok_auth_token"] = ngrok_auth_token
+        os.environ["NGROK_AUTHTOKEN"] = ngrok_auth_token
     
     # Determine if we're in Colab
     in_colab = is_in_colab()
     
-    # If in Colab, use simplified configuration
-    if in_colab:
-        # Set default values for Colab environment
-        config.setdefault("port", 8000)
-        config.setdefault("use_ngrok", True)
-        config.setdefault("model_id", os.environ.get("HUGGINGFACE_MODEL", DEFAULT_MODEL))
-        
-        # Use ngrok token from environment if available
-        if os.environ.get("NGROK_AUTH_TOKEN"):
-            config["ngrok_auth_token"] = os.environ.get("NGROK_AUTH_TOKEN")
-        elif ngrok_auth_token:
-            config["ngrok_auth_token"] = ngrok_auth_token
-            
-        # Set some reasonable defaults for Colab
-        config.setdefault("enable_quantization", True)
-        config.setdefault("quantization_type", "int8")
-        config.setdefault("enable_attention_slicing", True)
-        config.setdefault("enable_flash_attention", True)
-        config.setdefault("enable_better_transformer", True)
-        
-        return config
-    
-    # Check for GPU
-    has_gpu = False
-    gpu_memory = get_gpu_memory()
-    if gpu_memory:
-        has_gpu = True
-        total_gpu_memory, free_gpu_memory = gpu_memory
-        click.echo(f"🎮 GPU detected with {free_gpu_memory}MB free of {total_gpu_memory}MB total")
-    else:
-        click.echo("⚠️ No GPU detected. Running on CPU will be significantly slower.")
-    
-    # Get system memory
-    total_memory, free_memory = get_system_memory()
-    click.echo(f"💾 System memory: {free_memory}MB free of {total_memory}MB total")
-    
-    # Check for missing required environment variables
-    missing_vars = get_missing_required_env_vars()
-    
-    # Check if we have all required configuration and not forcing reconfiguration
-    has_model = "model_id" in config or os.environ.get("HUGGINGFACE_MODEL") or os.environ.get("DEFAULT_MODEL")
-    has_port = "port" in config or port is not None
-    has_ngrok_config = not in_colab or not config.get("use_ngrok", use_ngrok) or "ngrok_auth_token" in config or ngrok_auth_token is not None or os.environ.get("NGROK_AUTH_TOKEN")
-    
-    # If we have all required config and not forcing reconfiguration, return early
-    if not force_reconfigure and has_model and has_port and has_ngrok_config and not missing_vars:
-        # Ensure port is set in config
-        if "port" not in config and port is not None:
-            config["port"] = port
-        # Ensure use_ngrok is set in config
-        if "use_ngrok" not in config and use_ngrok is not None:
-            config["use_ngrok"] = use_ngrok
-        # Ensure ngrok_auth_token is set in config if needed
-        if config.get("use_ngrok", False) and "ngrok_auth_token" not in config and ngrok_auth_token is not None:
-            config["ngrok_auth_token"] = ngrok_auth_token
-        
-        return config
-    
+    # If in Colab, ensure ngrok is enabled by default
+    if in_colab and "use_ngrok" not in config:
+        config["use_ngrok"] = True
+        os.environ["LOCALLAB_USE_NGROK"] = "true"
+
     click.echo("\n🚀 Welcome to LocalLab! Let's set up your server.\n")
     
-    # Always ask for model when reconfiguring or if not provided
+    # Basic Configuration
+    # ------------------
+    click.echo("\n📋 Basic Configuration")
+    click.echo("─────────────────────")
+    
+    # Model selection
     model_id = click.prompt(
         "📦 Which model would you like to use?",
         default=config.get("model_id", DEFAULT_MODEL)
     )
-    os.environ["HUGGINGFACE_MODEL"] = model_id
     config["model_id"] = model_id
     
-    # Always ask for port when reconfiguring or if not provided
+    # Port configuration
     port = click.prompt(
         "🔌 Which port would you like to run on?",
         default=config.get("port", 8000),
         type=int
     )
     config["port"] = port
     
-    # Ask about ngrok
+    # Model Optimization Settings
+    # -------------------------
+    click.echo("\n⚡ Model Optimization Settings")
+    click.echo("─────────────────────────────")
+    
+    config["enable_quantization"] = click.confirm(
+        "Enable model quantization?",
+        default=config.get("enable_quantization", ENABLE_QUANTIZATION)
+    )
+    
+    if config["enable_quantization"]:
+        config["quantization_type"] = click.prompt(
+            "Quantization type (fp16/int8/int4)",
+            default=config.get("quantization_type", QUANTIZATION_TYPE),
+            type=click.Choice(["fp16", "int8", "int4"])
+        )
+    
+    config["enable_cpu_offloading"] = click.confirm(
+        "Enable CPU offloading?",
+        default=config.get("enable_cpu_offloading", ENABLE_CPU_OFFLOADING)
+    )
+    
+    config["enable_attention_slicing"] = click.confirm(
+        "Enable attention slicing?",
+        default=config.get("enable_attention_slicing", ENABLE_ATTENTION_SLICING)
+    )
+    
+    config["enable_flash_attention"] = click.confirm(
+        "Enable flash attention?",
+        default=config.get("enable_flash_attention", ENABLE_FLASH_ATTENTION)
+    )
+    
+    config["enable_better_transformer"] = click.confirm(
+        "Enable better transformer?",
+        default=config.get("enable_bettertransformer", ENABLE_BETTERTRANSFORMER)
+    )
+    
+    # Advanced Settings
+    # ----------------
+    click.echo("\n⚙️ Advanced Settings")
+    click.echo("──────────────────")
+    
+    config["model_timeout"] = click.prompt(
+        "Model timeout (seconds)",
+        default=config.get("model_timeout", 3600),
+        type=int
+    )
+    
+    # Cache Settings
+    # -------------
+    click.echo("\n💾 Cache Settings")
+    click.echo("────────────────")
+    
+    config["enable_cache"] = click.confirm(
+        "Enable response caching?",
+        default=config.get("enable_cache", True)
+    )
+    
+    if config["enable_cache"]:
+        config["cache_ttl"] = click.prompt(
+            "Cache TTL (seconds)",
+            default=config.get("cache_ttl", 3600),
+            type=int
+        )
+    
+    # Logging Settings
+    # ---------------
+    click.echo("\n📝 Logging Settings")
+    click.echo("──────────────────")
+    
+    config["log_level"] = click.prompt(
+        "Log level",
+        default=config.get("log_level", "INFO"),
+        type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"])
+    )
+    
+    config["enable_file_logging"] = click.confirm(
+        "Enable file logging?",
+        default=config.get("enable_file_logging", False)
+    )
+    
+    if config["enable_file_logging"]:
+        config["log_file"] = click.prompt(
+            "Log file path",
+            default=config.get("log_file", "locallab.log")
+        )
+    
+    # Ngrok Configuration
+    # ------------------
+    click.echo("\n🌐 Ngrok Configuration")
+    click.echo("────────────────────")
+    
     use_ngrok = click.confirm(
-        "🌐 Do you want to enable public access via ngrok?",
+        "Enable public access via ngrok?",
         default=config.get("use_ngrok", in_colab)
     )
     config["use_ngrok"] = use_ngrok
+    os.environ["LOCALLAB_USE_NGROK"] = str(use_ngrok).lower()
     
     if use_ngrok:
-        # Show current token if exists
         current_token = config.get("ngrok_auth_token") or get_env_var(NGROK_TOKEN_ENV)
         if current_token:
             click.echo(f"\nCurrent ngrok token: {current_token}")
             
         ngrok_auth_token = click.prompt(
-            "🔑 Enter your ngrok auth token (get one at https://dashboard.ngrok.com/get-started/your-authtoken)",
+            "Enter your ngrok auth token (get one at https://dashboard.ngrok.com/get-started/your-authtoken)",
             default=current_token,
             type=str,
             show_default=True
@@ -166,18 +214,26 @@ def prompt_for_config(use_ngrok: bool = None, port: int = None, ngrok_auth_token
         if ngrok_auth_token:
             token_str = str(ngrok_auth_token).strip()
             config["ngrok_auth_token"] = token_str
-            set_env_var(NGROK_TOKEN_ENV, token_str)
-            click.echo(f"✅ Ngrok token saved: {token_str}")
+            # Set both environment variables to ensure compatibility
+            os.environ["NGROK_AUTHTOKEN"] = token_str
+            os.environ["LOCALLAB_NGROK_AUTH_TOKEN"] = token_str
+            
+            # Save immediately to ensure persistence
+            from .config import save_config
+            save_config(config)
+            click.echo(f"✅ Ngrok token saved and activated")
+
+    # HuggingFace Token
+    # ----------------
+    click.echo("\n🤗 HuggingFace Token")
+    click.echo("──────────────────")
     
-    # Ask about HuggingFace token
     current_hf_token = config.get("huggingface_token") or get_env_var(HF_TOKEN_ENV)
     if current_hf_token:
-        click.echo(f"\nCurrent HuggingFace token: {current_hf_token}")
+        click.echo(f"Current HuggingFace token: {current_hf_token}")
         
     if not current_hf_token or force_reconfigure:
-        click.echo("\n🔑 HuggingFace Token Configuration")
-        click.echo("───────────────────────────────")
-        click.echo("A token is required to download models like microsoft/phi-2")
+        click.echo("\nA token is required to download models.")
         click.echo("Get your token from: https://huggingface.co/settings/tokens")
         
         hf_token = click.prompt(
@@ -190,16 +246,14 @@ def prompt_for_config(use_ngrok: bool = None, port: int = None, ngrok_auth_token
         if hf_token:
             if len(hf_token) < 20:
                 click.echo("❌ Invalid token format. Token should be longer than 20 characters.")
-                return config
-            
-            token_str = str(hf_token).strip()
-            config["huggingface_token"] = token_str
-            set_env_var(HF_TOKEN_ENV, token_str)
-            click.echo(f"✅ HuggingFace token saved: {token_str}")
-            
-            # Save immediately
-            from .config import save_config
-            save_config(config)
+            else:
+                token_str = str(hf_token).strip()
+                config["huggingface_token"] = token_str
+                set_env_var(HF_TOKEN_ENV, token_str)
+                
+                # Save immediately
+                from .config import save_config
+                save_config(config)
         else:
             click.echo("\n⚠️  No token provided. Some models may not be accessible.")
 
diff --git a/locallab/server.py b/locallab/server.py
@@ -639,81 +639,60 @@ def _initialize_lifespan(self):
 
 
 def start_server(use_ngrok: bool = None, port: int = None, ngrok_auth_token: Optional[str] = None):
-    
     try:
         set_server_status("initializing")
         
         print_initializing_banner(__version__)
         
+        # Load configuration
         from .cli.config import load_config, set_config_value
         
         try:
             saved_config = load_config()
         except Exception as e:
             logger.warning(f"Error loading configuration: {str(e)}. Using defaults.")
             saved_config = {}
+            
+        # Set up ngrok configuration
+        use_ngrok = (
+            use_ngrok if use_ngrok is not None 
+            else saved_config.get("use_ngrok", False) 
+            or os.environ.get("LOCALLAB_USE_NGROK", "").lower() == "true"
+        )
         
-        for key, value in saved_config.items():
-            if key == "model_id":
-                os.environ["HUGGINGFACE_MODEL"] = str(value)
-            elif key == "ngrok_auth_token":
-                os.environ["NGROK_AUTH_TOKEN"] = str(value)
-            elif key == "huggingface_token":
-                os.environ["HUGGINGFACE_TOKEN"] = str(value)
-            elif key in ["enable_quantization", "enable_attention_slicing", "enable_flash_attention", 
-                        "enable_better_transformer", "enable_cpu_offloading", "enable_cache", 
-                        "enable_file_logging"]:
-                env_key = f"LOCALLAB_{key.upper()}"
-                os.environ[env_key] = str(value).lower()
-            elif key in ["quantization_type", "model_timeout", "cache_ttl", "log_level", "log_file"]:
-                env_key = f"LOCALLAB_{key.upper()}"
-                os.environ[env_key] = str(value)
-        
-        config = prompt_for_config(use_ngrok, port, ngrok_auth_token)
-        
-        save_config(config)
-        
-        use_ngrok = config.get("use_ngrok", use_ngrok)
-        port = config.get("port", port or 8000)
-        ngrok_auth_token = config.get("ngrok_auth_token", ngrok_auth_token)
-        
-        if is_port_in_use(port):
-            logger.warning(f"Port {port} is already in use. Trying to find another port...")
-            for p in range(port+1, port+100):
-                if not is_port_in_use(p):
-                    port = p
-                    logger.info(f"Using alternative port: {port}")
-                    break
-            else:
-                raise RuntimeError(f"Could not find an available port in range {port}-{port+100}")
+        # Get port configuration
+        port = port or saved_config.get("port", None) or int(os.environ.get("LOCALLAB_PORT", "8000"))
         
+        # Handle ngrok auth token
+        if ngrok_auth_token:
+            os.environ["NGROK_AUTHTOKEN"] = ngrok_auth_token
+        elif saved_config.get("ngrok_auth_token"):
+            os.environ["NGROK_AUTHTOKEN"] = saved_config["ngrok_auth_token"]
+            
+        # Set up ngrok if enabled
         public_url = None
         if use_ngrok:
             os.environ["LOCALLAB_USE_NGROK"] = "true"
             
-            if not ngrok_auth_token and not os.environ.get("NGROK_AUTH_TOKEN"):
+            if not os.environ.get("NGROK_AUTHTOKEN"):
                 logger.error("Ngrok auth token is required for public access. Please set it in the configuration.")
                 logger.info("You can get a free token from: https://dashboard.ngrok.com/get-started/your-authtoken")
                 raise ValueError("Ngrok auth token is required for public access")
                 
-            logger.info(f"{Fore.CYAN}Setting up ngrok tunnel to port {port}...{Style.RESET_ALL}")
-            public_url = setup_ngrok(port=port, auth_token=ngrok_auth_token or os.environ.get("NGROK_AUTH_TOKEN"))
+            logger.info(f"Setting up ngrok tunnel to port {port}...")
+            public_url = setup_ngrok(port)
+            
             if public_url:
                 os.environ["LOCALLAB_NGROK_URL"] = public_url
                 
                 ngrok_section = f"\n{Fore.CYAN}┌────────────────────────── Ngrok Tunnel Details ─────────────────────────────┐{Style.RESET_ALL}\n│\n│  🚀 Ngrok Public URL: {Fore.GREEN}{public_url}{Style.RESET_ALL}\n│\n{Fore.CYAN}└──────────────────────────────────────────────────────────────────────────────┘{Style.RESET_ALL}\n"
                 print(ngrok_section)
-            
             else:
-
- 
-
-                logger.warning(f"{Fore.YELLOW}Failed to set up ngrok tunnel. Server will run locally on port {port}.{Style.RESET_ALL}")
- 
-
+                logger.error(f"Failed to set up ngrok tunnel. Server will run locally on port {port}.")
+                raise RuntimeError("Failed to set up ngrok tunnel")
         else:
-            # Set environment variable to indicate ngrok is not enabled
             os.environ["LOCALLAB_USE_NGROK"] = "false"
+
         # Set environment variable with the port
         os.environ["LOCALLAB_PORT"] = str(port)
         # Server info section
diff --git a/locallab/utils/networking.py b/locallab/utils/networking.py
