Improved UX of HF token

Author: UtkarshTheDev

locallab/cli/interactive.py

@@ -269,30 +269,49 @@ def prompt_for_config(use_ngrok: bool = None, port: int = None, ngrok_auth_token
             os.environ["LOCALLAB_LOG_FILE"] = log_file
             config["log_file"] = log_file
 
-    # Ask about HuggingFace token
+    # Ask about HuggingFace token with improved UX
     hf_token = config.get("huggingface_token") or os.environ.get("HUGGINGFACE_TOKEN")
     if not hf_token or force_reconfigure:
-        click.echo("\n🔑 Enter your HuggingFace token (optional)")
-        click.echo("   Get your token from: https://huggingface.co/settings/tokens")
-        click.echo("   Press Enter to skip or paste your token (it will be hidden): ", nl=False)
+        click.echo("\n�� HuggingFace Token Configuration")
+        click.echo("───────────────────────────────")
+        click.echo("A token is required to download models like microsoft/phi-2")
+        click.echo("Get your token from: https://huggingface.co/settings/tokens")
 
-        # Read token character by character
+        if hf_token:
+            click.echo(f"\nCurrent token: {hf_token[:4]}...{hf_token[-4:]}")
+            if not click.confirm("Would you like to update your token?", default=False):
+                click.echo("Keeping existing token...")
+                return config
+        
+        click.echo("\nEnter your HuggingFace token (press Enter to skip): ", nl=False)
+        
+        # Read token character by character for secure input
         chars = []
         while True:
             char = click.getchar()
             if char in ('\r', '\n'):
                 break
             chars.append(char)
-            click.echo('*', nl=False)  # Show * for each character
+            click.echo('*', nl=False)
 
         hf_token = ''.join(chars)
 
         if hf_token:
-            click.echo("\n✅ Token saved!")
+            # Validate token format
+            if len(hf_token) < 20:  # Basic validation
+                click.echo("\n❌ Invalid token format. Token should be longer than 20 characters.")
+                click.echo("Please check your token and try again.")
+                return config
+                
+            click.echo("\n✅ Token saved successfully!")
             os.environ["HUGGINGFACE_TOKEN"] = hf_token
             config["huggingface_token"] = hf_token
+            
+            # Save immediately to ensure it's persisted
+            from .config import save_config
+            save_config(config)
         else:
-            click.echo("\nSkipping HuggingFace token...")
-    
+            click.echo("\n⚠️  No token provided. Some models may not be accessible.")
+
     click.echo("\n✅ Configuration complete!\n")
     return config

locallab/config.py

@@ -531,7 +531,7 @@ def get_hf_token(interactive: bool = False) -> Optional[str]:
     # If interactive and still no token, prompt user
     if not token and interactive:
         try:
-            click.echo("\n🔑 HuggingFace token is recommended for better model access.")
+            click.echo("\n🔑 HuggingFace token is required for accessing this model.")
             click.echo("Get your token from: https://huggingface.co/settings/tokens")
             click.echo("Enter token (press Enter to skip): ", nl=False)
 
@@ -546,8 +546,16 @@ def get_hf_token(interactive: bool = False) -> Optional[str]:
 
             token = ''.join(chars)
             if token:
+                if len(token) < 20:  # Basic validation
+                    click.echo("\n❌ Invalid token format. Please check your token.")
+                    return None
+                    
                 click.echo("\n✅ Token saved!")
                 os.environ["HUGGINGFACE_TOKEN"] = token
+                
+                # Save to config
+                from .cli.config import set_config_value
+                set_config_value("huggingface_token", token)
             else:
                 click.echo("\nSkipping token...")
         except:

locallab/model_manager.py

@@ -255,8 +255,17 @@ async def load_model(self, model_id: str) -> bool:
             start_time = time.time()
             logger.info(f"\n{Fore.CYAN}Loading model: {model_id}{Style.RESET_ALL}")
 
+            # Get and validate HuggingFace token
             from .config import get_hf_token
-            hf_token = get_hf_token(interactive=True)
+            hf_token = get_hf_token(interactive=False)  # Don't prompt during model loading
+            
+            if not hf_token and model_id in ["microsoft/phi-2"]:  # Add other gated models here
+                logger.error(f"{Fore.RED}This model requires authentication. Please configure your HuggingFace token first.{Style.RESET_ALL}")
+                logger.info(f"{Fore.YELLOW}You can set your token by running: locallab config{Style.RESET_ALL}")
+                raise HTTPException(
+                    status_code=401,
+                    detail="HuggingFace token required for this model. Run 'locallab config' to set up."
+                )
 
             if self.model is not None:
                 prev_model = self.current_model

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="locallab",
-    version="0.4.29",
+    version="0.4.30",
     packages=find_packages(include=["locallab", "locallab.*"]),
     install_requires=[
         "fastapi>=0.95.0,<1.0.0",