Coverage for src/lilbee/app/models.py: 100%
200 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +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
112 safety_stripped: bool = False
114 @classmethod
115 def from_catalog_model(cls, entry: CatalogModel) -> CatalogEntryData:
116 return cls(
117 ref=entry.ref,
118 display_name=entry.display_name,
119 hf_repo=entry.hf_repo,
120 gguf_filename=entry.gguf_filename,
121 size_gb=entry.size_gb,
122 min_ram_gb=entry.min_ram_gb,
123 description=entry.description,
124 task=entry.task,
125 featured=entry.featured,
126 architecture=entry.architecture,
127 compat=entry.compat,
128 safety_stripped=entry.safety_stripped,
129 )
132class ManifestData(BaseModel):
133 ref: str
134 display_name: str
135 task: ModelTask
136 size_gb: float
137 size_bytes: int
138 hf_repo: str
139 gguf_filename: str
140 downloaded_at: str
142 @classmethod
143 def from_manifest(cls, manifest: ModelManifest) -> ManifestData:
144 from lilbee.catalog import clean_display_name
146 return cls(
147 ref=manifest.ref,
148 display_name=clean_display_name(manifest.hf_repo),
149 task=manifest.task,
150 size_gb=_bytes_to_gb(manifest.disk_size_bytes),
151 size_bytes=manifest.disk_size_bytes,
152 hf_repo=manifest.hf_repo,
153 gguf_filename=manifest.gguf_filename,
154 downloaded_at=manifest.downloaded_at,
155 )
158class ShowModelResult(BaseModel):
159 command: str = ModelCommand.SHOW
160 model: str
161 catalog: CatalogEntryData | None = None
162 installed: bool = False
163 source: str | None = None
164 path: str | None = None
165 manifest: ManifestData | None = None
168class PullResult(BaseModel):
169 command: str = ModelCommand.PULL
170 model: str
171 source: str
172 status: PullStatus
173 path: str | None = None
176class PullProgressEvent(BaseModel):
177 command: str = ModelCommand.PULL
178 event: str = PullEvent.PROGRESS
179 model: str
180 percent: float
181 detail: str
182 cache_hit: bool
185class RemoveResult(BaseModel):
186 command: str = ModelCommand.RM
187 model: str
188 deleted: bool
189 freed_gb: float = Field(default=0.0)
192class AdoptResult(BaseModel):
193 """Outcome of adopting a downloaded index's embedder."""
195 model: str
196 status: AdoptStatus
197 reindex_required: bool = False
200def adopt_embedder(ref: str) -> AdoptResult:
201 """Switch lilbee to embedder *ref*, downloading it first if missing.
203 Makes a downloaded index searchable under its own embedder without a
204 rebuild: the persisted vectors already match *ref*, so the switch routes
205 through the settings boundary and ``reindex_required`` stays false.
206 """
207 from lilbee.app.settings import apply_settings_update
208 from lilbee.catalog.types import ModelSource
210 manager = get_services().model_manager
211 installed = manager.is_installed(ref, ModelSource.NATIVE)
212 already_active = cfg.embedding_model == ref and installed
213 if not installed:
214 pull_model_data(ref, ModelSource.NATIVE)
215 result = apply_settings_update({"embedding_model": ref})
216 return AdoptResult(
217 model=ref,
218 status=AdoptStatus.ALREADY_ACTIVE if already_active else AdoptStatus.ADOPTED,
219 reindex_required=result.reindex_required,
220 )
223def installed_chat_model_refs() -> list[str]:
224 """Return sorted refs for every chat-task model in the registry."""
225 registry = get_services().registry
226 return sorted(m.ref for m in registry.list_installed() if m.task == ModelTask.CHAT)
229def _native_manifest_index() -> dict[str, ModelManifest]:
230 """Map ref string ('hf_repo/filename') to manifest for every installed native model."""
231 registry = ModelRegistry(cfg.models_dir)
232 return {m.ref: m for m in registry.list_installed()}
235def _resolve_native_path(ref: str) -> str | None:
236 """Return the on-disk path of an installed native model, if resolvable.
238 Swallows ``KeyError`` (manifest present but blob missing) and
239 ``ValueError`` (malformed ref) so callers can treat the path as
240 optional metadata.
241 """
242 try:
243 return str(ModelRegistry(cfg.models_dir).resolve(ref))
244 except (KeyError, ValueError):
245 return None
248def _collect_native_entries() -> list[ModelEntry]:
249 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
250 from lilbee.catalog.types import ModelSource
252 manifests = _native_manifest_index()
253 refs = get_services().model_manager.list_installed(source=ModelSource.NATIVE)
254 return [ModelEntry.from_native(ref, manifests.get(ref)) for ref in refs]
257def _collect_backend_entries() -> list[ModelEntry]:
258 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
259 from lilbee.catalog.types import ModelSource
260 from lilbee.modelhub.model_manager import classify_all_remote_models
261 from lilbee.providers.local_servers import local_server_for_label
263 def _source(remote: RemoteModel) -> ModelSource:
264 spec = local_server_for_label(remote.provider)
265 return ModelSource(spec.key) if spec is not None else ModelSource.REMOTE
267 remote_by_name = {
268 rm.name: rm for rm in classify_all_remote_models(timeout=_BACKEND_LIST_TIMEOUT_S)
269 }
270 return [
271 ModelEntry.from_backend(name, rm, _source(rm))
272 for name, rm in sorted(remote_by_name.items())
273 ]
276def list_models_data(
277 source: ModelSource | None = None,
278 task: ModelTask | None = None,
279) -> ListModelsResult:
280 """Build the list of installed models with source and task metadata.
282 Discovers remote models via a single HTTP call with a short timeout
283 so the command stays responsive when the backend is down.
284 """
285 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
286 from lilbee.catalog.types import ModelSource
288 entries: list[ModelEntry] = []
289 if source is None or source is ModelSource.NATIVE:
290 entries.extend(_collect_native_entries())
291 if source is not ModelSource.NATIVE:
292 backend = _collect_backend_entries()
293 # A specific local-server source (ollama/lm_studio/frontier) narrows the
294 # backend list; REMOTE and None keep every backend entry.
295 if source is not None and source is not ModelSource.REMOTE:
296 backend = [e for e in backend if e.source == source.value]
297 entries.extend(backend)
298 if task:
299 entries = [e for e in entries if e.task == task]
300 return ListModelsResult(models=entries, total=len(entries))
303def show_model_data(ref: str) -> ShowModelResult:
304 """Return catalog and install metadata for *ref*.
306 Raises :class:`~lilbee.modelhub.model_manager.ModelNotFoundError` if the ref
307 is unknown to both the catalog and the installed set.
308 """
309 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
310 from lilbee.catalog import find_pick
311 from lilbee.modelhub.model_manager import ModelNotFoundError
313 entry = find_pick(ref)
314 source = get_services().model_manager.get_source(ref)
315 if entry is None and source is None:
316 raise ModelNotFoundError(f"model not found: {ref}")
317 manifest = _native_manifest_index().get(ref)
318 return ShowModelResult(
319 model=ref,
320 catalog=CatalogEntryData.from_catalog_model(entry) if entry else None,
321 installed=source is not None,
322 source=source.value if source else None,
323 manifest=ManifestData.from_manifest(manifest) if manifest else None,
324 path=_resolve_native_path(ref) if manifest is not None else None,
325 )
328def _vision_projector_missing(ref: str) -> bool:
329 """True when *ref*'s mmproj projector does not resolve on disk."""
330 from lilbee.providers.base import ProviderError
331 from lilbee.providers.engine_params import resolve_model_path
332 from lilbee.providers.gguf_meta import find_mmproj_for_model
334 try:
335 find_mmproj_for_model(resolve_model_path(ref))
336 except (ProviderError, OSError, ValueError, KeyError):
337 return True
338 return False
341def _ensure_vision_projector(ref: str) -> None:
342 """Fetch a vision model's mmproj projector when a cached install lacks it.
344 No-op for non-vision refs and when the projector already resolves on disk,
345 so pulling a complete install never touches the network.
346 """
347 from lilbee.catalog import download_mmproj, resolve_pull_target
349 entry = resolve_pull_target(ref)
350 if entry is not None and entry.task is ModelTask.VISION and _vision_projector_missing(ref):
351 download_mmproj(entry)
354def pull_model_data(
355 ref: str,
356 source: ModelSource,
357 *,
358 on_update: Callable[[DownloadProgress], None] | None = None,
359 allow_unsupported: bool = False,
360 cancel: CancelSignal | None = None,
361) -> PullResult:
362 """Pull *ref* from *source* and return a typed result.
364 Only native models are downloadable; a non-native *source* is refused by
365 :meth:`ModelManager.pull`. Progress updates are throttled by
366 :func:`~lilbee.catalog.make_download_callback`, so callers see at most
367 roughly 10 Hz of progress events. A *cancel* signal makes the download
368 cancellable mid-transfer; see :func:`~lilbee.catalog.download_model`.
369 """
370 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
371 from lilbee.catalog import make_download_callback
373 manager = get_services().model_manager
375 if manager.is_installed(ref, source):
376 # A cached vision install may carry the main GGUF but not its mmproj
377 # projector; without it llama-server can't serve OCR, so ensure it before
378 # reporting already-installed.
379 _ensure_vision_projector(ref)
380 return PullResult(model=ref, source=source.value, status=PullStatus.ALREADY_INSTALLED)
382 bytes_cb = make_download_callback(on_update) if on_update is not None else None
383 path = manager.pull(
384 ref,
385 source,
386 on_bytes=bytes_cb,
387 allow_unsupported=allow_unsupported,
388 cancel=cancel,
389 )
390 return PullResult(
391 model=ref,
392 source=source.value,
393 status=PullStatus.OK,
394 path=str(path) if path is not None else None,
395 )
398def _legacy_disk_size(ref: str, *, fallback: int) -> int:
399 """Sum a split GGUF's shard sizes on disk; *fallback* on any failure/single file."""
400 import contextlib
402 with contextlib.suppress(Exception):
403 shards = get_services().registry.shard_paths(ref)
404 if len(shards) > 1:
405 return sum(path.stat().st_size for path in shards)
406 return fallback
409def remove_model_data(
410 ref: str,
411 source: ModelSource | None = None,
412) -> RemoveResult:
413 """Remove *ref* and return a typed result with freed size."""
414 manager = get_services().model_manager
415 manifests = _native_manifest_index()
416 # disk_size_bytes is the full multi-shard total; size_bytes alone would report
417 # only the first shard for a split GGUF.
418 manifest = manifests.get(ref)
419 size_bytes = manifest.disk_size_bytes if manifest is not None else 0
420 if manifest is not None and manifest.total_size_bytes is None:
421 # Legacy manifest without shard accounting: removal still frees every
422 # shard, so recover the true on-disk total from the shards for the report.
423 size_bytes = _legacy_disk_size(ref, fallback=size_bytes)
424 removed = manager.remove(ref, source=source)
425 return RemoveResult(
426 model=ref,
427 deleted=removed,
428 freed_gb=_bytes_to_gb(size_bytes),
429 )