Coverage for src/lilbee/server/handlers/models.py: 100%
239 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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.models import list_models_data
14from lilbee.app.services import get_services
15from lilbee.app.settings import apply_settings_update
16from lilbee.catalog import (
17 CatalogResult,
18 ModelFamily,
19 enrich_catalog,
20 get_catalog,
21 get_families,
22 get_picks,
23 page_window,
24 picks_for,
25)
26from lilbee.catalog.refs import hf_repo_from_ref, is_bare_hf_repo
27from lilbee.catalog.types import CatalogSize, CatalogSort, KeyStatus, ModelSource, ModelTask
28from lilbee.core.config import cfg
29from lilbee.modelhub.model_manager import classify_all_remote_models, discover_api_models
30from lilbee.modelhub.model_manager.types import RemoteModel
31from lilbee.modelhub.role_validator import MODEL_FIELD_TO_TASK, validate_model_task_assignment
32from lilbee.providers.key_check import provider_key_status
33from lilbee.providers.local_servers import local_server_for_label
34from lilbee.providers.model_ref import format_remote_ref, parse_model_ref
35from lilbee.providers.sdk_backend import PROVIDER_KEYS
36from lilbee.runtime.cancellation import TaskCancelledError
37from lilbee.runtime.hardware import (
38 FitLevel,
39 SizeVariantInfo,
40 available_memory_for_fit,
41 family_size_variants,
42 fit_for_size,
43 make_fit_filter,
44)
45from lilbee.runtime.progress import SseEvent
46from lilbee.server.handlers.sse import SseStream, sse_error, sse_event
47from lilbee.server.models import (
48 CatalogEntryResponse,
49 ExternalModelsResponse,
50 ModelsCatalogResponse,
51 ModelsDeleteResponse,
52 ModelsInstalledResponse,
53 ModelsShowResponse,
54 SetModelResponse,
55)
57if TYPE_CHECKING:
58 from lilbee.catalog import CatalogModel
59 from lilbee.catalog.formatting import EnrichedModel
61log = logging.getLogger(__name__)
64class ModelCatalogEntry(BaseModel):
65 """A single model in the catalog."""
67 name: str
68 size_gb: float
69 min_ram_gb: float
70 description: str
71 installed: bool
74class ModelCatalogSection(BaseModel):
75 """A single-role catalog section with active model and installed list."""
77 active: str
78 catalog: list[ModelCatalogEntry]
79 installed: list[str]
82class ModelsResponse(BaseModel):
83 """Response for GET /api/models: one catalog section per role."""
85 chat: ModelCatalogSection
86 embedding: ModelCatalogSection
87 vision: ModelCatalogSection
88 reranker: ModelCatalogSection
91# ``ModelTask.RERANK.value`` is ``"rerank"`` but the route is ``/api/models/reranker``,
92# so this mapping is needed to build correct redirect URLs in 422 responses.
93TASK_ENDPOINT_PATH: dict[ModelTask, str] = {
94 ModelTask.CHAT: "chat",
95 ModelTask.EMBEDDING: "embedding",
96 ModelTask.VISION: "vision",
97 ModelTask.RERANK: "reranker",
98}
101def format_task_mismatch(ref: str, entry_task: ModelTask, expected_task: ModelTask) -> str:
102 """Build the 422 body when a role slot is assigned a model of the wrong task."""
103 endpoint = TASK_ENDPOINT_PATH[entry_task]
104 return (
105 f"Model '{ref}' is a {entry_task} model, not {expected_task}. "
106 f"Set it via PUT /api/models/{endpoint} instead."
107 )
110def _catalog_section(
111 featured: tuple[CatalogModel, ...],
112 active: str,
113 installed: set[str],
114) -> ModelCatalogSection:
115 """Build a ModelCatalogSection from a featured-catalog tuple.
117 A featured row is "installed" when at least one quant of its
118 ``hf_repo`` has a manifest. Installed refs are full
119 ``hf_repo/filename`` strings, so the membership test compares the
120 leading ``hf_repo`` segment. Bare ``hf_repo`` entries are accepted
121 too (e.g. older clients that report just the repo).
122 """
123 installed_repos = {hf_repo_from_ref(ref) for ref in installed}
124 return ModelCatalogSection(
125 active=active,
126 catalog=[
127 ModelCatalogEntry(
128 name=m.display_name,
129 size_gb=m.size_gb,
130 min_ram_gb=m.min_ram_gb,
131 description=m.description,
132 installed=m.hf_repo in installed_repos,
133 )
134 for m in featured
135 ],
136 installed=sorted(installed),
137 )
140async def list_models() -> ModelsResponse:
141 """Return per-role catalogs (chat, embedding, vision, reranker) with active selections.
143 Uses the unfiltered installed set so a single ref lights up in every
144 catalog section it legitimately matches.
145 """
146 # list_installed walks the model filesystem; offload it like list_external_models.
147 installed = set(await asyncio.to_thread(get_services().model_manager.list_installed))
148 # Resolving the picks hits HuggingFace on the first call of the process, so
149 # offload it too rather than stalling the event loop. The four picks_for
150 # reads below are served from the memo this fills.
151 await asyncio.to_thread(get_picks)
153 return ModelsResponse(
154 chat=_catalog_section(picks_for(ModelTask.CHAT), cfg.chat_model, installed),
155 embedding=_catalog_section(picks_for(ModelTask.EMBEDDING), cfg.embedding_model, installed),
156 vision=_catalog_section(picks_for(ModelTask.VISION), cfg.vision_model, installed),
157 reranker=_catalog_section(picks_for(ModelTask.RERANK), cfg.reranker_model, installed),
158 )
161async def _set_model(
162 field: Literal["chat_model", "embedding_model", "vision_model", "reranker_model"],
163 model: str,
164) -> SetModelResponse:
165 """Persist a model field through the shared write boundary."""
166 result = await asyncio.to_thread(apply_settings_update, {field: model})
167 return SetModelResponse(model=model, warnings=list(result.warnings))
170def _resolve_via_available_repo(model: str, available: set[str]) -> str | None:
171 """Resolve a bare ``hf_repo`` to whichever quant of it *available* lists.
173 Sorted scan so the pick is deterministic when several quants are installed.
174 """
175 if not is_bare_hf_repo(model):
176 return None
177 return next((ref for ref in sorted(available) if ref.startswith(f"{model}/")), None)
180def _resolve_via_parse(model: str, available: set[str]) -> str | None:
181 """Resolve a provider-prefixed ref against *available*.
183 The backend lists hosted models under bare names while selections carry the
184 routing prefix. When the bare name is visible, return the prefixed ref so it
185 keeps provider routing instead of falling through to the native check.
186 """
187 try:
188 parsed = parse_model_ref(model)
189 except ValueError:
190 return None
191 return model if parsed.name in available else None
194def _resolve_via_provider_key(model: str) -> str | None:
195 """Accept an API-provider-prefixed ref when that provider's key is set and not rejected.
197 Frontier models surface through ``discover_api_models``, not the default
198 ``list_models()``, so they never appear in *available*. With a usable key,
199 the SDK routes the ref (and validates the model name at call time).
200 """
201 try:
202 parsed = parse_model_ref(model)
203 except ValueError:
204 return None
205 if parsed.is_api and provider_key_status(parsed.provider) is KeyStatus.READY:
206 return model
207 return None
210def _require_model_available(model: str) -> str:
211 """Return the installed-and-routable form of *model*, or raise."""
212 not_available = ValueError(
213 f"Model '{model}' is not available. Pull it first or check the name."
214 )
215 if not model:
216 raise not_available
217 available = set(get_services().provider.list_models())
218 if model in available:
219 return model
220 hit = (
221 _resolve_via_available_repo(model, available)
222 or _resolve_via_parse(model, available)
223 or _resolve_via_provider_key(model)
224 )
225 if hit is None:
226 raise not_available
227 return hit
230def _build_task_to_field() -> dict[ModelTask, str]:
231 """Invert ``MODEL_FIELD_TO_TASK`` so the two maps stay in sync."""
232 return {ModelTask(task): field for field, task in MODEL_FIELD_TO_TASK.items()}
235_TASK_TO_FIELD: dict[ModelTask, str] = _build_task_to_field()
238def _require_model_for_task(model: str, expected: ModelTask, *, allow_empty: bool = False) -> str:
239 """Validate *model* is installed locally AND passes the catalog task check.
241 Empty string unsets the role when *allow_empty* is True. Catalog +
242 task validation delegates to ``validate_model_task_assignment`` so
243 the handler and config paths share a single implementation.
244 """
245 if allow_empty and not model.strip():
246 return ""
247 normalized = _require_model_available(model)
248 return validate_model_task_assignment(_TASK_TO_FIELD[expected], normalized, allow_bypass=False)
251async def set_chat_model(model: str) -> SetModelResponse:
252 """Switch active chat model. Validates installation and catalog task."""
253 normalized = await asyncio.to_thread(_require_model_for_task, model, ModelTask.CHAT)
254 return await _set_model("chat_model", normalized)
257async def set_embedding_model(model: str) -> SetModelResponse:
258 """Switch embedding model. Validates installation and catalog task.
260 Returns ``reindex_required=True`` when the new model differs from the
261 embedding model that built the persisted vector store. The caller is
262 expected to trigger a rebuild (``lilbee rebuild`` or ``POST /api/sync``
263 with ``force_rebuild=true``). Search and ingest will refuse to operate
264 until that happens. The settings boundary pins legacy store meta to
265 the OLD ref before the write and computes ``reindex_required`` after.
266 """
267 normalized = await asyncio.to_thread(_require_model_for_task, model, ModelTask.EMBEDDING)
268 result = await asyncio.to_thread(apply_settings_update, {"embedding_model": normalized})
269 return SetModelResponse(model=normalized, reindex_required=result.reindex_required)
272async def set_vision_model(model: str) -> SetModelResponse:
273 """Switch vision OCR model. Empty string unsets it (vision OCR disabled)."""
274 normalized = await asyncio.to_thread(
275 _require_model_for_task, model, ModelTask.VISION, allow_empty=True
276 )
277 return await _set_model("vision_model", normalized)
280async def set_reranker_model(model: str) -> SetModelResponse:
281 """Switch reranker model. Empty string unsets it (reranking disabled)."""
282 normalized = await asyncio.to_thread(
283 _require_model_for_task, model, ModelTask.RERANK, allow_empty=True
284 )
285 return await _set_model("reranker_model", normalized)
288async def models_show(model: str) -> ModelsShowResponse:
289 """Return model metadata/parameters. Returns empty model if unavailable."""
290 provider = get_services().provider
291 # show_model dispatches to the SDK backend, which does a network call to the
292 # local server; offload so a hung backend doesn't block the event loop.
293 result = await asyncio.to_thread(provider.show_model, model)
294 return ModelsShowResponse(**(result or {}))
297def _parse_source(source: str) -> ModelSource:
298 """Convert a source string to ModelSource enum."""
299 return ModelSource(source)
302def _row_fit(enriched: EnrichedModel, available_bytes: int | None) -> FitLevel | None:
303 """Fit level for *enriched*, or None when host memory or row size can't be measured."""
304 if enriched.source != ModelSource.NATIVE.value:
305 return None
306 return fit_for_size(enriched.size_gb, available_bytes)
309def _families_by_repo() -> dict[str, ModelFamily]:
310 """Index featured ModelFamilies by every variant's ``hf_repo`` for size-variant lookup."""
311 index: dict[str, ModelFamily] = {}
312 for family in get_families():
313 for variant in family.variants:
314 index[variant.hf_repo] = family
315 return index
318def _row_size_variants(
319 enriched: EnrichedModel, families_by_repo: dict[str, ModelFamily]
320) -> list[SizeVariantInfo]:
321 """Size-variant strip for *enriched*; empty when the row isn't part of a family."""
322 family = families_by_repo.get(enriched.hf_repo)
323 if family is None:
324 return []
325 return family_size_variants(family)
328def _build_catalog_entry(
329 enriched: EnrichedModel,
330 *,
331 available_bytes: int | None,
332 families_by_repo: dict[str, ModelFamily],
333) -> CatalogEntryResponse:
334 """Translate one enriched catalog model into its HTTP response row."""
335 return CatalogEntryResponse(
336 hf_repo=enriched.hf_repo,
337 gguf_filename=enriched.gguf_filename,
338 task=enriched.task,
339 display_name=enriched.display_name,
340 param_count=enriched.param_count,
341 size_gb=enriched.size_gb,
342 min_ram_gb=enriched.min_ram_gb,
343 description=enriched.description,
344 quality_tier=enriched.quality_tier,
345 featured=enriched.featured,
346 downloads=enriched.downloads,
347 installed=enriched.installed,
348 source=enriched.source,
349 fit=_row_fit(enriched, available_bytes),
350 size_variants=_row_size_variants(enriched, families_by_repo),
351 architecture=enriched.architecture,
352 compat=enriched.compat,
353 safety_stripped=enriched.safety_stripped,
354 )
357def _hosted_entry(rm: RemoteModel, source: ModelSource) -> CatalogEntryResponse:
358 """Build a selectable, no-download catalog row for a discovered hosted model."""
359 return CatalogEntryResponse(
360 hf_repo=format_remote_ref(rm.name, rm.provider),
361 gguf_filename="",
362 task=rm.task,
363 display_name=rm.name,
364 param_count=rm.parameter_size,
365 size_gb=0,
366 min_ram_gb=0,
367 description="",
368 quality_tier="",
369 featured=False,
370 downloads=0,
371 installed=True,
372 source=source,
373 fit=None,
374 size_variants=[],
375 provider=rm.provider,
376 key_status=KeyStatus.READY if source is ModelSource.FRONTIER else None,
377 )
380_HOSTED_MODELS_TTL = 60
383# Single-entry TTL caches keyed on the config tuple that produced the value:
384# maxsize=1 means a lookup under a new key evicts the old one.
385_hosted_cache: TTLCache[str, list[CatalogEntryResponse]] = TTLCache(
386 maxsize=1, ttl=_HOSTED_MODELS_TTL
387)
390def _discover_hosted_sync() -> list[CatalogEntryResponse]:
391 """All hosted rows (frontier + the configured local server), unfiltered.
393 Blocking; call via to_thread. Local rows take the detected server's source
394 (Ollama or LM Studio). Both discovery calls fail soft when no keys are set
395 or the endpoint is unreachable, so the catalog degrades to native-only.
396 """
397 rows: list[CatalogEntryResponse] = []
398 for models in discover_api_models().values():
399 rows.extend(_hosted_entry(rm, ModelSource.FRONTIER) for rm in models)
400 for rm in classify_all_remote_models():
401 spec = local_server_for_label(rm.provider)
402 source = ModelSource(spec.key) if spec is not None else ModelSource.REMOTE
403 rows.append(_hosted_entry(rm, source))
404 return rows
407def _hosted_cache_key() -> str:
408 """Cache key over the inputs that change discovery output.
410 Enumerates configured provider-key fields generically from
411 ``PROVIDER_KEYS`` so adding a provider does not silently reuse a
412 stale cache entry.
413 """
414 keys = ":".join(getattr(cfg, field) or "" for _, field, *_ in PROVIDER_KEYS)
415 return f"{cfg.ollama_base_url}:{cfg.lm_studio_base_url}:{keys}"
418async def _collect_hosted_entries(
419 *, task: ModelTask | None, search: str, source: ModelSource | None = None
420) -> list[CatalogEntryResponse]:
421 """Hosted catalog rows filtered by task/search/source, off the event loop + TTL-cached."""
422 key = _hosted_cache_key()
423 rows = _hosted_cache.get(key)
424 if rows is None:
425 rows = await asyncio.to_thread(_discover_hosted_sync)
426 _hosted_cache[key] = rows
427 if task is not None:
428 rows = [r for r in rows if r.task == task]
429 if source is not None:
430 rows = [r for r in rows if r.source == source]
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 max_fit: str | None = None,
444 sort: str = "featured",
445 limit: int = 20,
446 offset: int = 0,
447 source: str | None = None,
448) -> ModelsCatalogResponse:
449 """One page of the catalog rows that pass every filter, hosted rows first."""
450 # Validate every closed-set param at the HTTP boundary instead of
451 # letting unknown values silently short-circuit the filter inside.
452 parsed_task = ModelTask(task) if task else None
453 parsed_size = CatalogSize(size) if size else None
454 parsed_sort = CatalogSort(sort)
455 parsed_max_fit = FitLevel(max_fit) if max_fit else None
456 parsed_source = ModelSource.parse(source)
458 # Hosted rows (frontier + ollama) lead the listing; none for featured-only,
459 # installed=False, or source=native.
460 hosted: list[CatalogEntryResponse] = []
461 if not featured and installed is not False and parsed_source is not ModelSource.NATIVE:
462 hosted = await _collect_hosted_entries(
463 task=parsed_task, search=search, source=parsed_source
464 )
465 window = page_window(len(hosted), offset, limit)
467 available_bytes = available_memory_for_fit()
468 # A listing filtered to one hosted source has no native part.
469 result = CatalogResult(total=0, limit=0, offset=0, models=[])
470 if parsed_source in (None, ModelSource.NATIVE):
471 # get_catalog resolves the picks, which is HTTP on the first call.
472 result = await asyncio.to_thread(
473 get_catalog,
474 task=parsed_task,
475 search=search,
476 size=parsed_size,
477 installed=installed,
478 featured=featured,
479 fit_filter=make_fit_filter(parsed_max_fit, available_bytes),
480 sort=parsed_sort,
481 limit=window.rest_limit,
482 offset=window.rest_offset,
483 model_manager=get_services().model_manager,
484 )
486 registry = get_services().registry
487 installed_refs = {m.ref for m in registry.list_installed()}
488 enriched = enrich_catalog(result, installed_refs)
490 families_by_repo = _families_by_repo()
492 native_rows = [
493 _build_catalog_entry(e, available_bytes=available_bytes, families_by_repo=families_by_repo)
494 for e in enriched
495 ]
496 has_more = len(hosted) > offset + limit or result.has_more
497 return ModelsCatalogResponse(
498 total=None if result.total is None else len(hosted) + result.total,
499 limit=limit,
500 offset=offset,
501 has_more=has_more,
502 next_offset=offset + limit if has_more else None,
503 truncated=result.truncated,
504 models=hosted[offset : offset + limit] + native_rows,
505 )
508async def models_installed(task: ModelTask | None = None) -> ModelsInstalledResponse:
509 """Installed models across every source, only those of *task* when given."""
510 # Walks the model filesystem and queries the local servers over HTTP.
511 result = await asyncio.to_thread(list_models_data, task=task)
512 return ModelsInstalledResponse(models=result.models)
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))