Coverage for src/lilbee/app/models.py: 100%
199 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""Surface-agnostic model lifecycle use-cases (list / show / pull / remove)."""
3from __future__ import annotations
5from enum import StrEnum
6from typing import TYPE_CHECKING
8from pydantic import BaseModel, Field
10from lilbee.app.services import get_services
11from lilbee.catalog.types import ModelCompat, ModelTask
12from lilbee.core.config import cfg
13from lilbee.modelhub.registry import ModelRegistry
15if TYPE_CHECKING:
16 from collections.abc import Callable
18 from lilbee.catalog import CatalogModel, DownloadProgress
19 from lilbee.catalog.types import ModelSource
20 from lilbee.modelhub.model_manager import RemoteModel
21 from lilbee.modelhub.registry import ModelManifest
22 from lilbee.runtime.cancellation import CancelSignal
25_BYTES_PER_GB = 1024**3 # Model sizes are reported to users in GiB.
26_BACKEND_LIST_TIMEOUT_S = 2.0 # Keep `model list` snappy when backend is down.
29def _bytes_to_gb(n: int) -> float:
30 """Convert bytes to GiB rounded to 2 decimals for user display."""
31 return round(n / _BYTES_PER_GB, 2)
34class ModelCommand(StrEnum):
35 """Command field values for model sub-app JSON output."""
37 LIST = "model list"
38 SHOW = "model show"
39 PULL = "model pull"
40 RM = "model rm"
43class PullStatus(StrEnum):
44 OK = "ok"
45 ALREADY_INSTALLED = "already_installed"
48class AdoptStatus(StrEnum):
49 ADOPTED = "adopted"
50 ALREADY_ACTIVE = "already_active"
53class PullEvent(StrEnum):
54 PROGRESS = "progress"
55 DONE = "done"
58class ModelEntry(BaseModel):
59 """One row of `lilbee model list` output."""
61 name: str
62 source: str
63 task: ModelTask | None = None
64 size_gb: float | None = None
65 display_name: str = ""
67 @classmethod
68 def from_native(cls, ref: str, manifest: ModelManifest | None) -> ModelEntry:
69 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
70 from lilbee.catalog import clean_display_name
71 from lilbee.catalog.types import ModelSource
73 return cls(
74 name=ref,
75 source=ModelSource.NATIVE.value,
76 task=manifest.task if manifest else None,
77 size_gb=_bytes_to_gb(manifest.disk_size_bytes) if manifest else None,
78 display_name=clean_display_name(manifest.hf_repo) if manifest else "",
79 )
81 @classmethod
82 def from_backend(cls, ref: str, remote: RemoteModel | None, source: ModelSource) -> ModelEntry:
83 from lilbee.providers.local_servers import canonical_local_ref
85 return cls(
86 name=canonical_local_ref(ref, source.value),
87 source=source.value,
88 task=remote.task if remote else None,
89 size_gb=None,
90 display_name=remote.parameter_size if remote else "",
91 )
94class ListModelsResult(BaseModel):
95 command: str = ModelCommand.LIST
96 models: list[ModelEntry]
97 total: int
100class CatalogEntryData(BaseModel):
101 ref: str
102 display_name: str
103 hf_repo: str
104 gguf_filename: str
105 size_gb: float
106 min_ram_gb: float
107 description: str
108 task: ModelTask
109 featured: bool
110 architecture: str = ""
111 compat: ModelCompat = ModelCompat.UNKNOWN
113 @classmethod
114 def from_catalog_model(cls, entry: CatalogModel) -> CatalogEntryData:
115 return cls(
116 ref=entry.ref,
117 display_name=entry.display_name,
118 hf_repo=entry.hf_repo,
119 gguf_filename=entry.gguf_filename,
120 size_gb=entry.size_gb,
121 min_ram_gb=entry.min_ram_gb,
122 description=entry.description,
123 task=entry.task,
124 featured=entry.featured,
125 architecture=entry.architecture,
126 compat=entry.compat,
127 )
130class ManifestData(BaseModel):
131 ref: str
132 display_name: str
133 task: ModelTask
134 size_gb: float
135 size_bytes: int
136 hf_repo: str
137 gguf_filename: str
138 downloaded_at: str
140 @classmethod
141 def from_manifest(cls, manifest: ModelManifest) -> ManifestData:
142 from lilbee.catalog import clean_display_name
144 return cls(
145 ref=manifest.ref,
146 display_name=clean_display_name(manifest.hf_repo),
147 task=manifest.task,
148 size_gb=_bytes_to_gb(manifest.disk_size_bytes),
149 size_bytes=manifest.disk_size_bytes,
150 hf_repo=manifest.hf_repo,
151 gguf_filename=manifest.gguf_filename,
152 downloaded_at=manifest.downloaded_at,
153 )
156class ShowModelResult(BaseModel):
157 command: str = ModelCommand.SHOW
158 model: str
159 catalog: CatalogEntryData | None = None
160 installed: bool = False
161 source: str | None = None
162 path: str | None = None
163 manifest: ManifestData | None = None
166class PullResult(BaseModel):
167 command: str = ModelCommand.PULL
168 model: str
169 source: str
170 status: PullStatus
171 path: str | None = None
174class PullProgressEvent(BaseModel):
175 command: str = ModelCommand.PULL
176 event: str = PullEvent.PROGRESS
177 model: str
178 percent: float
179 detail: str
180 cache_hit: bool
183class RemoveResult(BaseModel):
184 command: str = ModelCommand.RM
185 model: str
186 deleted: bool
187 freed_gb: float = Field(default=0.0)
190class AdoptResult(BaseModel):
191 """Outcome of adopting a downloaded index's embedder."""
193 model: str
194 status: AdoptStatus
195 reindex_required: bool = False
198def adopt_embedder(ref: str) -> AdoptResult:
199 """Switch lilbee to embedder *ref*, downloading it first if missing.
201 Makes a downloaded index searchable under its own embedder without a
202 rebuild: the persisted vectors already match *ref*, so the switch routes
203 through the settings boundary and ``reindex_required`` stays false.
204 """
205 from lilbee.app.settings import apply_settings_update
206 from lilbee.catalog.types import ModelSource
208 manager = get_services().model_manager
209 installed = manager.is_installed(ref, ModelSource.NATIVE)
210 already_active = cfg.embedding_model == ref and installed
211 if not installed:
212 pull_model_data(ref, ModelSource.NATIVE)
213 result = apply_settings_update({"embedding_model": ref})
214 return AdoptResult(
215 model=ref,
216 status=AdoptStatus.ALREADY_ACTIVE if already_active else AdoptStatus.ADOPTED,
217 reindex_required=result.reindex_required,
218 )
221def installed_chat_model_refs() -> list[str]:
222 """Return sorted refs for every chat-task model in the registry."""
223 registry = get_services().registry
224 return sorted(m.ref for m in registry.list_installed() if m.task == ModelTask.CHAT)
227def _native_manifest_index() -> dict[str, ModelManifest]:
228 """Map ref string ('hf_repo/filename') to manifest for every installed native model."""
229 registry = ModelRegistry(cfg.models_dir)
230 return {m.ref: m for m in registry.list_installed()}
233def _resolve_native_path(ref: str) -> str | None:
234 """Return the on-disk path of an installed native model, if resolvable.
236 Swallows ``KeyError`` (manifest present but blob missing) and
237 ``ValueError`` (malformed ref) so callers can treat the path as
238 optional metadata.
239 """
240 try:
241 return str(ModelRegistry(cfg.models_dir).resolve(ref))
242 except (KeyError, ValueError):
243 return None
246def _collect_native_entries() -> list[ModelEntry]:
247 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
248 from lilbee.catalog.types import ModelSource
250 manifests = _native_manifest_index()
251 refs = get_services().model_manager.list_installed(source=ModelSource.NATIVE)
252 return [ModelEntry.from_native(ref, manifests.get(ref)) for ref in refs]
255def _collect_backend_entries() -> list[ModelEntry]:
256 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
257 from lilbee.catalog.types import ModelSource
258 from lilbee.modelhub.model_manager import classify_all_remote_models
259 from lilbee.providers.local_servers import local_server_for_label
261 def _source(remote: RemoteModel) -> ModelSource:
262 spec = local_server_for_label(remote.provider)
263 return ModelSource(spec.key) if spec is not None else ModelSource.REMOTE
265 remote_by_name = {
266 rm.name: rm for rm in classify_all_remote_models(timeout=_BACKEND_LIST_TIMEOUT_S)
267 }
268 return [
269 ModelEntry.from_backend(name, rm, _source(rm))
270 for name, rm in sorted(remote_by_name.items())
271 ]
274def list_models_data(
275 source: ModelSource | None = None,
276 task: ModelTask | None = None,
277) -> ListModelsResult:
278 """Build the list of installed models with source and task metadata.
280 Discovers remote models via a single HTTP call with a short timeout
281 so the command stays responsive when the backend is down.
282 """
283 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
284 from lilbee.catalog.types import ModelSource
286 entries: list[ModelEntry] = []
287 if source is None or source is ModelSource.NATIVE:
288 entries.extend(_collect_native_entries())
289 if source is not ModelSource.NATIVE:
290 backend = _collect_backend_entries()
291 # A specific local-server source (ollama/lm_studio/frontier) narrows the
292 # backend list; REMOTE and None keep every backend entry.
293 if source is not None and source is not ModelSource.REMOTE:
294 backend = [e for e in backend if e.source == source.value]
295 entries.extend(backend)
296 if task:
297 entries = [e for e in entries if e.task == task]
298 return ListModelsResult(models=entries, total=len(entries))
301def show_model_data(ref: str) -> ShowModelResult:
302 """Return catalog and install metadata for *ref*.
304 Raises :class:`~lilbee.modelhub.model_manager.ModelNotFoundError` if the ref
305 is unknown to both the catalog and the installed set.
306 """
307 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
308 from lilbee.catalog import find_pick
309 from lilbee.modelhub.model_manager import ModelNotFoundError
311 entry = find_pick(ref)
312 source = get_services().model_manager.get_source(ref)
313 if entry is None and source is None:
314 raise ModelNotFoundError(f"model not found: {ref}")
315 manifest = _native_manifest_index().get(ref)
316 return ShowModelResult(
317 model=ref,
318 catalog=CatalogEntryData.from_catalog_model(entry) if entry else None,
319 installed=source is not None,
320 source=source.value if source else None,
321 manifest=ManifestData.from_manifest(manifest) if manifest else None,
322 path=_resolve_native_path(ref) if manifest is not None else None,
323 )
326def _vision_projector_missing(ref: str) -> bool:
327 """True when *ref*'s mmproj projector does not resolve on disk."""
328 from lilbee.providers.base import ProviderError
329 from lilbee.providers.engine_params import resolve_model_path
330 from lilbee.providers.gguf_meta import find_mmproj_for_model
332 try:
333 find_mmproj_for_model(resolve_model_path(ref))
334 except (ProviderError, OSError, ValueError, KeyError):
335 return True
336 return False
339def _ensure_vision_projector(ref: str) -> None:
340 """Fetch a vision model's mmproj projector when a cached install lacks it.
342 No-op for non-vision refs and when the projector already resolves on disk,
343 so pulling a complete install never touches the network.
344 """
345 from lilbee.catalog import download_mmproj, resolve_pull_target
347 entry = resolve_pull_target(ref)
348 if entry is not None and entry.task is ModelTask.VISION and _vision_projector_missing(ref):
349 download_mmproj(entry)
352def pull_model_data(
353 ref: str,
354 source: ModelSource,
355 *,
356 on_update: Callable[[DownloadProgress], None] | None = None,
357 allow_unsupported: bool = False,
358 cancel: CancelSignal | None = None,
359) -> PullResult:
360 """Pull *ref* from *source* and return a typed result.
362 Only native models are downloadable; a non-native *source* is refused by
363 :meth:`ModelManager.pull`. Progress updates are throttled by
364 :func:`~lilbee.catalog.make_download_callback`, so callers see at most
365 roughly 10 Hz of progress events. A *cancel* signal makes the download
366 cancellable mid-transfer; see :func:`~lilbee.catalog.download_model`.
367 """
368 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
369 from lilbee.catalog import make_download_callback
371 manager = get_services().model_manager
373 if manager.is_installed(ref, source):
374 # A cached vision install may carry the main GGUF but not its mmproj
375 # projector; without it llama-server can't serve OCR, so ensure it before
376 # reporting already-installed.
377 _ensure_vision_projector(ref)
378 return PullResult(model=ref, source=source.value, status=PullStatus.ALREADY_INSTALLED)
380 bytes_cb = make_download_callback(on_update) if on_update is not None else None
381 path = manager.pull(
382 ref,
383 source,
384 on_bytes=bytes_cb,
385 allow_unsupported=allow_unsupported,
386 cancel=cancel,
387 )
388 return PullResult(
389 model=ref,
390 source=source.value,
391 status=PullStatus.OK,
392 path=str(path) if path is not None else None,
393 )
396def _legacy_disk_size(ref: str, *, fallback: int) -> int:
397 """Sum a split GGUF's shard sizes on disk; *fallback* on any failure/single file."""
398 import contextlib
400 with contextlib.suppress(Exception):
401 shards = get_services().registry.shard_paths(ref)
402 if len(shards) > 1:
403 return sum(path.stat().st_size for path in shards)
404 return fallback
407def remove_model_data(
408 ref: str,
409 source: ModelSource | None = None,
410) -> RemoveResult:
411 """Remove *ref* and return a typed result with freed size."""
412 manager = get_services().model_manager
413 manifests = _native_manifest_index()
414 # disk_size_bytes is the full multi-shard total; size_bytes alone would report
415 # only the first shard for a split GGUF.
416 manifest = manifests.get(ref)
417 size_bytes = manifest.disk_size_bytes if manifest is not None else 0
418 if manifest is not None and manifest.total_size_bytes is None:
419 # Legacy manifest without shard accounting: removal still frees every
420 # shard, so recover the true on-disk total from the shards for the report.
421 size_bytes = _legacy_disk_size(ref, fallback=size_bytes)
422 removed = manager.remove(ref, source=source)
423 return RemoveResult(
424 model=ref,
425 deleted=removed,
426 freed_gb=_bytes_to_gb(size_bytes),
427 )