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