Coverage for src / lilbee / cli / model.py: 100%
298 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-04-29 19:16 +0000
1"""`lilbee model` sub-app: list/show/pull/rm/browse for installed models.
3Thin CLI and shared data helpers over
4:class:`lilbee.model_manager.ModelManager`. The ``*_data`` functions
5return Pydantic models so MCP tools and CLI commands share a single,
6typed implementation.
8Heavy imports (:mod:`lilbee.catalog`, :mod:`lilbee.model_manager`,
9:mod:`lilbee.registry`, :mod:`lilbee.cli.tui`) are deferred to function
10bodies so importing this module at CLI startup stays cheap.
11"""
13from __future__ import annotations
15from enum import StrEnum
16from pathlib import Path
17from typing import TYPE_CHECKING, Any
19import typer
20from pydantic import BaseModel, Field
21from rich.console import Console
22from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
23from rich.table import Table
25from lilbee.cli import theme
26from lilbee.cli.app import (
27 apply_overrides,
28 console,
29 data_dir_option,
30 global_option,
31)
32from lilbee.cli.helpers import json_output
33from lilbee.config import cfg
35if TYPE_CHECKING:
36 from collections.abc import Callable
38 from lilbee.catalog import CatalogModel, DownloadProgress
39 from lilbee.model_manager import ModelSource, RemoteModel
40 from lilbee.registry import ModelManifest
43_BYTES_PER_GB = 1024**3 # Model sizes are reported to users in GiB.
44_BACKEND_LIST_TIMEOUT_S = 2.0 # Keep `model list` snappy when backend is down.
47def _bytes_to_gb(n: int) -> float:
48 """Convert bytes to GiB rounded to 2 decimals for user display."""
49 return round(n / _BYTES_PER_GB, 2)
52class ModelCommand(StrEnum):
53 """Command field values for model sub-app JSON output."""
55 LIST = "model list"
56 SHOW = "model show"
57 PULL = "model pull"
58 RM = "model rm"
61class PullStatus(StrEnum):
62 OK = "ok"
63 ALREADY_INSTALLED = "already_installed"
66class PullEvent(StrEnum):
67 PROGRESS = "progress"
68 DONE = "done"
71class ModelEntry(BaseModel):
72 """One row of `lilbee model list` output."""
74 name: str
75 source: str
76 task: str | None = None
77 size_gb: float | None = None
78 display_name: str = ""
80 @classmethod
81 def from_native(cls, ref: str, manifest: ModelManifest | None) -> ModelEntry:
82 from lilbee.catalog import clean_display_name
83 from lilbee.model_manager import ModelSource
85 return cls(
86 name=ref,
87 source=ModelSource.NATIVE.value,
88 task=manifest.task if manifest else None,
89 size_gb=_bytes_to_gb(manifest.size_bytes) if manifest else None,
90 display_name=clean_display_name(manifest.hf_repo) if manifest else "",
91 )
93 @classmethod
94 def from_backend(cls, ref: str, remote: RemoteModel | None) -> ModelEntry:
95 from lilbee.model_manager import ModelSource
97 return cls(
98 name=ref,
99 source=ModelSource.REMOTE.value,
100 task=remote.task if remote else None,
101 size_gb=None,
102 display_name=remote.parameter_size if remote else "",
103 )
106class ListModelsResult(BaseModel):
107 command: str = ModelCommand.LIST
108 models: list[ModelEntry]
109 total: int
111 def __rich__(self) -> Table:
112 table = Table(title="Installed models")
113 table.add_column("Name", style=theme.ACCENT)
114 table.add_column("Source", style=theme.MUTED)
115 table.add_column("Task")
116 table.add_column("Size", justify="right")
117 for entry in self.models:
118 size = f"{entry.size_gb:.2f} GB" if entry.size_gb is not None else ""
119 table.add_row(entry.name, entry.source, entry.task or "", size)
120 return table
123class CatalogEntryData(BaseModel):
124 ref: str
125 display_name: str
126 hf_repo: str
127 gguf_filename: str
128 size_gb: float
129 min_ram_gb: float
130 description: str
131 task: str
132 featured: bool
133 recommended: bool
135 @classmethod
136 def from_catalog_model(cls, entry: CatalogModel) -> CatalogEntryData:
137 return cls(
138 ref=entry.ref,
139 display_name=entry.display_name,
140 hf_repo=entry.hf_repo,
141 gguf_filename=entry.gguf_filename,
142 size_gb=entry.size_gb,
143 min_ram_gb=entry.min_ram_gb,
144 description=entry.description,
145 task=entry.task,
146 featured=entry.featured,
147 recommended=entry.recommended,
148 )
151class ManifestData(BaseModel):
152 ref: str
153 display_name: str
154 task: str
155 size_gb: float
156 size_bytes: int
157 hf_repo: str
158 gguf_filename: str
159 downloaded_at: str
161 @classmethod
162 def from_manifest(cls, manifest: ModelManifest) -> ManifestData:
163 from lilbee.catalog import clean_display_name
165 return cls(
166 ref=manifest.ref,
167 display_name=clean_display_name(manifest.hf_repo),
168 task=manifest.task,
169 size_gb=_bytes_to_gb(manifest.size_bytes),
170 size_bytes=manifest.size_bytes,
171 hf_repo=manifest.hf_repo,
172 gguf_filename=manifest.gguf_filename,
173 downloaded_at=manifest.downloaded_at,
174 )
177class ShowModelResult(BaseModel):
178 command: str = ModelCommand.SHOW
179 model: str
180 catalog: CatalogEntryData | None = None
181 installed: bool = False
182 source: str | None = None
183 path: str | None = None
184 manifest: ManifestData | None = None
186 def __rich__(self) -> str:
187 lines = [f"[{theme.ACCENT}]{self.model}[/{theme.ACCENT}]"]
188 if self.catalog is not None:
189 lines.extend(
190 [
191 f" display_name: {self.catalog.display_name}",
192 f" task: {self.catalog.task}",
193 f" size_gb: {self.catalog.size_gb}",
194 f" min_ram_gb: {self.catalog.min_ram_gb}",
195 f" hf_repo: {self.catalog.hf_repo}",
196 f" description: {self.catalog.description}",
197 ]
198 )
199 lines.append(f" installed: {self.installed}")
200 if self.source:
201 lines.append(f" source: {self.source}")
202 if self.path:
203 lines.append(f" path: {self.path}")
204 if self.manifest is not None:
205 lines.append(f" downloaded: {self.manifest.downloaded_at}")
206 return "\n".join(lines)
209class PullResult(BaseModel):
210 command: str = ModelCommand.PULL
211 model: str
212 source: str
213 status: str
214 path: str | None = None
217class PullProgressEvent(BaseModel):
218 command: str = ModelCommand.PULL
219 event: str = PullEvent.PROGRESS
220 model: str
221 percent: float
222 detail: str
223 cache_hit: bool
226class RemoveResult(BaseModel):
227 command: str = ModelCommand.RM
228 model: str
229 deleted: bool
230 freed_gb: float = Field(default=0.0)
233def _native_manifest_index() -> dict[str, ModelManifest]:
234 """Map ref string ('hf_repo/filename') to manifest for every installed native model."""
235 from lilbee.registry import ModelRegistry
237 registry = ModelRegistry(cfg.models_dir)
238 return {m.ref: m for m in registry.list_installed()}
241def _resolve_native_path(ref: str) -> str | None:
242 """Return the on-disk path of an installed native model, if resolvable.
244 Swallows ``KeyError`` (manifest present but blob missing) and
245 ``ValueError`` (malformed ref) so callers can treat the path as
246 optional metadata.
247 """
248 from lilbee.registry import ModelRegistry
250 try:
251 return str(ModelRegistry(cfg.models_dir).resolve(ref))
252 except (KeyError, ValueError):
253 return None
256def _collect_native_entries() -> list[ModelEntry]:
257 from lilbee.model_manager import ModelSource, get_model_manager
259 manifests = _native_manifest_index()
260 refs = get_model_manager().list_installed(source=ModelSource.NATIVE)
261 return [ModelEntry.from_native(ref, manifests.get(ref)) for ref in refs]
264def _collect_backend_entries() -> list[ModelEntry]:
265 from lilbee.model_manager import classify_remote_models
267 remote_list = classify_remote_models(cfg.remote_base_url, timeout=_BACKEND_LIST_TIMEOUT_S)
268 remote_by_name = {rm.name: rm for rm in remote_list}
269 return [ModelEntry.from_backend(ref, remote_by_name[ref]) for ref in sorted(remote_by_name)]
272def list_models_data(
273 source: ModelSource | None = None,
274 task: str | None = None,
275) -> ListModelsResult:
276 """Build the list of installed models with source and task metadata.
278 Discovers remote models via a single HTTP call with a short timeout
279 so the command stays responsive when the backend is down.
280 """
281 from lilbee.model_manager import ModelSource
283 entries: list[ModelEntry] = []
284 if source is None or source is ModelSource.NATIVE:
285 entries.extend(_collect_native_entries())
286 if source is None or source is ModelSource.REMOTE:
287 entries.extend(_collect_backend_entries())
288 if task:
289 entries = [e for e in entries if e.task == task]
290 return ListModelsResult(models=entries, total=len(entries))
293def show_model_data(ref: str) -> ShowModelResult:
294 """Return catalog and install metadata for *ref*.
296 Raises :class:`~lilbee.model_manager.ModelNotFoundError` if the ref
297 is unknown to both the catalog and the installed set.
298 """
299 from lilbee.catalog import find_catalog_entry
300 from lilbee.model_manager import ModelNotFoundError, get_model_manager
302 entry = find_catalog_entry(ref)
303 source = get_model_manager().get_source(ref)
304 if entry is None and source is None:
305 raise ModelNotFoundError(f"model not found: {ref}")
306 manifest = _native_manifest_index().get(ref)
307 return ShowModelResult(
308 model=ref,
309 catalog=CatalogEntryData.from_catalog_model(entry) if entry else None,
310 installed=source is not None,
311 source=source.value if source else None,
312 manifest=ManifestData.from_manifest(manifest) if manifest else None,
313 path=_resolve_native_path(ref) if manifest is not None else None,
314 )
317def _backend_event_to_progress(
318 on_update: Callable[[DownloadProgress], None],
319 event: dict[str, Any],
320) -> None:
321 """Adapt an Ollama-style dict event into a DownloadProgress call."""
322 from lilbee.catalog import DownloadProgress
324 total = event.get("total", 0) or 0
325 completed = event.get("completed", 0) or 0
326 detail = event.get("status", "") or ""
327 pct = int(completed * 100 / total) if total > 0 else 0
328 on_update(DownloadProgress(percent=pct, detail=detail, is_cache_hit=False))
331def _build_pull_callbacks(
332 on_update: Callable[[DownloadProgress], None] | None,
333) -> tuple[Callable[[dict[str, Any]], None] | None, Callable[[int, int], None] | None]:
334 """Build the (dict_cb, bytes_cb) pair for ModelManager.pull from on_update."""
335 import functools
337 from lilbee.catalog import make_download_callback
339 if on_update is None:
340 return None, None
341 dict_cb = functools.partial(_backend_event_to_progress, on_update)
342 bytes_cb = make_download_callback(on_update)
343 return dict_cb, bytes_cb
346def pull_model_data(
347 ref: str,
348 source: ModelSource,
349 *,
350 on_update: Callable[[DownloadProgress], None] | None = None,
351) -> PullResult:
352 """Pull *ref* from *source* and return a typed result.
354 Progress updates are throttled by
355 :func:`~lilbee.catalog.make_download_callback`, so callers see at
356 most roughly 10 Hz of progress events.
357 """
358 from lilbee.model_manager import get_model_manager
360 manager = get_model_manager()
362 if manager.is_installed(ref, source):
363 return PullResult(model=ref, source=source.value, status=PullStatus.ALREADY_INSTALLED)
365 dict_cb, bytes_cb = _build_pull_callbacks(on_update)
366 path = manager.pull(ref, source, on_progress=dict_cb, on_bytes=bytes_cb)
367 return PullResult(
368 model=ref,
369 source=source.value,
370 status=PullStatus.OK,
371 path=str(path) if path is not None else None,
372 )
375def remove_model_data(
376 ref: str,
377 source: ModelSource | None = None,
378) -> RemoveResult:
379 """Remove *ref* and return a typed result with freed size."""
380 from lilbee.model_manager import get_model_manager
382 manager = get_model_manager()
383 manifests = _native_manifest_index()
384 size_bytes = manifests[ref].size_bytes if ref in manifests else 0
385 removed = manager.remove(ref, source=source)
386 return RemoveResult(
387 model=ref,
388 deleted=removed,
389 freed_gb=_bytes_to_gb(size_bytes),
390 )
393model_app = typer.Typer(
394 name="model",
395 help="Manage installed and available models (pull / list / show / rm / browse).",
396 no_args_is_help=True,
397)
399_source_option = typer.Option(
400 None,
401 "--source",
402 "-s",
403 help="Filter by source: 'native' or 'remote' (default: all).",
404)
405_task_option = typer.Option(
406 None,
407 "--task",
408 "-t",
409 help="Filter by task: 'chat', 'embedding', 'vision', or 'rerank'.",
410)
411_yes_option = typer.Option(
412 False,
413 "--yes",
414 "-y",
415 help="Skip confirmation prompt.",
416)
419def _parse_source_or_bad_param(value: str | None) -> ModelSource | None:
420 """Parse a CLI --source value, raising typer.BadParameter on bad input."""
421 from lilbee.model_manager import ModelSource
423 try:
424 return ModelSource.parse(value)
425 except ValueError as exc:
426 if cfg.json_mode:
427 json_output({"error": str(exc)})
428 raise SystemExit(1) from None
429 raise typer.BadParameter(str(exc)) from exc
432@model_app.command("list")
433def list_cmd(
434 source: str | None = _source_option,
435 task: str | None = _task_option,
436 data_dir: Path | None = data_dir_option,
437 use_global: bool = global_option,
438) -> None:
439 """List installed models across all sources."""
440 apply_overrides(data_dir=data_dir, use_global=use_global)
441 data = list_models_data(source=_parse_source_or_bad_param(source), task=task)
442 if cfg.json_mode:
443 json_output(data.model_dump())
444 return
445 if not data.models:
446 console.print("No models installed.")
447 return
448 console.print(data)
451@model_app.command("show")
452def show_cmd(
453 ref: str = typer.Argument(..., help="Model ref (e.g. 'Qwen/Qwen3-0.6B-GGUF')."),
454 data_dir: Path | None = data_dir_option,
455 use_global: bool = global_option,
456) -> None:
457 """Show catalog and installed metadata for a model."""
458 from lilbee.model_manager import ModelNotFoundError
460 apply_overrides(data_dir=data_dir, use_global=use_global)
461 try:
462 data = show_model_data(ref)
463 except ModelNotFoundError as exc:
464 if cfg.json_mode:
465 json_output({"error": str(exc)})
466 else:
467 console.print(f"[{theme.ERROR}]{exc}[/{theme.ERROR}]")
468 raise typer.Exit(1) from None
469 if cfg.json_mode:
470 json_output(data.model_dump())
471 return
472 console.print(data)
475def _run_pull(
476 ref: str,
477 src: ModelSource,
478 on_update: Callable[[DownloadProgress], None],
479) -> PullResult:
480 """Invoke ``pull_model_data`` and translate known errors to typer.Exit."""
481 try:
482 return pull_model_data(ref, src, on_update=on_update)
483 except (RuntimeError, PermissionError) as exc:
484 if cfg.json_mode:
485 json_output({"error": str(exc)})
486 else:
487 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")
488 raise typer.Exit(1) from None
491def _pull_json_stream(ref: str, src: ModelSource) -> None:
492 """Emit newline-delimited JSON progress events, then the final result."""
494 def on_update(p: DownloadProgress) -> None:
495 event = PullProgressEvent(
496 model=ref, percent=p.percent, detail=p.detail, cache_hit=p.is_cache_hit
497 )
498 json_output(event.model_dump())
500 final = _run_pull(ref, src, on_update)
501 json_output({**final.model_dump(), "event": PullEvent.DONE.value})
504def _pull_rich_progress(ref: str, src: ModelSource) -> None:
505 """Drive a Rich progress bar during a native HuggingFace download."""
506 err_console = Console(stderr=True)
507 with Progress(
508 TextColumn("[progress.description]{task.description}"),
509 BarColumn(),
510 TextColumn("{task.percentage:>3.0f}%"),
511 TextColumn("{task.fields[detail]}"),
512 TimeRemainingColumn(),
513 console=err_console,
514 transient=False,
515 ) as progress:
516 task_id = progress.add_task(f"Downloading {ref}", total=100, detail="")
518 def on_update(p: DownloadProgress) -> None:
519 progress.update(task_id, completed=p.percent, detail=p.detail)
521 final = _run_pull(ref, src, on_update)
523 if final.status == PullStatus.ALREADY_INSTALLED:
524 console.print(f"{ref} is already installed.")
525 else:
526 console.print(f"Pulled [{theme.ACCENT}]{ref}[/{theme.ACCENT}].")
529@model_app.command("pull")
530def pull_cmd(
531 ref: str = typer.Argument(..., help="Model ref to download (e.g. 'Qwen/Qwen3-0.6B-GGUF')."),
532 source: str = typer.Option(
533 "native",
534 "--source",
535 "-s",
536 help="Pull from 'native' (HuggingFace GGUF) or 'remote' (SDK-managed).",
537 ),
538 data_dir: Path | None = data_dir_option,
539 use_global: bool = global_option,
540) -> None:
541 """Download a model."""
542 from lilbee.model_manager import ModelSource
544 apply_overrides(data_dir=data_dir, use_global=use_global)
545 src = _parse_source_or_bad_param(source) or ModelSource.NATIVE
546 if cfg.json_mode:
547 _pull_json_stream(ref, src)
548 else:
549 _pull_rich_progress(ref, src)
552def _confirm_remove_or_exit(ref: str, yes: bool) -> None:
553 if yes or cfg.json_mode:
554 return
555 if not typer.confirm(f"Remove {ref}?", default=False):
556 console.print("Aborted.")
557 raise typer.Exit(0)
560@model_app.command("rm")
561def rm_cmd(
562 ref: str = typer.Argument(..., help="Model ref to remove."),
563 source: str | None = _source_option,
564 yes: bool = _yes_option,
565 data_dir: Path | None = data_dir_option,
566 use_global: bool = global_option,
567) -> None:
568 """Remove an installed model."""
569 apply_overrides(data_dir=data_dir, use_global=use_global)
570 src = _parse_source_or_bad_param(source)
571 _confirm_remove_or_exit(ref, yes)
572 data = remove_model_data(ref, source=src)
573 if cfg.json_mode:
574 json_output(data.model_dump())
575 if not data.deleted:
576 raise typer.Exit(1)
577 return
578 if not data.deleted:
579 console.print(f"[{theme.WARNING}]Not found: {ref}[/{theme.WARNING}]")
580 raise typer.Exit(1)
581 suffix = f" ({data.freed_gb:.2f} GB freed)" if data.freed_gb else ""
582 console.print(f"Removed [{theme.ACCENT}]{ref}[/{theme.ACCENT}]{suffix}.")
585def _is_interactive_terminal() -> bool:
586 """Return True when both stdin and stdout are connected to a TTY.
588 Extracted as a module-level helper so tests can patch it deterministically;
589 CliRunner replaces ``sys.stdin`` during invoke which makes direct
590 monkey-patching of ``sys.stdin.isatty`` unreliable.
591 """
592 import sys
594 return sys.stdin.isatty() and sys.stdout.isatty()
597@model_app.command("browse")
598def browse_cmd(
599 data_dir: Path | None = data_dir_option,
600 use_global: bool = global_option,
601) -> None:
602 """Open the Textual TUI directly on the model catalog screen.
604 Exit codes follow the project convention: 2 for invalid flag
605 combinations (``--json`` with an interactive-only command), 1 for
606 runtime environment failures (no TTY).
607 """
608 apply_overrides(data_dir=data_dir, use_global=use_global)
609 if cfg.json_mode:
610 json_output({"error": "model browse is interactive, not available in --json mode"})
611 raise typer.Exit(2)
612 if not _is_interactive_terminal():
613 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] model browse requires a terminal.")
614 raise typer.Exit(1)
616 # Browsing the catalog does not depend on documents, so skip auto-sync.
617 from lilbee.cli.tui import run_tui
619 run_tui(auto_sync=False, initial_view="Catalog")