Coverage for src/lilbee/modelhub/model_manager/core.py: 100%
146 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""ModelManager: native and SDK-backed model lifecycle operations."""
3import logging
4import time
5from collections.abc import Callable
6from pathlib import Path
8from lilbee.catalog.types import ModelSource
9from lilbee.core.config import DEFAULT_HTTP_TIMEOUT
10from lilbee.core.security import validate_path_within
11from lilbee.modelhub.model_manager.types import ModelNotFoundError
12from lilbee.modelhub.registry import ModelRegistry
13from lilbee.providers.local_servers import LOCAL_SERVERS, local_server_for_key
14from lilbee.providers.model_ref import parse_model_ref
16log = logging.getLogger(__name__)
18_INSTALLED_CACHE_TTL_SECONDS = 30.0
21def _prefixed_source(model: str) -> ModelSource | None:
22 """Map a provider-prefixed ref to its source, or ``None`` for a bare ref.
24 Local-server prefixes (``ollama/``, ``lm_studio/``) map to that server's
25 source; API-provider prefixes are FRONTIER. A bare name returns ``None``
26 so the caller falls back to backend membership.
27 """
28 for spec in LOCAL_SERVERS:
29 if model.startswith(spec.wire_prefix):
30 return ModelSource(spec.key)
31 try:
32 ref = parse_model_ref(model)
33 except ValueError:
34 return None
35 return ModelSource.FRONTIER if ref.is_api else None
38class ModelManager:
39 """Manages model lifecycle with distinct sources."""
41 def __init__(self, models_dir: Path) -> None:
42 self._models_dir = models_dir
43 self._registry = ModelRegistry(self._models_dir)
44 # Memoize list_installed results to avoid walking the registry
45 # filesystem and hitting the backend HTTP endpoint on every call.
46 # The catalog filter path fires this per request. Time-based TTL
47 # plus explicit invalidation on pull/remove keeps freshness.
48 self._installed_cache: dict[ModelSource | None, tuple[float, list[str]]] = {}
49 # Identity cache: refs + hf_repos of installed natives. The catalog
50 # screen reads this to mark rows as installed without re-walking
51 # the registry on every screen mount (~150-300 ms saved).
52 self._native_identities_cache: tuple[float, frozenset[str]] | None = None
54 def list_installed(self, source: ModelSource | None = None) -> list[str]:
55 """List installed model names. ``source=None`` lists all sources.
57 Memoized with a ``_INSTALLED_CACHE_TTL_SECONDS`` TTL and
58 invalidated eagerly by ``pull``/``remove``.
59 """
60 now = time.monotonic()
61 cached = self._installed_cache.get(source)
62 if cached is not None:
63 cached_at, cached_result = cached
64 if now - cached_at < _INSTALLED_CACHE_TTL_SECONDS:
65 return cached_result
67 if source is None:
68 native = set(self._list_native())
69 remote = set(self._list_remote())
70 result = sorted(native | remote)
71 elif source is ModelSource.NATIVE:
72 result = self._list_native()
73 else:
74 result = self._list_remote()
76 self._installed_cache[source] = (now, result)
77 return result
79 def list_native_identities(self) -> frozenset[str]:
80 """Return refs + hf_repos of installed native models.
82 Same TTL as ``list_installed``. The catalog screen reads this to
83 mark catalog rows as installed without re-walking the registry
84 on every screen mount.
85 """
86 now = time.monotonic()
87 if self._native_identities_cache is not None:
88 cached_at, cached_result = self._native_identities_cache
89 if now - cached_at < _INSTALLED_CACHE_TTL_SECONDS:
90 return cached_result
91 identities: set[str] = set()
92 try:
93 for m in self._registry.list_installed():
94 identities.add(m.ref)
95 identities.add(m.hf_repo)
96 except Exception:
97 log.debug("ModelRegistry.list_installed failed", exc_info=True)
98 result = frozenset(identities)
99 self._native_identities_cache = (now, result)
100 return result
102 def _invalidate_installed_cache(self) -> None:
103 """Drop all cached list_installed results and the route-layer cache."""
104 self._installed_cache.clear()
105 self._native_identities_cache = None
106 from lilbee.app.services import peek_services
108 # peek_services is None for a standalone ModelManager (test setup);
109 # the route isn't running so there's nothing to invalidate.
110 services = peek_services()
111 if services is not None:
112 services.known_models.invalidate()
114 def _list_native(self) -> list[str]:
115 """List native models from the registry only."""
116 return sorted(m.ref for m in self._registry.list_installed())
118 def _list_remote(self) -> list[str]:
119 """List model names across every configured local server (Ollama, LM Studio).
121 Reuses the discovery dispatch so each listing endpoint matches its
122 server (Ollama ``/api/tags`` vs LM Studio ``/v1/models``). Returns
123 ``[]`` when the backends are unreachable.
124 """
125 # circular: discovery -> app.services -> model_manager.__init__ -> core
126 from lilbee.modelhub.model_manager.discovery import classify_all_remote_models
128 models = classify_all_remote_models(timeout=DEFAULT_HTTP_TIMEOUT)
129 return [m.name for m in models]
131 def is_installed(self, model: str, source: ModelSource | None = None) -> bool:
132 """Check if model exists in specified source."""
133 if source is None:
134 return self._is_native(model) or self._is_remote(model)
135 if source is ModelSource.NATIVE:
136 return self._is_native(model)
137 return self._is_remote(model)
139 def _is_native(self, model: str) -> bool:
140 if self._registry.is_installed(model):
141 return True
142 try:
143 validate_path_within(self._models_dir / model, self._models_dir)
144 except ValueError:
145 return False
146 return (self._models_dir / model).is_file()
148 def _is_remote(self, model: str) -> bool:
149 return model in self.list_installed(ModelSource.REMOTE)
151 def get_source(self, model: str) -> ModelSource | None:
152 """Return the granular source a model lives in. Native takes precedence.
154 A provider-prefixed ref classifies without a network call; a bare name
155 that a backend reports installed is ``REMOTE`` (the prefix is what names
156 the specific server). ``None`` when the model is in no known source.
157 """
158 if self._is_native(model):
159 return ModelSource.NATIVE
160 prefixed = _prefixed_source(model)
161 if prefixed is not None:
162 return prefixed
163 if self._is_remote(model):
164 return ModelSource.REMOTE
165 return None
167 def pull(
168 self,
169 model: str,
170 source: ModelSource,
171 *,
172 on_bytes: Callable[[int, int], None] | None = None,
173 allow_unsupported: bool = False,
174 ) -> Path | None:
175 """Download a native GGUF model and return its path.
177 lilbee pulls native models only. Local servers (Ollama, LM Studio)
178 are read-only: their models are managed in their own app and surface
179 here once present, so a non-native *source* is refused.
181 Native pulls of architectures the bundled llama.cpp doesn't support
182 are refused with ``UnsupportedArchError`` unless *allow_unsupported*
183 is True. *on_bytes* receives (downloaded_bytes, total_bytes) progress.
184 """
185 if source is not ModelSource.NATIVE:
186 spec = local_server_for_key(source.value)
187 where = spec.display_name if spec is not None else "the configured server"
188 raise ValueError(
189 f"lilbee runs {where} models but doesn't download them. "
190 f"Add the model in {where}, then pick it here."
191 )
192 if not allow_unsupported:
193 self.enforce_arch_compat(model)
194 try:
195 return self._pull_native(model, on_bytes=on_bytes)
196 finally:
197 self._invalidate_installed_cache()
199 def enforce_arch_compat(self, ref: str) -> None:
200 """Raise UnsupportedArchError if *ref*'s architecture isn't in the supported set.
202 Public because the pull preflight on the HTTP surface runs the same
203 check before starting a download, so a caller learns the model is
204 unsupported before any bytes move.
205 """
206 from lilbee.app.services import get_services
207 from lilbee.catalog.compat import (
208 ModelCompat,
209 UnsupportedArchError,
210 classify,
211 resolve_arch_for_pull,
212 )
214 arch = resolve_arch_for_pull(ref, get_services().hf_client)
215 if classify(arch) is ModelCompat.UNSUPPORTED:
216 raise UnsupportedArchError(ref, arch)
218 def _pull_native(
219 self,
220 model: str,
221 *,
222 on_bytes: Callable[[int, int], None] | None = None,
223 ) -> Path:
224 """Download a featured or ad-hoc HuggingFace model to the native GGUF directory."""
225 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
226 from lilbee.catalog import download_model, resolve_pull_target
227 from lilbee.modelhub.registry import register_downloaded_model
229 entry = resolve_pull_target(model)
230 if entry is None:
231 raise ModelNotFoundError(
232 f"Model '{model}' not recognized. "
233 "Pass a HuggingFace repo id (owner/name) or a featured model name."
234 )
235 path = download_model(entry, on_progress=on_bytes, on_complete=register_downloaded_model)
236 log.info("Downloaded %s to %s", model, path)
237 return path
239 def remove(self, model: str, source: ModelSource | None = None) -> bool:
240 """Remove an installed native model. Returns True if removed.
242 lilbee removes only native GGUF models it downloaded. Local servers
243 (Ollama, LM Studio) are read-only: a model that lives on one is refused
244 (mirrors ``pull``), since its lifecycle is managed in that app. A bare
245 ``source`` is resolved so a local-server ref is caught either way.
246 """
247 effective = source if source is not None else self.get_source(model)
248 if effective is not None and effective is not ModelSource.NATIVE:
249 spec = local_server_for_key(effective.value)
250 where = spec.display_name if spec is not None else "the configured server"
251 raise ValueError(
252 f"lilbee runs {where} models but doesn't remove them. "
253 f"Manage them in {where} instead."
254 )
255 try:
256 return self._remove_native(model)
257 finally:
258 self._invalidate_installed_cache()
260 def _remove_native(self, model: str) -> bool:
261 if self._registry.remove(model):
262 log.info("Removed native model %s from registry", model)
263 return True
264 try:
265 path = validate_path_within(self._models_dir / model, self._models_dir)
266 except ValueError:
267 log.warning("Path traversal blocked: %s escapes %s", model, self._models_dir)
268 return False
269 if path.is_file():
270 path.unlink()
271 log.info("Removed native model %s", model)
272 return True
273 return False