Coverage for src / lilbee / app / models.py: 100%
171 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-06-28 01:01 +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
24_BYTES_PER_GB = 1024**3 # Model sizes are reported to users in GiB.
25_BACKEND_LIST_TIMEOUT_S = 2.0 # Keep `model list` snappy when backend is down.
28def _bytes_to_gb(n: int) -> float:
29 """Convert bytes to GiB rounded to 2 decimals for user display."""
30 return round(n / _BYTES_PER_GB, 2)
33class ModelCommand(StrEnum):
34 """Command field values for model sub-app JSON output."""
36 LIST = "model list"
37 SHOW = "model show"
38 PULL = "model pull"
39 RM = "model rm"
42class PullStatus(StrEnum):
43 OK = "ok"
44 ALREADY_INSTALLED = "already_installed"
47class AdoptStatus(StrEnum):
48 ADOPTED = "adopted"
49 ALREADY_ACTIVE = "already_active"
52class PullEvent(StrEnum):
53 PROGRESS = "progress"
54 DONE = "done"
57class ModelEntry(BaseModel):
58 """One row of `lilbee model list` output."""
60 name: str
61 source: str
62 task: ModelTask | None = None
63 size_gb: float | None = None
64 display_name: str = ""
66 @classmethod
67 def from_native(cls, ref: str, manifest: ModelManifest | None) -> ModelEntry:
68 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
69 from lilbee.catalog import clean_display_name
70 from lilbee.catalog.types import ModelSource
72 return cls(
73 name=ref,
74 source=ModelSource.NATIVE.value,
75 task=manifest.task if manifest else None,
76 size_gb=_bytes_to_gb(manifest.size_bytes) if manifest else None,
77 display_name=clean_display_name(manifest.hf_repo) if manifest else "",
78 )
80 @classmethod
81 def from_backend(cls, ref: str, remote: RemoteModel | None, source: ModelSource) -> ModelEntry:
82 from lilbee.providers.local_servers import canonical_local_ref
84 return cls(
85 name=canonical_local_ref(ref, source.value),
86 source=source.value,
87 task=remote.task if remote else None,
88 size_gb=None,
89 display_name=remote.parameter_size if remote else "",
90 )
93class ListModelsResult(BaseModel):
94 command: str = ModelCommand.LIST
95 models: list[ModelEntry]
96 total: int
99class CatalogEntryData(BaseModel):
100 ref: str
101 display_name: str
102 hf_repo: str
103 gguf_filename: str
104 size_gb: float
105 min_ram_gb: float
106 description: str
107 task: ModelTask
108 featured: bool
109 recommended: 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 recommended=entry.recommended,
126 architecture=entry.architecture,
127 compat=entry.compat,
128 )
131class ManifestData(BaseModel):
132 ref: str
133 display_name: str
134 task: ModelTask
135 size_gb: float
136 size_bytes: int
137 hf_repo: str
138 gguf_filename: str
139 downloaded_at: str
141 @classmethod
142 def from_manifest(cls, manifest: ModelManifest) -> ManifestData:
143 from lilbee.catalog import clean_display_name
145 return cls(
146 ref=manifest.ref,
147 display_name=clean_display_name(manifest.hf_repo),
148 task=manifest.task,
149 size_gb=_bytes_to_gb(manifest.size_bytes),
150 size_bytes=manifest.size_bytes,
151 hf_repo=manifest.hf_repo,
152 gguf_filename=manifest.gguf_filename,
153 downloaded_at=manifest.downloaded_at,
154 )
157class ShowModelResult(BaseModel):
158 command: str = ModelCommand.SHOW
159 model: str
160 catalog: CatalogEntryData | None = None
161 installed: bool = False
162 source: str | None = None
163 path: str | None = None
164 manifest: ManifestData | None = None
167class PullResult(BaseModel):
168 command: str = ModelCommand.PULL
169 model: str
170 source: str
171 status: PullStatus
172 path: str | None = None
175class PullProgressEvent(BaseModel):
176 command: str = ModelCommand.PULL
177 event: str = PullEvent.PROGRESS
178 model: str
179 percent: float
180 detail: str
181 cache_hit: bool
184class RemoveResult(BaseModel):
185 command: str = ModelCommand.RM
186 model: str
187 deleted: bool
188 freed_gb: float = Field(default=0.0)
191class AdoptResult(BaseModel):
192 """Outcome of adopting a downloaded index's embedder."""
194 model: str
195 status: AdoptStatus
196 reindex_required: bool = False
199def adopt_embedder(ref: str) -> AdoptResult:
200 """Switch lilbee to embedder *ref*, downloading it first if missing.
202 Makes a downloaded index searchable under its own embedder without a
203 rebuild: the persisted vectors already match *ref*, so the switch routes
204 through the settings boundary and ``reindex_required`` stays false.
205 """
206 from lilbee.app.settings import apply_settings_update
207 from lilbee.catalog.types import ModelSource
209 manager = get_services().model_manager
210 already_active = cfg.embedding_model == ref and manager.is_installed(ref, ModelSource.NATIVE)
211 if not manager.is_installed(ref, ModelSource.NATIVE):
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 _native_manifest_index() -> dict[str, ModelManifest]:
222 """Map ref string ('hf_repo/filename') to manifest for every installed native model."""
223 registry = ModelRegistry(cfg.models_dir)
224 return {m.ref: m for m in registry.list_installed()}
227def _resolve_native_path(ref: str) -> str | None:
228 """Return the on-disk path of an installed native model, if resolvable.
230 Swallows ``KeyError`` (manifest present but blob missing) and
231 ``ValueError`` (malformed ref) so callers can treat the path as
232 optional metadata.
233 """
234 try:
235 return str(ModelRegistry(cfg.models_dir).resolve(ref))
236 except (KeyError, ValueError):
237 return None
240def _collect_native_entries() -> list[ModelEntry]:
241 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
242 from lilbee.catalog.types import ModelSource
244 manifests = _native_manifest_index()
245 refs = get_services().model_manager.list_installed(source=ModelSource.NATIVE)
246 return [ModelEntry.from_native(ref, manifests.get(ref)) for ref in refs]
249def _collect_backend_entries() -> list[ModelEntry]:
250 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
251 from lilbee.catalog.types import ModelSource
252 from lilbee.modelhub.model_manager import classify_all_remote_models
253 from lilbee.providers.local_servers import local_server_for_label
255 def _source(remote: RemoteModel) -> ModelSource:
256 spec = local_server_for_label(remote.provider)
257 return ModelSource(spec.key) if spec is not None else ModelSource.REMOTE
259 remote_by_name = {
260 rm.name: rm for rm in classify_all_remote_models(timeout=_BACKEND_LIST_TIMEOUT_S)
261 }
262 return [
263 ModelEntry.from_backend(name, rm, _source(rm))
264 for name, rm in sorted(remote_by_name.items())
265 ]
268def list_models_data(
269 source: ModelSource | None = None,
270 task: ModelTask | None = None,
271) -> ListModelsResult:
272 """Build the list of installed models with source and task metadata.
274 Discovers remote models via a single HTTP call with a short timeout
275 so the command stays responsive when the backend is down.
276 """
277 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
278 from lilbee.catalog.types import ModelSource
280 entries: list[ModelEntry] = []
281 if source is None or source is ModelSource.NATIVE:
282 entries.extend(_collect_native_entries())
283 if source is not ModelSource.NATIVE:
284 backend = _collect_backend_entries()
285 # A specific local-server source (ollama/lm_studio/frontier) narrows the
286 # backend list; REMOTE and None keep every backend entry.
287 if source is not None and source is not ModelSource.REMOTE:
288 backend = [e for e in backend if e.source == source.value]
289 entries.extend(backend)
290 if task:
291 entries = [e for e in entries if e.task == task]
292 return ListModelsResult(models=entries, total=len(entries))
295def show_model_data(ref: str) -> ShowModelResult:
296 """Return catalog and install metadata for *ref*.
298 Raises :class:`~lilbee.modelhub.model_manager.ModelNotFoundError` if the ref
299 is unknown to both the catalog and the installed set.
300 """
301 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
302 from lilbee.catalog import find_catalog_entry
303 from lilbee.modelhub.model_manager import ModelNotFoundError
305 entry = find_catalog_entry(ref)
306 source = get_services().model_manager.get_source(ref)
307 if entry is None and source is None:
308 raise ModelNotFoundError(f"model not found: {ref}")
309 manifest = _native_manifest_index().get(ref)
310 return ShowModelResult(
311 model=ref,
312 catalog=CatalogEntryData.from_catalog_model(entry) if entry else None,
313 installed=source is not None,
314 source=source.value if source else None,
315 manifest=ManifestData.from_manifest(manifest) if manifest else None,
316 path=_resolve_native_path(ref) if manifest is not None else None,
317 )
320def pull_model_data(
321 ref: str,
322 source: ModelSource,
323 *,
324 on_update: Callable[[DownloadProgress], None] | None = None,
325 allow_unsupported: bool = False,
326) -> PullResult:
327 """Pull *ref* from *source* and return a typed result.
329 Only native models are downloadable; a non-native *source* is refused by
330 :meth:`ModelManager.pull`. Progress updates are throttled by
331 :func:`~lilbee.catalog.make_download_callback`, so callers see at most
332 roughly 10 Hz of progress events.
333 """
334 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
335 from lilbee.catalog import make_download_callback
337 manager = get_services().model_manager
339 if manager.is_installed(ref, source):
340 return PullResult(model=ref, source=source.value, status=PullStatus.ALREADY_INSTALLED)
342 bytes_cb = make_download_callback(on_update) if on_update is not None else None
343 path = manager.pull(
344 ref,
345 source,
346 on_bytes=bytes_cb,
347 allow_unsupported=allow_unsupported,
348 )
349 return PullResult(
350 model=ref,
351 source=source.value,
352 status=PullStatus.OK,
353 path=str(path) if path is not None else None,
354 )
357def remove_model_data(
358 ref: str,
359 source: ModelSource | None = None,
360) -> RemoveResult:
361 """Remove *ref* and return a typed result with freed size."""
362 manager = get_services().model_manager
363 manifests = _native_manifest_index()
364 size_bytes = manifests[ref].size_bytes if ref in manifests else 0
365 removed = manager.remove(ref, source=source)
366 return RemoveResult(
367 model=ref,
368 deleted=removed,
369 freed_gb=_bytes_to_gb(size_bytes),
370 )