chore: import upstream snapshot with attribution

2026-07-13 13:30:03 +08:00
commit ec436095dd
1232 changed files with 404407 additions and 0 deletions
@@ -0,0 +1,250 @@
+"""
+Model Discovery Utilities
+
+Shared functions for discovering and registering new models across different commands.
+"""
+
+from typing import List, Optional, Tuple
+from pathlib import Path
+from rich.console import Console
+
+from kt_kernel.cli.utils.model_scanner import (
+    discover_models,
+    scan_directory_for_models,
+    ScannedModel,
+)
+from kt_kernel.cli.utils.user_model_registry import UserModelRegistry, UserModel
+
+
+console = Console()
+
+
+def discover_and_register_global(
+    min_size_gb: float = 2.0, max_depth: int = 6, show_progress: bool = True, lang: str = "en"
+) -> Tuple[int, int, List[UserModel]]:
+    """
+    Perform global model discovery and register new models.
+
+    Args:
+        min_size_gb: Minimum model size in GB
+        max_depth: Maximum search depth
+        show_progress: Whether to show progress messages
+        lang: Language for messages ("en" or "zh")
+
+    Returns:
+        Tuple of (total_found, new_found, registered_models)
+    """
+    registry = UserModelRegistry()
+
+    if show_progress:
+        if lang == "zh":
+            console.print("[dim]正在扫描系统中的模型权重，这可能需要30-60秒...[/dim]")
+        else:
+            console.print("[dim]Scanning system for model weights, this may take 30-60 seconds...[/dim]")
+
+    # Global scan
+    all_models = discover_models(mount_points=None, min_size_gb=min_size_gb, max_depth=max_depth)
+
+    # Filter out existing models
+    new_models = []
+    for model in all_models:
+        if not registry.find_by_path(model.path):
+            new_models.append(model)
+
+    # Register new models
+    registered = []
+    for model in new_models:
+        user_model = _create_and_register_model(registry, model)
+        if user_model:
+            registered.append(user_model)
+
+    return len(all_models), len(new_models), registered
+
+
+def discover_and_register_path(
+    path: str,
+    min_size_gb: float = 2.0,
+    existing_paths: Optional[set] = None,
+    show_progress: bool = True,
+    lang: str = "en",
+) -> Tuple[int, int, List[UserModel]]:
+    """
+    Discover models in a specific path and register new ones.
+
+    Args:
+        path: Directory path to scan
+        min_size_gb: Minimum model file size in GB
+        existing_paths: Set of already discovered paths in this session (optional)
+        show_progress: Whether to show progress messages
+        lang: Language for messages ("en" or "zh")
+
+    Returns:
+        Tuple of (total_found, new_found, registered_models)
+    """
+    registry = UserModelRegistry()
+
+    if show_progress:
+        if lang == "zh":
+            console.print(f"[dim]正在扫描 {path}...[/dim]")
+        else:
+            console.print(f"[dim]Scanning {path}...[/dim]")
+
+    # Scan directory
+    model_info = scan_directory_for_models(path, min_file_size_gb=min_size_gb)
+
+    if not model_info:
+        return 0, 0, []
+
+    # Convert to ScannedModel and filter
+    new_models = []
+    for dir_path, (format_type, size_bytes, file_count, files) in model_info.items():
+        # Check if already in registry
+        if registry.find_by_path(dir_path):
+            continue
+
+        # Check if already discovered in this session
+        if existing_paths and dir_path in existing_paths:
+            continue
+
+        model = ScannedModel(
+            path=dir_path, format=format_type, size_bytes=size_bytes, file_count=file_count, files=files
+        )
+        new_models.append(model)
+
+    # Register new models
+    registered = []
+    for model in new_models:
+        user_model = _create_and_register_model(registry, model)
+        if user_model:
+            registered.append(user_model)
+
+    return len(model_info), len(new_models), registered
+
+
+def _create_and_register_model(registry: UserModelRegistry, scanned_model: ScannedModel) -> Optional[UserModel]:
+    """
+    Create a UserModel from ScannedModel and register it.
+
+    Handles name conflicts by suggesting a unique name (e.g., model-2, model-3).
+    Automatically detects repo_id from README.md YAML frontmatter.
+    Automatically detects and caches MoE information for safetensors models.
+
+    Args:
+        registry: UserModelRegistry instance
+        scanned_model: ScannedModel to register
+
+    Returns:
+        Registered UserModel or None if failed
+    """
+    # Use suggest_name to get a unique name (adds -2, -3, etc. if needed)
+    unique_name = registry.suggest_name(scanned_model.folder_name)
+
+    user_model = UserModel(name=unique_name, path=scanned_model.path, format=scanned_model.format)
+
+    # Auto-detect repo_id from README.md (only YAML frontmatter)
+    try:
+        from kt_kernel.cli.utils.repo_detector import detect_repo_for_model
+
+        repo_info = detect_repo_for_model(scanned_model.path)
+        if repo_info:
+            repo_id, repo_type = repo_info
+            user_model.repo_id = repo_id
+            user_model.repo_type = repo_type
+    except Exception:
+        # Silently continue if detection fails
+        pass
+
+    # Auto-detect MoE information for safetensors models
+    if scanned_model.format == "safetensors":
+        try:
+            from kt_kernel.cli.utils.analyze_moe_model import analyze_moe_model
+
+            moe_result = analyze_moe_model(scanned_model.path, use_cache=True)
+            if moe_result and moe_result.get("is_moe"):
+                user_model.is_moe = True
+                user_model.moe_num_experts = moe_result.get("num_experts")
+                user_model.moe_num_experts_per_tok = moe_result.get("num_experts_per_tok")
+            else:
+                user_model.is_moe = False
+        except Exception:
+            # Silently continue if MoE detection fails
+            # is_moe will remain None
+            pass
+
+    try:
+        registry.add_model(user_model)
+        return user_model
+    except Exception:
+        # Should not happen since we used suggest_name, but handle gracefully
+        return None
+
+
+def format_discovery_summary(
+    total_found: int,
+    new_found: int,
+    registered: List[UserModel],
+    lang: str = "en",
+    show_models: bool = True,
+    max_show: int = 10,
+) -> None:
+    """
+    Print formatted discovery summary.
+
+    Args:
+        total_found: Total models found
+        new_found: New models found
+        registered: List of registered UserModel objects
+        lang: Language ("en" or "zh")
+        show_models: Whether to show model list
+        max_show: Maximum models to show
+    """
+    console.print()
+
+    if new_found == 0:
+        if total_found > 0:
+            if lang == "zh":
+                console.print(f"[green]✓[/green] 扫描完成：找到 {total_found} 个模型，所有模型均已在列表中")
+            else:
+                console.print(f"[green]✓[/green] Scan complete: found {total_found} models, all already in the list")
+        else:
+            if lang == "zh":
+                console.print("[yellow]未找到模型[/yellow]")
+            else:
+                console.print("[yellow]No models found[/yellow]")
+        return
+
+    # Show summary
+    if lang == "zh":
+        console.print(f"[green]✓[/green] 扫描完成：找到 {total_found} 个模型，其中 {new_found} 个为新模型")
+    else:
+        console.print(f"[green]✓[/green] Scan complete: found {total_found} models, {new_found} are new")
+
+    # Show registered count
+    if len(registered) > 0:
+        if lang == "zh":
+            console.print(f"[green]✓[/green] 成功添加 {len(registered)} 个新模型到列表")
+        else:
+            console.print(f"[green]✓[/green] Successfully added {len(registered)} new models to list")
+
+    # Show model list
+    if show_models and registered:
+        console.print()
+        if lang == "zh":
+            console.print(f"[dim]新发现的模型（前{max_show}个）:[/dim]")
+        else:
+            console.print(f"[dim]Newly discovered models (first {max_show}):[/dim]")
+
+        for i, model in enumerate(registered[:max_show], 1):
+            # Get size from registry or estimate
+            size_str = "?.? GB"
+            # Try to find the ScannedModel to get size
+            # For now just show name and path
+            console.print(f"  {i}. {model.name} ({model.format})")
+            console.print(f"     [dim]{model.path}[/dim]")
+
+        if len(registered) > max_show:
+            remaining = len(registered) - max_show
+            if lang == "zh":
+                console.print(f"  [dim]... 还有 {remaining} 个新模型[/dim]")
+            else:
+                console.print(f"  [dim]... and {remaining} more new models[/dim]")