Coverage for src/lilbee/server/handlers/models.py: 100%
241 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""Model catalog, role assignment, install/delete, and external listing handlers."""
3from __future__ import annotations
5import asyncio
6import logging
7from collections.abc import AsyncGenerator
8from typing import TYPE_CHECKING, Literal
10from cachetools import TTLCache
11from pydantic import BaseModel
13from lilbee.app.services import get_services
14from lilbee.app.settings import apply_settings_update
15from lilbee.catalog import (
16 ModelFamily,
17 enrich_catalog,
18 get_catalog,
19 get_families,
20 get_picks,
21 picks_for,
22)
23from lilbee.catalog.refs import hf_repo_from_ref, is_bare_hf_repo
24from lilbee.catalog.types import CatalogSize, CatalogSort, KeyStatus, ModelSource, ModelTask
25from lilbee.core.config import cfg
26from lilbee.modelhub.model_manager import classify_all_remote_models, discover_api_models
27from lilbee.modelhub.model_manager.types import RemoteModel
28from lilbee.modelhub.role_validator import MODEL_FIELD_TO_TASK, validate_model_task_assignment
29from lilbee.providers.local_servers import canonical_local_ref, local_server_for_label
30from lilbee.providers.model_ref import format_remote_ref, parse_model_ref
31from lilbee.providers.sdk_backend import PROVIDER_KEYS, get_provider_api_key
32from lilbee.runtime.cancellation import TaskCancelledError
33from lilbee.runtime.hardware import (
34 FitLevel,
35 SizeVariantInfo,
36 available_memory_for_fit,
37 compute_fit,
38 family_size_variants,
39)
40from lilbee.runtime.progress import SseEvent
41from lilbee.server.handlers.sse import SseStream, sse_error, sse_event
42from lilbee.server.models import (
43 CatalogEntryResponse,
44 ExternalModelsResponse,
45 InstalledModelEntry,
46 ModelsCatalogResponse,
47 ModelsDeleteResponse,
48 ModelsInstalledResponse,
49 ModelsShowResponse,
50 SetModelResponse,
51)
53if TYPE_CHECKING:
54 from lilbee.catalog import CatalogModel
55 from lilbee.catalog.formatting import EnrichedModel
57log = logging.getLogger(__name__)
60class ModelCatalogEntry(BaseModel):
61 """A single model in the catalog."""
63 name: str
64 size_gb: float
65 min_ram_gb: float
66 description: str
67 installed: bool
70class ModelCatalogSection(BaseModel):
71 """A single-role catalog section with active model and installed list."""
73 active: str
74 catalog: list[ModelCatalogEntry]
75 installed: list[str]
78class ModelsResponse(BaseModel):
79 """Response for GET /api/models: one catalog section per role."""
81 chat: ModelCatalogSection
82 embedding: ModelCatalogSection
83 vision: ModelCatalogSection
84 reranker: ModelCatalogSection
87# ``ModelTask.RERANK.value`` is ``"rerank"`` but the route is ``/api/models/reranker``,
88# so this mapping is needed to build correct redirect URLs in 422 responses.
89TASK_ENDPOINT_PATH: dict[ModelTask, str] = {
90 ModelTask.CHAT: "chat",
91 ModelTask.EMBEDDING: "embedding",
92 ModelTask.VISION: "vision",
93 ModelTask.RERANK: "reranker",
94}
97def format_task_mismatch(ref: str, entry_task: ModelTask, expected_task: ModelTask) -> str:
98 """Build the 422 body when a role slot is assigned a model of the wrong task."""
99 endpoint = TASK_ENDPOINT_PATH[entry_task]
100 return (
101 f"Model '{ref}' is a {entry_task} model, not {expected_task}. "
102 f"Set it via PUT /api/models/{endpoint} instead."
103 )
106def _catalog_section(
107 featured: tuple[CatalogModel, ...],
108 active: str,
109 installed: set[str],
110) -> ModelCatalogSection:
111 """Build a ModelCatalogSection from a featured-catalog tuple.
113 A featured row is "installed" when at least one quant of its
114 ``hf_repo`` has a manifest. Installed refs are full
115 ``hf_repo/filename`` strings, so the membership test compares the
116 leading ``hf_repo`` segment. Bare ``hf_repo`` entries are accepted
117 too (e.g. older clients that report just the repo).
118 """
119 installed_repos = {hf_repo_from_ref(ref) for ref in installed}
120 return ModelCatalogSection(
121 active=active,
122 catalog=[
123 ModelCatalogEntry(
124 name=m.display_name,
125 size_gb=m.size_gb,
126 min_ram_gb=m.min_ram_gb,
127 description=m.description,
128 installed=m.hf_repo in installed_repos,
129 )
130 for m in featured
131 ],
132 installed=sorted(installed),
133 )
136async def list_models() -> ModelsResponse:
137 """Return per-role catalogs (chat, embedding, vision, reranker) with active selections.
139 Uses the unfiltered installed set so a single ref lights up in every
140 catalog section it legitimately matches.
141 """
142 # list_installed walks the model filesystem; offload it like list_external_models.
143 installed = set(await asyncio.to_thread(get_services().model_manager.list_installed))
144 # Resolving the picks hits HuggingFace on the first call of the process, so
145 # offload it too rather than stalling the event loop. The four picks_for
146 # reads below are served from the memo this fills.
147 await asyncio.to_thread(get_picks)
149 return ModelsResponse(
150 chat=_catalog_section(picks_for(ModelTask.CHAT), cfg.chat_model, installed),
151 embedding=_catalog_section(picks_for(ModelTask.EMBEDDING), cfg.embedding_model, installed),
152 vision=_catalog_section(picks_for(ModelTask.VISION), cfg.vision_model, installed),
153 reranker=_catalog_section(picks_for(ModelTask.RERANK), cfg.reranker_model, installed),
154 )
157async def _set_model(
158 field: Literal["chat_model", "embedding_model", "vision_model", "reranker_model"],
159 model: str,
160) -> SetModelResponse:
161 """Persist a model field through the shared write boundary."""
162 apply_settings_update({field: model})
163 return SetModelResponse(model=model)
166def _resolve_via_available_repo(model: str, available: set[str]) -> str | None:
167 """Resolve a bare ``hf_repo`` to whichever quant of it *available* lists.
169 Sorted scan so the pick is deterministic when several quants are installed.
170 """
171 if not is_bare_hf_repo(model):
172 return None
173 return next((ref for ref in sorted(available) if ref.startswith(f"{model}/")), None)
176def _resolve_via_parse(model: str, available: set[str]) -> str | None:
177 """Resolve a provider-prefixed ref against *available*.
179 The backend lists hosted models under bare names while selections carry the
180 routing prefix. When the bare name is visible, return the prefixed ref so it
181 keeps provider routing instead of falling through to the native check.
182 """
183 try:
184 parsed = parse_model_ref(model)
185 except ValueError:
186 return None
187 return model if parsed.name in available else None
190def _resolve_via_provider_key(model: str) -> str | None:
191 """Accept an API-provider-prefixed ref when that provider's key is configured.
193 Frontier models surface through ``discover_api_models``, not the default
194 ``list_models()``, so they never appear in *available*. With the key set,
195 litellm routes the ref (and validates the model name at call time).
196 """
197 try:
198 parsed = parse_model_ref(model)
199 except ValueError:
200 return None
201 if parsed.is_api and get_provider_api_key(parsed.provider) is not None:
202 return model
203 return None
206def _require_model_available(model: str) -> str:
207 """Return the installed-and-routable form of *model*, or raise."""
208 not_available = ValueError(
209 f"Model '{model}' is not available. Pull it first or check the name."
210 )
211 if not model:
212 raise not_available
213 available = set(get_services().provider.list_models())
214 if model in available:
215 return model
216 hit = (
217 _resolve_via_available_repo(model, available)
218 or _resolve_via_parse(model, available)
219 or _resolve_via_provider_key(model)
220 )
221 if hit is None:
222 raise not_available
223 return hit
226def _build_task_to_field() -> dict[ModelTask, str]:
227 """Invert ``MODEL_FIELD_TO_TASK`` so the two maps stay in sync."""
228 return {ModelTask(task): field for field, task in MODEL_FIELD_TO_TASK.items()}
231_TASK_TO_FIELD: dict[ModelTask, str] = _build_task_to_field()
234def _require_model_for_task(model: str, expected: ModelTask, *, allow_empty: bool = False) -> str:
235 """Validate *model* is installed locally AND passes the catalog task check.
237 Empty string unsets the role when *allow_empty* is True. Catalog +
238 task validation delegates to ``validate_model_task_assignment`` so
239 the handler and config paths share a single implementation.
240 """
241 if allow_empty and not model.strip():
242 return ""
243 normalized = _require_model_available(model)
244 return validate_model_task_assignment(_TASK_TO_FIELD[expected], normalized, allow_bypass=False)
247async def set_chat_model(model: str) -> SetModelResponse:
248 """Switch active chat model. Validates installation and catalog task."""
249 normalized = await asyncio.to_thread(_require_model_for_task, model, ModelTask.CHAT)
250 return await _set_model("chat_model", normalized)
253async def set_embedding_model(model: str) -> SetModelResponse:
254 """Switch embedding model. Validates installation and catalog task.
256 Returns ``reindex_required=True`` when the new model differs from the
257 embedding model that built the persisted vector store. The caller is
258 expected to trigger a rebuild (``lilbee rebuild`` or ``POST /api/sync``
259 with ``force_rebuild=true``). Search and ingest will refuse to operate
260 until that happens. The settings boundary pins legacy store meta to
261 the OLD ref before the write and computes ``reindex_required`` after.
262 """
263 normalized = await asyncio.to_thread(_require_model_for_task, model, ModelTask.EMBEDDING)
264 result = apply_settings_update({"embedding_model": normalized})
265 return SetModelResponse(model=normalized, reindex_required=result.reindex_required)
268async def set_vision_model(model: str) -> SetModelResponse:
269 """Switch vision OCR model. Empty string unsets it (vision OCR disabled)."""
270 normalized = await asyncio.to_thread(
271 _require_model_for_task, model, ModelTask.VISION, allow_empty=True
272 )
273 return await _set_model("vision_model", normalized)
276async def set_reranker_model(model: str) -> SetModelResponse:
277 """Switch reranker model. Empty string unsets it (reranking disabled)."""
278 normalized = await asyncio.to_thread(
279 _require_model_for_task, model, ModelTask.RERANK, allow_empty=True
280 )
281 return await _set_model("reranker_model", normalized)
284async def models_show(model: str) -> ModelsShowResponse:
285 """Return model metadata/parameters. Returns empty model if unavailable."""
286 provider = get_services().provider
287 # show_model dispatches to the SDK backend, which does a network call to the
288 # local server; offload so a hung backend doesn't block the event loop.
289 result = await asyncio.to_thread(provider.show_model, model)
290 return ModelsShowResponse(**(result or {}))
293def _parse_source(source: str) -> ModelSource:
294 """Convert a source string to ModelSource enum."""
295 return ModelSource(source)
298_BYTES_PER_GB = 1024**3
301def _row_fit(enriched: EnrichedModel, available_bytes: int | None) -> FitLevel | None:
302 """Fit level for *enriched*, or None when host memory or row size can't be measured."""
303 if available_bytes is None:
304 return None
305 if enriched.source != ModelSource.NATIVE.value:
306 return None
307 if enriched.size_gb <= 0:
308 return None
309 return compute_fit(int(enriched.size_gb * _BYTES_PER_GB), available_bytes).level
312def _families_by_repo() -> dict[str, ModelFamily]:
313 """Index featured ModelFamilies by every variant's ``hf_repo`` for size-variant lookup."""
314 index: dict[str, ModelFamily] = {}
315 for family in get_families():
316 for variant in family.variants:
317 index[variant.hf_repo] = family
318 return index
321def _row_size_variants(
322 enriched: EnrichedModel, families_by_repo: dict[str, ModelFamily]
323) -> list[SizeVariantInfo]:
324 """Size-variant strip for *enriched*; empty when the row isn't part of a family."""
325 family = families_by_repo.get(enriched.hf_repo)
326 if family is None:
327 return []
328 return family_size_variants(family)
331def _build_catalog_entry(
332 enriched: EnrichedModel,
333 *,
334 available_bytes: int | None,
335 families_by_repo: dict[str, ModelFamily],
336) -> CatalogEntryResponse:
337 """Translate one enriched catalog model into its HTTP response row."""
338 return CatalogEntryResponse(
339 hf_repo=enriched.hf_repo,
340 gguf_filename=enriched.gguf_filename,
341 task=enriched.task,
342 display_name=enriched.display_name,
343 param_count=enriched.param_count,
344 size_gb=enriched.size_gb,
345 min_ram_gb=enriched.min_ram_gb,
346 description=enriched.description,
347 quality_tier=enriched.quality_tier,
348 featured=enriched.featured,
349 downloads=enriched.downloads,
350 installed=enriched.installed,
351 source=enriched.source,
352 fit=_row_fit(enriched, available_bytes),
353 size_variants=_row_size_variants(enriched, families_by_repo),
354 architecture=enriched.architecture,
355 compat=enriched.compat,
356 )
359def _hosted_entry(rm: RemoteModel, source: ModelSource) -> CatalogEntryResponse:
360 """Build a selectable, no-download catalog row for a discovered hosted model."""
361 return CatalogEntryResponse(
362 hf_repo=format_remote_ref(rm.name, rm.provider),
363 gguf_filename="",
364 task=rm.task,
365 display_name=rm.name,
366 param_count=rm.parameter_size,
367 size_gb=0,
368 min_ram_gb=0,
369 description="",
370 quality_tier="",
371 featured=False,
372 downloads=0,
373 installed=True,
374 source=source,
375 fit=None,
376 size_variants=[],
377 provider=rm.provider,
378 key_status=KeyStatus.READY if source is ModelSource.FRONTIER else None,
379 )
382_HOSTED_MODELS_TTL = 60
385# Single-entry TTL caches keyed on the config tuple that produced the value:
386# maxsize=1 means a lookup under a new key evicts the old one.
387_hosted_cache: TTLCache[str, list[CatalogEntryResponse]] = TTLCache(
388 maxsize=1, ttl=_HOSTED_MODELS_TTL
389)
392def _discover_hosted_sync() -> list[CatalogEntryResponse]:
393 """All hosted rows (frontier + the configured local server), unfiltered.
395 Blocking; call via to_thread. Local rows take the detected server's source
396 (Ollama or LM Studio). Both discovery calls fail soft when no keys are set
397 or the endpoint is unreachable, so the catalog degrades to native-only.
398 """
399 rows: list[CatalogEntryResponse] = []
400 for models in discover_api_models().values():
401 rows.extend(_hosted_entry(rm, ModelSource.FRONTIER) for rm in models)
402 for rm in classify_all_remote_models():
403 spec = local_server_for_label(rm.provider)
404 source = ModelSource(spec.key) if spec is not None else ModelSource.REMOTE
405 rows.append(_hosted_entry(rm, source))
406 return rows
409def _hosted_cache_key() -> str:
410 """Cache key over the inputs that change discovery output.
412 Enumerates configured provider-key fields generically from
413 ``PROVIDER_KEYS`` so adding a provider does not silently reuse a
414 stale cache entry.
415 """
416 keys = ":".join(getattr(cfg, field) or "" for _, field, *_ in PROVIDER_KEYS)
417 return f"{cfg.ollama_base_url}:{cfg.lm_studio_base_url}:{keys}"
420async def _collect_hosted_entries(
421 *, task: ModelTask | None, search: str
422) -> list[CatalogEntryResponse]:
423 """Hosted catalog rows filtered by task/search, off the event loop + TTL-cached."""
424 key = _hosted_cache_key()
425 rows = _hosted_cache.get(key)
426 if rows is None:
427 rows = await asyncio.to_thread(_discover_hosted_sync)
428 _hosted_cache[key] = rows
429 if task is not None:
430 rows = [r for r in rows if r.task == task]
431 if search:
432 needle = search.lower()
433 rows = [r for r in rows if needle in r.display_name.lower()]
434 return rows
437async def models_catalog(
438 task: str | None = None,
439 search: str = "",
440 size: str | None = None,
441 installed: bool | None = None,
442 featured: bool | None = None,
443 sort: str = "featured",
444 limit: int = 20,
445 offset: int = 0,
446) -> ModelsCatalogResponse:
447 """Return paginated model catalog with installed status."""
448 # Validate every closed-set param at the HTTP boundary instead of
449 # letting unknown values silently short-circuit the filter inside.
450 parsed_task = ModelTask(task) if task else None
451 parsed_size = CatalogSize(size) if size else None
452 parsed_sort = CatalogSort(sort)
453 # get_catalog resolves the picks, which is HTTP on the first call.
454 result = await asyncio.to_thread(
455 get_catalog,
456 task=parsed_task,
457 search=search,
458 size=parsed_size,
459 installed=installed,
460 featured=featured,
461 sort=parsed_sort,
462 limit=limit,
463 offset=offset,
464 model_manager=get_services().model_manager,
465 )
467 registry = get_services().registry
468 installed_refs = {m.ref for m in registry.list_installed()}
469 enriched = enrich_catalog(result, installed_refs)
471 available_bytes = available_memory_for_fit()
472 families_by_repo = _families_by_repo()
474 native_rows = [
475 _build_catalog_entry(e, available_bytes=available_bytes, families_by_repo=families_by_repo)
476 for e in enriched
477 ]
478 # Hosted rows (frontier + ollama) are selectable and download-free, shown
479 # on the first page only and skipped for featured-only / installed=False.
480 # They stay out of ``total``: as an unpaginated overlay they made page 1
481 # report a larger total than page 2 of the same listing.
482 hosted_rows: list[CatalogEntryResponse] = []
483 if offset == 0 and not featured and installed is not False:
484 hosted_rows = await _collect_hosted_entries(task=parsed_task, search=search)
486 return ModelsCatalogResponse(
487 total=result.total,
488 limit=result.limit,
489 offset=result.offset,
490 has_more=result.has_more,
491 models=hosted_rows + native_rows,
492 )
495def _installed_entries_sync() -> list[InstalledModelEntry]:
496 """Blocking body of :func:`models_installed`: list installs and their sources."""
497 manager = get_services().model_manager
498 entries = []
499 for name in manager.list_installed():
500 source = manager.get_source(name) or ModelSource.REMOTE
501 entries.append(
502 InstalledModelEntry(name=canonical_local_ref(name, source.value), source=source)
503 )
504 return entries
507async def models_installed() -> ModelsInstalledResponse:
508 """Return installed models with their granular source and canonical ref."""
509 # list_installed walks the model filesystem and, on TTL expiry, queries the
510 # configured local servers over HTTP; offload it like list_models does.
511 entries = await asyncio.to_thread(_installed_entries_sync)
512 return ModelsInstalledResponse(models=entries)
515async def enforce_pull_arch_compat(
516 model: str, *, source: str = "native", allow_unsupported: bool = False
517) -> None:
518 """Raise HTTP 409 for an unsupported architecture before the pull stream opens.
520 The route must await this BEFORE returning ``Stream(models_pull(...))``: a raise
521 inside the ``models_pull`` async generator fires only on first iteration, after
522 Litestar has already flushed the 200 SSE headers, so it can no longer set the
523 status. (``manager.pull`` re-enforces compatibility during the pull itself.)
524 """
525 from litestar.exceptions import HTTPException
527 from lilbee.catalog.compat import SUPPORTED_ARCHS, UnsupportedArchError
529 if _parse_source(source) is not ModelSource.NATIVE or allow_unsupported:
530 return
531 manager = get_services().model_manager
532 try:
533 await asyncio.to_thread(manager.enforce_arch_compat, model)
534 except UnsupportedArchError as exc:
535 raise HTTPException(
536 status_code=409,
537 detail="unsupported_arch",
538 extra={
539 "code": "unsupported_arch",
540 "arch": exc.architecture,
541 "ref": exc.ref,
542 "supported_examples": sorted(SUPPORTED_ARCHS)[:5],
543 "total_supported": len(SUPPORTED_ARCHS),
544 },
545 ) from exc
548async def models_pull(
549 model: str, *, source: str = "native", allow_unsupported: bool = False
550) -> AsyncGenerator[str, None]:
551 """Yield SSE progress events while pulling a model in real time.
552 Sets a cancel event on client disconnect so the pull stops.
554 Architecture compatibility is enforced by the route via
555 :func:`enforce_pull_arch_compat` before this stream opens; ``manager.pull``
556 re-enforces it during the pull.
557 """
558 manager = get_services().model_manager
559 src = _parse_source(source)
561 sse = SseStream()
563 def _pull_blocking() -> None:
564 def _on_bytes(downloaded: int, total: int) -> None:
565 # Raise, not return: the pull runs in a worker thread asyncio
566 # cannot interrupt, so returning left a multi-GB download running
567 # for a client that had gone.
568 if sse.cancel.is_set():
569 raise TaskCancelledError
570 payload = sse_event(SseEvent.PROGRESS, {"current": downloaded, "total": total})
571 sse.loop.call_soon_threadsafe(sse.queue.put_event_nowait, payload, SseEvent.PROGRESS)
573 try:
574 manager.pull(
575 model,
576 src,
577 on_bytes=_on_bytes,
578 allow_unsupported=allow_unsupported,
579 cancel=sse.cancel,
580 )
581 except TaskCancelledError:
582 log.info("Model pull for %s aborted: client disconnected", model)
583 except Exception as exc:
584 sse.loop.call_soon_threadsafe(sse.queue.put_nowait, sse_error(str(exc)))
585 finally:
586 sse.loop.call_soon_threadsafe(sse.queue.put_nowait, None)
588 task = asyncio.ensure_future(asyncio.to_thread(_pull_blocking))
589 try:
590 async for event in sse.drain(task, "Model pull stream"):
591 yield event
592 finally:
593 # Closing this generator means the client is gone. Set it here, not in
594 # drain's cleanup, which runs only once that generator is collected;
595 # until then the worker thread keeps downloading.
596 sse.cancel.set()
599async def models_delete(model: str, *, source: str = "native") -> ModelsDeleteResponse:
600 """Delete a model. Returns deletion status, model name, and freed space.
602 lilbee removes only native models it downloaded; removing a read-only
603 local-server model (Ollama, LM Studio) is refused with a 409.
604 """
605 from litestar.exceptions import HTTPException
607 from lilbee.app.models import remove_model_data
609 src = _parse_source(source)
610 try:
611 # Delegate to the shared remove path so REST reports the same freed size
612 # (full multi-shard total) as the CLI and MCP, not a hardcoded 0.
613 result = remove_model_data(model, src)
614 except ValueError as exc:
615 raise HTTPException(status_code=409, detail=str(exc)) from exc
616 return ModelsDeleteResponse(
617 deleted=result.deleted, model=result.model, freed_gb=result.freed_gb
618 )
621_EXTERNAL_MODELS_TTL = 60
624_external_cache: TTLCache[str, ExternalModelsResponse] = TTLCache(
625 maxsize=1, ttl=_EXTERNAL_MODELS_TTL
626)
629async def list_external_models() -> ExternalModelsResponse:
630 """Query the provider for available models via its list_models() API."""
631 key = f"{cfg.ollama_base_url}:{cfg.lm_studio_base_url}:{cfg.llm_api_key or ''}"
632 # ``is not None``, not truthiness: an empty model list is a real answer.
633 cached = _external_cache.get(key)
634 if cached is not None:
635 return cached
637 try:
638 models = await asyncio.to_thread(get_services().provider.list_models)
639 result = ExternalModelsResponse(models=models)
640 _external_cache[key] = result
641 return result
642 except Exception as exc:
643 log.warning("Failed to list external models: %s", exc)
644 return ExternalModelsResponse(models=[], error=str(exc))