Coverage for src / lilbee / app / models.py: 100%
160 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-05-15 20:55 +0000
1"""Surface-agnostic model lifecycle use-cases (list / show / pull / remove)."""
3from __future__ import annotations
5import functools
6from enum import StrEnum
7from typing import TYPE_CHECKING, Any
9from pydantic import BaseModel, Field
11from lilbee.app.services import get_services
12from lilbee.catalog.types import ModelTask
13from lilbee.core.config import cfg
14from lilbee.modelhub.registry import ModelRegistry
16if TYPE_CHECKING:
17 from collections.abc import Callable
19 from lilbee.catalog import CatalogModel, DownloadProgress
20 from lilbee.catalog.types import ModelSource
21 from lilbee.modelhub.model_manager import RemoteModel
22 from lilbee.modelhub.registry import ModelManifest
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 PullEvent(StrEnum):
49 PROGRESS = "progress"
50 DONE = "done"
53class ModelEntry(BaseModel):
54 """One row of `lilbee model list` output."""
56 name: str
57 source: str
58 task: ModelTask | None = None
59 size_gb: float | None = None
60 display_name: str = ""
62 @classmethod
63 def from_native(cls, ref: str, manifest: ModelManifest | None) -> ModelEntry:
64 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
65 from lilbee.catalog import clean_display_name
66 from lilbee.catalog.types import ModelSource
68 return cls(
69 name=ref,
70 source=ModelSource.NATIVE.value,
71 task=manifest.task if manifest else None,
72 size_gb=_bytes_to_gb(manifest.size_bytes) if manifest else None,
73 display_name=clean_display_name(manifest.hf_repo) if manifest else "",
74 )
76 @classmethod
77 def from_backend(cls, ref: str, remote: RemoteModel | None) -> ModelEntry:
78 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
79 from lilbee.catalog.types import ModelSource
81 return cls(
82 name=ref,
83 source=ModelSource.REMOTE.value,
84 task=remote.task if remote else None,
85 size_gb=None,
86 display_name=remote.parameter_size if remote else "",
87 )
90class ListModelsResult(BaseModel):
91 command: str = ModelCommand.LIST
92 models: list[ModelEntry]
93 total: int
96class CatalogEntryData(BaseModel):
97 ref: str
98 display_name: str
99 hf_repo: str
100 gguf_filename: str
101 size_gb: float
102 min_ram_gb: float
103 description: str
104 task: ModelTask
105 featured: bool
106 recommended: bool
108 @classmethod
109 def from_catalog_model(cls, entry: CatalogModel) -> CatalogEntryData:
110 return cls(
111 ref=entry.ref,
112 display_name=entry.display_name,
113 hf_repo=entry.hf_repo,
114 gguf_filename=entry.gguf_filename,
115 size_gb=entry.size_gb,
116 min_ram_gb=entry.min_ram_gb,
117 description=entry.description,
118 task=entry.task,
119 featured=entry.featured,
120 recommended=entry.recommended,
121 )
124class ManifestData(BaseModel):
125 ref: str
126 display_name: str
127 task: ModelTask
128 size_gb: float
129 size_bytes: int
130 hf_repo: str
131 gguf_filename: str
132 downloaded_at: str
134 @classmethod
135 def from_manifest(cls, manifest: ModelManifest) -> ManifestData:
136 from lilbee.catalog import clean_display_name
138 return cls(
139 ref=manifest.ref,
140 display_name=clean_display_name(manifest.hf_repo),
141 task=manifest.task,
142 size_gb=_bytes_to_gb(manifest.size_bytes),
143 size_bytes=manifest.size_bytes,
144 hf_repo=manifest.hf_repo,
145 gguf_filename=manifest.gguf_filename,
146 downloaded_at=manifest.downloaded_at,
147 )
150class ShowModelResult(BaseModel):
151 command: str = ModelCommand.SHOW
152 model: str
153 catalog: CatalogEntryData | None = None
154 installed: bool = False
155 source: str | None = None
156 path: str | None = None
157 manifest: ManifestData | None = None
160class PullResult(BaseModel):
161 command: str = ModelCommand.PULL
162 model: str
163 source: str
164 status: PullStatus
165 path: str | None = None
168class PullProgressEvent(BaseModel):
169 command: str = ModelCommand.PULL
170 event: str = PullEvent.PROGRESS
171 model: str
172 percent: float
173 detail: str
174 cache_hit: bool
177class RemoveResult(BaseModel):
178 command: str = ModelCommand.RM
179 model: str
180 deleted: bool
181 freed_gb: float = Field(default=0.0)
184def _native_manifest_index() -> dict[str, ModelManifest]:
185 """Map ref string ('hf_repo/filename') to manifest for every installed native model."""
186 registry = ModelRegistry(cfg.models_dir)
187 return {m.ref: m for m in registry.list_installed()}
190def _resolve_native_path(ref: str) -> str | None:
191 """Return the on-disk path of an installed native model, if resolvable.
193 Swallows ``KeyError`` (manifest present but blob missing) and
194 ``ValueError`` (malformed ref) so callers can treat the path as
195 optional metadata.
196 """
197 try:
198 return str(ModelRegistry(cfg.models_dir).resolve(ref))
199 except (KeyError, ValueError):
200 return None
203def _collect_native_entries() -> list[ModelEntry]:
204 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
205 from lilbee.catalog.types import ModelSource
207 manifests = _native_manifest_index()
208 refs = get_services().model_manager.list_installed(source=ModelSource.NATIVE)
209 return [ModelEntry.from_native(ref, manifests.get(ref)) for ref in refs]
212def _collect_backend_entries() -> list[ModelEntry]:
213 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
214 from lilbee.modelhub.model_manager import classify_remote_models
216 remote_list = classify_remote_models(cfg.remote_base_url, timeout=_BACKEND_LIST_TIMEOUT_S)
217 remote_by_name = {rm.name: rm for rm in remote_list}
218 return [ModelEntry.from_backend(ref, remote_by_name[ref]) for ref in sorted(remote_by_name)]
221def list_models_data(
222 source: ModelSource | None = None,
223 task: ModelTask | None = None,
224) -> ListModelsResult:
225 """Build the list of installed models with source and task metadata.
227 Discovers remote models via a single HTTP call with a short timeout
228 so the command stays responsive when the backend is down.
229 """
230 # heavy: lilbee.modelhub.model_manager (>50ms; huggingface_hub fanout)
231 from lilbee.catalog.types import ModelSource
233 entries: list[ModelEntry] = []
234 if source is None or source is ModelSource.NATIVE:
235 entries.extend(_collect_native_entries())
236 if source is None or source is ModelSource.REMOTE:
237 entries.extend(_collect_backend_entries())
238 if task:
239 entries = [e for e in entries if e.task == task]
240 return ListModelsResult(models=entries, total=len(entries))
243def show_model_data(ref: str) -> ShowModelResult:
244 """Return catalog and install metadata for *ref*.
246 Raises :class:`~lilbee.modelhub.model_manager.ModelNotFoundError` if the ref
247 is unknown to both the catalog and the installed set.
248 """
249 # heavy: lilbee.catalog (>50ms; huggingface_hub) + lilbee.modelhub.model_manager (>50ms)
250 from lilbee.catalog import find_catalog_entry
251 from lilbee.modelhub.model_manager import ModelNotFoundError
253 entry = find_catalog_entry(ref)
254 source = get_services().model_manager.get_source(ref)
255 if entry is None and source is None:
256 raise ModelNotFoundError(f"model not found: {ref}")
257 manifest = _native_manifest_index().get(ref)
258 return ShowModelResult(
259 model=ref,
260 catalog=CatalogEntryData.from_catalog_model(entry) if entry else None,
261 installed=source is not None,
262 source=source.value if source else None,
263 manifest=ManifestData.from_manifest(manifest) if manifest else None,
264 path=_resolve_native_path(ref) if manifest is not None else None,
265 )
268def _backend_event_to_progress(
269 on_update: Callable[[DownloadProgress], None],
270 event: dict[str, Any],
271) -> None:
272 """Adapt an Ollama-style dict event into a DownloadProgress call."""
273 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
274 from lilbee.catalog import DownloadProgress
276 total = event.get("total", 0) or 0
277 completed = event.get("completed", 0) or 0
278 detail = event.get("status", "") or ""
279 pct = int(completed * 100 / total) if total > 0 else 0
280 on_update(DownloadProgress(percent=pct, detail=detail, is_cache_hit=False))
283def _build_pull_callbacks(
284 on_update: Callable[[DownloadProgress], None] | None,
285) -> tuple[Callable[[dict[str, Any]], None] | None, Callable[[int, int], None] | None]:
286 """Build the (dict_cb, bytes_cb) pair for ModelManager.pull from on_update."""
287 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
288 from lilbee.catalog import make_download_callback
290 if on_update is None:
291 return None, None
292 dict_cb = functools.partial(_backend_event_to_progress, on_update)
293 bytes_cb = make_download_callback(on_update)
294 return dict_cb, bytes_cb
297def pull_model_data(
298 ref: str,
299 source: ModelSource,
300 *,
301 on_update: Callable[[DownloadProgress], None] | None = None,
302) -> PullResult:
303 """Pull *ref* from *source* and return a typed result.
305 Progress updates are throttled by
306 :func:`~lilbee.catalog.make_download_callback`, so callers see at
307 most roughly 10 Hz of progress events.
308 """
309 manager = get_services().model_manager
311 if manager.is_installed(ref, source):
312 return PullResult(model=ref, source=source.value, status=PullStatus.ALREADY_INSTALLED)
314 dict_cb, bytes_cb = _build_pull_callbacks(on_update)
315 path = manager.pull(ref, source, on_progress=dict_cb, on_bytes=bytes_cb)
316 return PullResult(
317 model=ref,
318 source=source.value,
319 status=PullStatus.OK,
320 path=str(path) if path is not None else None,
321 )
324def remove_model_data(
325 ref: str,
326 source: ModelSource | None = None,
327) -> RemoveResult:
328 """Remove *ref* and return a typed result with freed size."""
329 manager = get_services().model_manager
330 manifests = _native_manifest_index()
331 size_bytes = manifests[ref].size_bytes if ref in manifests else 0
332 removed = manager.remove(ref, source=source)
333 return RemoveResult(
334 model=ref,
335 deleted=removed,
336 freed_gb=_bytes_to_gb(size_bytes),
337 )