ktransformers/kt-kernel/python/cli/utils/model_discovery.py

"""
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]")