Coverage for src/lilbee/app/models.py: 100%
204 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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.query import reclassify_by_name
12from lilbee.catalog.types import ModelCompat, ModelSource, ModelTask
13from lilbee.core.config import cfg
14from lilbee.modelhub.registry import ModelRegistry
15from lilbee.providers.model_ref import parse_model_ref
17if TYPE_CHECKING:
18 from collections.abc import Callable
20 from lilbee.catalog import CatalogModel, DownloadProgress
21 from lilbee.modelhub.model_manager import RemoteModel
22 from lilbee.modelhub.registry import ModelManifest
23 from lilbee.runtime.cancellation import CancelSignal
26_BYTES_PER_GB = 1024**3 # Model sizes are reported to users in GiB.
27_BACKEND_LIST_TIMEOUT_S = 2.0 # Keep `model list` snappy when backend is down.
30def _bytes_to_gb(n: int) -> float:
31 """Convert bytes to GiB rounded to 2 decimals for user display."""
32 return round(n / _BYTES_PER_GB, 2)
35class ModelCommand(StrEnum):
36 """Command field values for model sub-app JSON output."""
38 LIST = "model list"
39 SHOW = "model show"
40 PULL = "model pull"
41 RM = "model rm"
44class PullStatus(StrEnum):
45 OK = "ok"
46 ALREADY_INSTALLED = "already_installed"
49class AdoptStatus(StrEnum):
50 ADOPTED = "adopted"
51 ALREADY_ACTIVE = "already_active"
54class PullEvent(StrEnum):
55 PROGRESS = "progress"
56 DONE = "done"
59class ModelEntry(BaseModel):
60 """One installed model, as `lilbee model list`, MCP and the installed route report it."""
62 name: str
63 source: ModelSource
64 task: ModelTask | None = None
65 size_gb: float | None = None
66 display_name: str = ""
68 @classmethod
69 def from_native(cls, ref: str, manifest: ModelManifest | None) -> ModelEntry:
70 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
71 from lilbee.catalog import clean_display_name
73 return cls(
74 name=ref,
75 source=ModelSource.NATIVE,
76 task=ModelTask(reclassify_by_name(ref, 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,
88 task=remote.task if remote else None,
89 size_gb=None,
90 display_name=ref,
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, which normally reports ``reindex_required``
206 false. The result carries the boundary's verdict as is: true when the index
207 still does not match *ref*, and the CLI then shows the rebuild hint.
208 """
209 from lilbee.app.settings import apply_settings_update
210 from lilbee.catalog.types import ModelSource
212 manager = get_services().model_manager
213 installed = manager.is_installed(ref, ModelSource.NATIVE)
214 already_active = cfg.embedding_model == ref and installed
215 if not installed:
216 pull_model_data(ref, ModelSource.NATIVE)
217 result = apply_settings_update({"embedding_model": ref})
218 return AdoptResult(
219 model=ref,
220 status=AdoptStatus.ALREADY_ACTIVE if already_active else AdoptStatus.ADOPTED,
221 reindex_required=result.reindex_required,
222 )
225def installed_chat_model_refs() -> list[str]:
226 """Return sorted refs for every chat-task model in the registry."""
227 registry = get_services().registry
228 return sorted(m.ref for m in registry.list_installed() if m.task == ModelTask.CHAT)
231def with_configured_chat(refs: list[str], configured: str) -> list[str]:
232 """Return *refs* with *configured* first when this server serves it but *refs* omit it.
234 A remote-configured model (``ollama/...``, ``openai/...``) and a model set by
235 the path of a GGUF file have no registry entry. The chat routes serve both
236 through the known-model lookup, so the lookup decides, and every listing
237 that advertises models to a client agrees with the routes. An empty
238 *configured* means the role is unconfigured.
239 """
240 if not configured or configured in refs:
241 return list(refs)
242 if parse_model_ref(configured).is_remote or (
243 get_services().known_models.resolve(configured) is not None
244 ):
245 return [configured, *refs]
246 return list(refs)
249def _native_manifest_index() -> dict[str, ModelManifest]:
250 """Map ref string ('hf_repo/filename') to manifest for every installed native model."""
251 registry = ModelRegistry(cfg.models_dir)
252 return {m.ref: m for m in registry.list_installed()}
255def _resolve_native_path(ref: str) -> str | None:
256 """Return the on-disk path of an installed native model, if resolvable.
258 Swallows ``KeyError`` (manifest present but blob missing) and
259 ``ValueError`` (malformed ref) so callers can treat the path as
260 optional metadata.
261 """
262 try:
263 return str(ModelRegistry(cfg.models_dir).resolve(ref))
264 except (KeyError, ValueError):
265 return None
268def _collect_native_entries() -> list[ModelEntry]:
269 manifests = _native_manifest_index()
270 refs = get_services().model_manager.list_installed(source=ModelSource.NATIVE)
271 return [ModelEntry.from_native(ref, manifests.get(ref)) for ref in refs]
274def _collect_backend_entries() -> list[ModelEntry]:
275 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
276 from lilbee.modelhub.model_manager import classify_all_remote_models
277 from lilbee.providers.local_servers import local_server_for_label
279 def _source(remote: RemoteModel) -> ModelSource:
280 spec = local_server_for_label(remote.provider)
281 return ModelSource(spec.key) if spec is not None else ModelSource.REMOTE
283 remote_by_name = {
284 rm.name: rm for rm in classify_all_remote_models(timeout=_BACKEND_LIST_TIMEOUT_S)
285 }
286 return [
287 ModelEntry.from_backend(name, rm, _source(rm))
288 for name, rm in sorted(remote_by_name.items())
289 ]
292def list_models_data(
293 source: ModelSource | None = None,
294 task: ModelTask | None = None,
295) -> ListModelsResult:
296 """Build the list of installed models with source and task metadata.
298 Discovers remote models via a single HTTP call with a short timeout
299 so the listing stays responsive when the backend is down.
300 """
301 entries: list[ModelEntry] = []
302 if source is None or source is ModelSource.NATIVE:
303 entries.extend(_collect_native_entries())
304 if source is not ModelSource.NATIVE:
305 backend = _collect_backend_entries()
306 # A specific local-server source (ollama/lm_studio/frontier) narrows the
307 # backend list; REMOTE and None keep every backend entry.
308 if source is not None and source is not ModelSource.REMOTE:
309 backend = [e for e in backend if e.source is source]
310 entries.extend(backend)
311 if task:
312 entries = [e for e in entries if e.task == task]
313 return ListModelsResult(models=entries, total=len(entries))
316def show_model_data(ref: str) -> ShowModelResult:
317 """Return catalog and install metadata for *ref*.
319 Raises :class:`~lilbee.modelhub.model_manager.ModelNotFoundError` if the ref
320 is unknown to both the catalog and the installed set.
321 """
322 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
323 from lilbee.catalog import find_pick
324 from lilbee.modelhub.model_manager import ModelNotFoundError
326 entry = find_pick(ref)
327 source = get_services().model_manager.get_source(ref)
328 if entry is None and source is None:
329 raise ModelNotFoundError(f"model not found: {ref}")
330 manifest = _native_manifest_index().get(ref)
331 return ShowModelResult(
332 model=ref,
333 catalog=CatalogEntryData.from_catalog_model(entry) if entry else None,
334 installed=source is not None,
335 source=source.value if source else None,
336 manifest=ManifestData.from_manifest(manifest) if manifest else None,
337 path=_resolve_native_path(ref) if manifest is not None else None,
338 )
341def _vision_projector_missing(ref: str) -> bool:
342 """True when *ref*'s mmproj projector does not resolve on disk."""
343 from lilbee.providers.base import ProviderError
344 from lilbee.providers.engine_params import resolve_model_path
345 from lilbee.providers.gguf_meta import find_mmproj_for_model
347 try:
348 find_mmproj_for_model(resolve_model_path(ref))
349 except (ProviderError, OSError, ValueError, KeyError):
350 return True
351 return False
354def _ensure_vision_projector(ref: str) -> None:
355 """Fetch a vision model's mmproj projector when a cached install lacks it.
357 No-op for non-vision refs and when the projector already resolves on disk,
358 so pulling a complete install never touches the network.
359 """
360 from lilbee.catalog import download_mmproj, resolve_pull_target
362 entry = resolve_pull_target(ref)
363 if entry is not None and entry.task is ModelTask.VISION and _vision_projector_missing(ref):
364 download_mmproj(entry)
367def pull_model_data(
368 ref: str,
369 source: ModelSource,
370 *,
371 on_update: Callable[[DownloadProgress], None] | None = None,
372 allow_unsupported: bool = False,
373 cancel: CancelSignal | None = None,
374) -> PullResult:
375 """Pull *ref* from *source* and return a typed result.
377 Only native models are downloadable; a non-native *source* is refused by
378 :meth:`ModelManager.pull`. Progress updates are throttled by
379 :func:`~lilbee.catalog.make_download_callback`, so callers see at most
380 roughly 10 Hz of progress events. A *cancel* signal makes the download
381 cancellable mid-transfer; see :func:`~lilbee.catalog.download_model`.
382 """
383 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
384 from lilbee.catalog import make_download_callback
386 manager = get_services().model_manager
388 # A GGUF found only in the HF cache is not listed until a manifest records it.
389 # Pulling again registers it, or fails with the reason it cannot.
390 if manager.is_installed(ref, source) and (
391 source is not ModelSource.NATIVE or manager.is_registered(ref)
392 ):
393 # A cached vision install may carry the main GGUF but not its mmproj
394 # projector; without it llama-server can't serve OCR, so ensure it before
395 # reporting already-installed.
396 _ensure_vision_projector(ref)
397 return PullResult(model=ref, source=source.value, status=PullStatus.ALREADY_INSTALLED)
399 bytes_cb = make_download_callback(on_update) if on_update is not None else None
400 path = manager.pull(
401 ref,
402 source,
403 on_bytes=bytes_cb,
404 allow_unsupported=allow_unsupported,
405 cancel=cancel,
406 )
407 return PullResult(
408 model=ref,
409 source=source.value,
410 status=PullStatus.OK,
411 path=str(path) if path is not None else None,
412 )
415def _legacy_disk_size(ref: str, *, fallback: int) -> int:
416 """Sum a split GGUF's shard sizes on disk; *fallback* on any failure/single file."""
417 import contextlib
419 with contextlib.suppress(Exception):
420 shards = get_services().registry.shard_paths(ref)
421 if len(shards) > 1:
422 return sum(path.stat().st_size for path in shards)
423 return fallback
426def remove_model_data(
427 ref: str,
428 source: ModelSource | None = None,
429) -> RemoveResult:
430 """Remove *ref* and return a typed result with freed size."""
431 manager = get_services().model_manager
432 manifests = _native_manifest_index()
433 # disk_size_bytes is the full multi-shard total; size_bytes alone would report
434 # only the first shard for a split GGUF.
435 manifest = manifests.get(ref)
436 size_bytes = manifest.disk_size_bytes if manifest is not None else 0
437 if manifest is not None and manifest.total_size_bytes is None:
438 # Legacy manifest without shard accounting: removal still frees every
439 # shard, so recover the true on-disk total from the shards for the report.
440 size_bytes = _legacy_disk_size(ref, fallback=size_bytes)
441 removed = manager.remove(ref, source=source)
442 return RemoveResult(
443 model=ref,
444 deleted=removed,
445 freed_gb=_bytes_to_gb(size_bytes),
446 )