Coverage for src/lilbee/modelhub/model_manager/core.py: 100%
167 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-31 21:55 +0000
1"""ModelManager: native and SDK-backed model lifecycle operations."""
3import logging
4import time
5from collections.abc import Callable
6from dataclasses import replace
7from pathlib import Path
9from lilbee.catalog.compat import UnsupportedQuantError
10from lilbee.catalog.models import CatalogModel
11from lilbee.catalog.types import ModelSource
12from lilbee.core.config import DEFAULT_HTTP_TIMEOUT
13from lilbee.core.security import validate_path_within
14from lilbee.modelhub.model_manager.types import ModelNotFoundError
15from lilbee.modelhub.registry import ModelRegistry
16from lilbee.providers.local_servers import LOCAL_SERVERS, local_server_for_key
17from lilbee.providers.model_ref import parse_model_ref
18from lilbee.runtime.cancellation import CancelSignal
20log = logging.getLogger(__name__)
22_INSTALLED_CACHE_TTL_SECONDS = 30.0
25def _prefixed_source(model: str) -> ModelSource | None:
26 """Map a provider-prefixed ref to its source, or ``None`` for a bare ref.
28 Local-server prefixes (``ollama/``, ``lm_studio/``) map to that server's
29 source; API-provider prefixes are FRONTIER. A bare name returns ``None``
30 so the caller falls back to backend membership.
31 """
32 for spec in LOCAL_SERVERS:
33 if model.startswith(spec.wire_prefix):
34 return ModelSource(spec.key)
35 try:
36 ref = parse_model_ref(model)
37 except ValueError:
38 return None
39 return ModelSource.FRONTIER if ref.is_api else None
42class ModelManager:
43 """Manages model lifecycle with distinct sources."""
45 def __init__(self, models_dir: Path) -> None:
46 self._models_dir = models_dir
47 self._registry = ModelRegistry(self._models_dir)
48 # Memoize list_installed results to avoid walking the registry
49 # filesystem and hitting the backend HTTP endpoint on every call.
50 # The catalog filter path fires this per request. Time-based TTL
51 # plus explicit invalidation on pull/remove keeps freshness.
52 self._installed_cache: dict[ModelSource | None, tuple[float, list[str]]] = {}
53 # Identity cache: refs + hf_repos of installed natives. The catalog
54 # screen reads this to mark rows as installed without re-walking
55 # the registry on every screen mount (~150-300 ms saved).
56 self._native_identities_cache: tuple[float, frozenset[str]] | None = None
58 def list_installed(self, source: ModelSource | None = None) -> list[str]:
59 """List installed model names. ``source=None`` lists all sources.
61 Memoized with a ``_INSTALLED_CACHE_TTL_SECONDS`` TTL and
62 invalidated eagerly by ``pull``/``remove``.
63 """
64 now = time.monotonic()
65 cached = self._installed_cache.get(source)
66 if cached is not None:
67 cached_at, cached_result = cached
68 if now - cached_at < _INSTALLED_CACHE_TTL_SECONDS:
69 return cached_result
71 if source is None:
72 native = set(self._list_native())
73 remote = set(self._list_remote())
74 result = sorted(native | remote)
75 elif source is ModelSource.NATIVE:
76 result = self._list_native()
77 else:
78 result = self._list_remote()
80 self._installed_cache[source] = (now, result)
81 return result
83 def list_native_identities(self) -> frozenset[str]:
84 """Return refs + hf_repos of installed native models.
86 Same TTL as ``list_installed``. The catalog screen reads this to
87 mark catalog rows as installed without re-walking the registry
88 on every screen mount.
89 """
90 now = time.monotonic()
91 if self._native_identities_cache is not None:
92 cached_at, cached_result = self._native_identities_cache
93 if now - cached_at < _INSTALLED_CACHE_TTL_SECONDS:
94 return cached_result
95 identities: set[str] = set()
96 try:
97 for m in self._registry.list_installed():
98 identities.add(m.ref)
99 identities.add(m.hf_repo)
100 except Exception:
101 log.debug("ModelRegistry.list_installed failed", exc_info=True)
102 result = frozenset(identities)
103 self._native_identities_cache = (now, result)
104 return result
106 def _invalidate_installed_cache(self) -> None:
107 """Drop all cached list_installed results and the route-layer cache."""
108 self._installed_cache.clear()
109 self._native_identities_cache = None
110 from lilbee.app.services import peek_services
112 # peek_services is None for a standalone ModelManager (test setup);
113 # the route isn't running so there's nothing to invalidate.
114 services = peek_services()
115 if services is not None:
116 services.known_models.invalidate()
118 def _list_native(self) -> list[str]:
119 """List native models from the registry only."""
120 return sorted(m.ref for m in self._registry.list_installed())
122 def _list_remote(self) -> list[str]:
123 """List model names across every configured local server (Ollama, LM Studio).
125 Reuses the discovery dispatch so each listing endpoint matches its
126 server (Ollama ``/api/tags`` vs LM Studio ``/v1/models``). Returns
127 ``[]`` when the backends are unreachable.
128 """
129 # circular: discovery -> app.services -> model_manager.__init__ -> core
130 from lilbee.modelhub.model_manager.discovery import classify_all_remote_models
132 models = classify_all_remote_models(timeout=DEFAULT_HTTP_TIMEOUT)
133 return [m.name for m in models]
135 def is_installed(self, model: str, source: ModelSource | None = None) -> bool:
136 """Check if model exists in specified source."""
137 if source is None:
138 return self._is_native(model) or self._is_remote(model)
139 if source is ModelSource.NATIVE:
140 return self._is_native(model)
141 return self._is_remote(model)
143 def _is_native(self, model: str) -> bool:
144 if self._registry.is_installed(model):
145 return True
146 try:
147 validate_path_within(self._models_dir / model, self._models_dir)
148 except ValueError:
149 return False
150 return (self._models_dir / model).is_file()
152 def _is_remote(self, model: str) -> bool:
153 return model in self.list_installed(ModelSource.REMOTE)
155 def get_source(self, model: str) -> ModelSource | None:
156 """Return the granular source a model lives in. Native takes precedence.
158 A provider-prefixed ref classifies without a network call; a bare name
159 that a backend reports installed is ``REMOTE`` (the prefix is what names
160 the specific server). ``None`` when the model is in no known source.
161 """
162 if self._is_native(model):
163 return ModelSource.NATIVE
164 prefixed = _prefixed_source(model)
165 if prefixed is not None:
166 return prefixed
167 if self._is_remote(model):
168 return ModelSource.REMOTE
169 return None
171 def pull(
172 self,
173 model: str,
174 source: ModelSource,
175 *,
176 on_bytes: Callable[[int, int], None] | None = None,
177 allow_unsupported: bool = False,
178 cancel: CancelSignal | None = None,
179 ) -> Path | None:
180 """Download a native GGUF model and return its path.
182 lilbee pulls native models only. Local servers (Ollama, LM Studio)
183 are read-only: their models are managed in their own app and surface
184 here once present, so a non-native *source* is refused.
186 Native pulls of architectures the bundled llama.cpp doesn't support
187 are refused with ``UnsupportedArchError``, and files whose tensors it
188 cannot decode with ``UnsupportedQuantError``, unless *allow_unsupported*
189 is True. *on_bytes* receives (downloaded_bytes, total_bytes) progress.
190 A *cancel* signal makes the download cancellable mid-transfer; see
191 :func:`~lilbee.catalog.download_model`.
192 """
193 if source is not ModelSource.NATIVE:
194 spec = local_server_for_key(source.value)
195 where = spec.display_name if spec is not None else "the configured server"
196 raise ValueError(
197 f"lilbee runs {where} models but doesn't download them. "
198 f"Add the model in {where}, then pick it here."
199 )
200 if not allow_unsupported:
201 self.enforce_arch_compat(model)
202 try:
203 return self._pull_native(
204 model, on_bytes=on_bytes, cancel=cancel, allow_unsupported=allow_unsupported
205 )
206 finally:
207 self._invalidate_installed_cache()
209 def enforce_arch_compat(self, ref: str) -> None:
210 """Raise UnsupportedArchError if *ref*'s architecture isn't in the supported set.
212 Public because the pull preflight on the HTTP surface runs the same
213 check before starting a download, so a caller learns the model is
214 unsupported before any bytes move.
215 """
216 from lilbee.app.services import get_services
217 from lilbee.catalog.compat import (
218 ModelCompat,
219 UnsupportedArchError,
220 classify,
221 resolve_arch_for_pull,
222 )
224 arch = resolve_arch_for_pull(ref, get_services().hf_client)
225 if classify(arch) is ModelCompat.UNSUPPORTED:
226 raise UnsupportedArchError(ref, arch)
228 def _pull_native(
229 self,
230 model: str,
231 *,
232 on_bytes: Callable[[int, int], None] | None = None,
233 cancel: CancelSignal | None = None,
234 allow_unsupported: bool = False,
235 ) -> Path:
236 """Download a featured or ad-hoc HuggingFace model to the native GGUF directory."""
237 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
238 from lilbee.catalog import download_model, resolve_pull_target
239 from lilbee.modelhub.registry import register_downloaded_model
241 entry = resolve_pull_target(model)
242 if entry is None:
243 raise ModelNotFoundError(
244 f"Model '{model}' not recognized. "
245 "Pass a HuggingFace repo id (owner/name) or a featured model name."
246 )
247 entry = self._resolved_entry(entry, verify=not allow_unsupported)
248 path = download_model(
249 entry, on_progress=on_bytes, on_complete=register_downloaded_model, cancel=cancel
250 )
251 log.info("Downloaded %s to %s", model, path)
252 return path
254 @staticmethod
255 def _resolved_entry(entry: CatalogModel, *, verify: bool) -> CatalogModel:
256 """*entry* carrying the file the pull will fetch.
258 Resolution happens here because the engine's verdict is part of choosing
259 the file, not a check on the file already chosen. A repo can publish the
260 same weights in several packings, only some of which this build reads,
261 and the resolver walks its ranking until one of them answers. Passing the
262 resolved name down means the download does not work it out again.
264 A repo that will not resolve is not this check's verdict to deliver. The
265 download reports the gated repo or the missing file properly, so a
266 resolution failure here is left to it.
267 """
268 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
269 from lilbee.catalog.download import resolve_filename
270 from lilbee.catalog.hf_client import hf_token
271 from lilbee.providers.fleet.loadability import assert_engine_can_load
273 token = hf_token()
275 def _can_load(hf_repo: str, filename: str) -> None:
276 assert_engine_can_load(hf_repo, filename, token)
278 try:
279 filename = resolve_filename(entry, can_load=_can_load if verify else None)
280 except (PermissionError, RuntimeError) as exc:
281 if isinstance(exc, UnsupportedQuantError):
282 raise
283 log.debug("Cannot name a file to fetch for %s: %s", entry.hf_repo, exc)
284 return entry
285 return replace(entry, gguf_filename=filename)
287 def remove(self, model: str, source: ModelSource | None = None) -> bool:
288 """Remove an installed native model. Returns True if removed.
290 lilbee removes only native GGUF models it downloaded. Local servers
291 (Ollama, LM Studio) are read-only: a model that lives on one is refused
292 (mirrors ``pull``), since its lifecycle is managed in that app. A bare
293 ``source`` is resolved so a local-server ref is caught either way.
294 """
295 effective = source if source is not None else self.get_source(model)
296 if effective is not None and effective is not ModelSource.NATIVE:
297 spec = local_server_for_key(effective.value)
298 where = spec.display_name if spec is not None else "the configured server"
299 raise ValueError(
300 f"lilbee runs {where} models but doesn't remove them. "
301 f"Manage them in {where} instead."
302 )
303 try:
304 return self._remove_native(model)
305 finally:
306 self._invalidate_installed_cache()
308 def _remove_native(self, model: str) -> bool:
309 if self._registry.remove(model):
310 log.info("Removed native model %s from registry", model)
311 return True
312 try:
313 path = validate_path_within(self._models_dir / model, self._models_dir)
314 except ValueError:
315 log.warning("Path traversal blocked: %s escapes %s", model, self._models_dir)
316 return False
317 if path.is_file():
318 path.unlink()
319 log.info("Removed native model %s", model)
320 return True
321 return False