Function _search_model

src/whichllm/cli.py:993–1050

Search for a model by name/ID. Returns single model or exits.

(models: list, model_name: str)

def _search_model(models: list, model_name: str):
    """Search for a model by name/ID. Returns single model or exits."""
    query_lower = model_name.lower()
    terms = query_lower.split()
    size_b = None

    matches = [m for m in models if m.id.lower() == query_lower]
    if not matches:
        matches = [m for m in models if m.id.lower().endswith("/" + query_lower)]
    if not matches:
        text_terms, size_b = _parse_size_tokens(terms)
        matches = [
            m
            for m in models
            if all(t in m.id.lower() for t in text_terms)
            and (size_b is None or _size_compatible(m, size_b))
        ]

    if not matches:
        console.print(f"[red]No model found matching '{model_name}'.[/]")
        suggestions = [m for m in models if any(t in m.id.lower() for t in terms)]
        if suggestions:
            suggestions.sort(key=lambda m: m.downloads, reverse=True)
            console.print("\n[yellow]Did you mean:[/]")
            for m in suggestions[:5]:
                p = (
                    f"{m.parameter_count / 1e9:.1f}B"
                    if m.parameter_count >= 1e9
                    else f"{m.parameter_count / 1e6:.0f}M"
                )
                console.print(f"  • {m.id} ({p})")
        raise typer.Exit(code=1)

    if size_b is not None:

        def _sort_key(m: ModelInfo) -> tuple:
            id_size = _extract_id_size_b(m.id)
            has_id_size = id_size is not None
            id_dist = abs(id_size - size_b) if has_id_size else float("inf")
            pc_dist = (
                abs(m.parameter_count / 1e9 - size_b)
                if m.parameter_count > 0
                else float("inf")
            )
            return (
                0 if (has_id_size or m.parameter_count > 0) else 1,
                id_dist,
                pc_dist,
                -m.downloads,
            )

        matches.sort(key=_sort_key)
    else:
        matches.sort(key=lambda m: m.downloads, reverse=True)
    model = matches[0]
    if len(matches) > 1:
        console.print(f"[dim]Found {len(matches)} matches, using: {model.id}[/]")
    return model