Coverage for src/lilbee/cli/tui/screens/catalog_utils.py: 100%
142 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Catalog data types, row builders, and formatting helpers.
3The catalog renders two distinct row shapes side by side: locally
4installed / installable GGUFs (``LocalCatalogRow``) and cloud chat
5models accessed through a provider's API (``FrontierCatalogRow``).
6They share enough surface area that grouping and search reuse the
7same helpers, but they carry different metadata and pull from
8different sources, so they're separate types under a sealed
9``CatalogRow`` union rather than a single optional-fields dataclass.
10"""
12from __future__ import annotations
14import re
15from collections.abc import Callable
16from dataclasses import dataclass, field
17from enum import StrEnum
18from typing import Any, Literal
20from lilbee.catalog import PARAM_COUNT_RE, CatalogModel, ModelFamily, ModelVariant, extract_quant
21from lilbee.catalog.types import KeyStatus, ModelCompat, ModelTask
22from lilbee.modelhub.model_manager import RemoteModel
23from lilbee.providers.model_ref import format_remote_ref
24from lilbee.runtime.hardware import FitChip
26# Backend label for local GGUF models. Renderers drop this pill (the backend is
27# implied for local models) and only show it for non-native SDK backends.
28NATIVE_BACKEND = "native"
31class CatalogRowKind(StrEnum):
32 """Discriminator for the sealed CatalogRow union."""
34 LOCAL = "local"
35 FRONTIER = "frontier"
38# Tab IDs for the 6-tab catalog shell. Discover is the curated landing,
39# the four task tabs each render a single per-task grid, Library is the
40# personal-encyclopedia view of installed local + activated cloud APIs.
41TAB_DISCOVER = "discover"
42TAB_CHAT = "chat"
43TAB_EMBED = "embed"
44TAB_VISION = "vision"
45TAB_RERANK = "rerank"
46TAB_LIBRARY = "library"
48# Order matters: numbered shortcuts 1-6 follow this sequence.
49ALL_TAB_IDS: tuple[str, ...] = (
50 TAB_DISCOVER,
51 TAB_CHAT,
52 TAB_EMBED,
53 TAB_VISION,
54 TAB_RERANK,
55 TAB_LIBRARY,
56)
57TASK_TAB_IDS: tuple[str, ...] = (TAB_CHAT, TAB_EMBED, TAB_VISION, TAB_RERANK)
59# Maps ModelTask -> the per-task tab id that renders its rows. Featured
60# items still appear pinned at the top of their task tab; cross-task
61# discovery happens on the Discover landing.
62TASK_TO_TAB_ID: dict[ModelTask, str] = {
63 ModelTask.CHAT: TAB_CHAT,
64 ModelTask.EMBEDDING: TAB_EMBED,
65 ModelTask.VISION: TAB_VISION,
66 ModelTask.RERANK: TAB_RERANK,
67}
68TAB_ID_TO_TASK: dict[str, ModelTask] = {v: k for k, v in TASK_TO_TAB_ID.items()}
71class SourceMode(StrEnum):
72 """Per-task tab filter for which row source backs the visible cards.
74 LOCAL hides frontier rows (the default; mirrors the legacy mega-grid
75 behavior). CLOUD shows only frontier rows for the active task. BOTH
76 unions them so users can compare a local Llama with a cloud Llama
77 side by side. Cycled via the ``c`` keybinding.
78 """
80 LOCAL = "local"
81 CLOUD = "cloud"
82 BOTH = "both"
85_SOURCE_MODE_CYCLE: tuple[SourceMode, ...] = (SourceMode.LOCAL, SourceMode.CLOUD, SourceMode.BOTH)
88def next_source_mode(current: SourceMode) -> SourceMode:
89 """Return the next SourceMode in the LOCAL -> CLOUD -> BOTH -> LOCAL cycle."""
90 idx = _SOURCE_MODE_CYCLE.index(current)
91 return _SOURCE_MODE_CYCLE[(idx + 1) % len(_SOURCE_MODE_CYCLE)]
94def task_to_tab_id(task: ModelTask | str) -> str:
95 """Return the per-task tab id for a ModelTask or its string value.
97 Accepts the string form because catalog rows carry ``task`` as a
98 raw string (matching how HF API and the row builders return it),
99 while the routing tables are keyed on the enum.
100 """
101 if isinstance(task, ModelTask):
102 return TASK_TO_TAB_ID[task]
103 try:
104 return TASK_TO_TAB_ID[ModelTask(task)]
105 except (KeyError, ValueError) as exc:
106 raise KeyError(f"unknown task: {task!r}") from exc
109# SI thresholds for short download counts ("12.3M" / "456K") and binary
110# thresholds for sizes ("4.2 GB" / "768 MB").
111_DOWNLOADS_PER_M = 1_000_000
112_DOWNLOADS_PER_K = 1_000
113_MB_PER_GB = 1024
116@dataclass(frozen=True)
117class SizeVariant:
118 """One size/quant variant of a model family for the family-as-card strip.
120 ``label`` renders inline on the card (e.g. "8B Q4_K_M"). ``ref`` is
121 the canonical pull target for this specific variant. ``fit`` is the
122 fit chip computed against the host's available memory; ``None``
123 when the hardware probe has not yet run.
124 """
126 label: str
127 quant: str
128 size_gb: float
129 ref: str
130 fit: FitChip | None = None
133@dataclass
134class LocalCatalogRow:
135 """A row in the catalog backed by a local GGUF (installable or installed).
137 ``name`` is the human-readable display label (e.g. "Qwen3 0.6B").
138 ``ref`` is the canonical identifier used for config persistence:
139 ``hf_repo`` for catalog rows, ``hf_repo/filename`` for installed
140 native models, and the provider's ref shape for remote/API rows.
141 ``size_variants`` carries every quant for a family-aggregated row,
142 so the card can render an inline chip strip and the detail drawer
143 can list all sizes. ``fit`` is the chip for the row's primary
144 variant.
145 """
147 name: str
148 task: str
149 params: str
150 size: str
151 quant: str
152 downloads: str
153 featured: bool
154 installed: bool
155 sort_downloads: int
156 sort_size: float
157 ref: str = ""
158 backend: str = ""
159 variant: ModelVariant | None = None
160 family: ModelFamily | None = None
161 catalog_model: CatalogModel | None = None
162 remote_model: RemoteModel | None = None
163 size_variants: list[SizeVariant] = field(default_factory=list)
164 fit: FitChip | None = None
165 compat: ModelCompat = ModelCompat.UNKNOWN
166 kind: Literal[CatalogRowKind.LOCAL] = CatalogRowKind.LOCAL
169@dataclass
170class FrontierCatalogRow:
171 """A row in the catalog backed by a cloud provider's chat API.
173 Frontier rows skip the local-model fields (size on disk, quant,
174 GGUF filename) because they don't apply: the model lives on the
175 provider's infrastructure.
176 """
178 name: str
179 ref: str
180 task: str
181 provider: str # Display label, e.g. "Gemini" / "OpenAI" / "Anthropic".
182 provider_id: str # Canonical id used for the API key field, e.g. "gemini".
183 key_status: KeyStatus
184 kind: Literal[CatalogRowKind.FRONTIER] = CatalogRowKind.FRONTIER
187# Sealed union discriminated on .kind. Pattern-match (or compare) on row.kind
188# to dispatch instead of isinstance, so adding a new row type is one place.
189CatalogRow = LocalCatalogRow | FrontierCatalogRow
192def parse_param_label(name: str) -> str:
193 """Extract parameter count label from model name (e.g. '8B', '0.6B')."""
194 from lilbee.catalog import PARAM_COUNT_RE
196 match = PARAM_COUNT_RE.search(name)
197 return match.group(1).upper() if match else "--"
200def _format_downloads(n: int) -> str:
201 if n >= _DOWNLOADS_PER_M:
202 return f"{n / _DOWNLOADS_PER_M:.1f}M"
203 if n >= _DOWNLOADS_PER_K:
204 return f"{n / _DOWNLOADS_PER_K:.0f}K"
205 return str(n)
208def _format_size_mb(size_mb: int) -> str:
209 """Format size in MB to a human-readable string."""
210 if size_mb == 0:
211 return "--"
212 if size_mb >= _MB_PER_GB:
213 return f"{size_mb / _MB_PER_GB:.1f} GB"
214 return f"{size_mb} MB"
217def format_size_gb(size_gb: float) -> str:
218 """Format size in GB to a human-readable string."""
219 if size_gb <= 0:
220 return "--"
221 return f"{size_gb:.1f} GB"
224def _is_param_count(label: str) -> bool:
225 """True when label looks like a parameter count (e.g. '8B', '0.6B')."""
226 return bool(PARAM_COUNT_RE.fullmatch(label))
229def family_to_size_variants(family: ModelFamily) -> list[SizeVariant]:
230 """Build the size-chip strip for a featured ModelFamily.
232 Variants are returned in increasing size order so the chip strip
233 reads compact-to-large left-to-right. ``fit`` is left ``None``;
234 the catalog screen fills it in once the hardware probe has run.
235 """
236 variants = sorted(family.variants, key=lambda v: v.size_mb)
237 return [
238 SizeVariant(
239 label=_size_variant_label(v),
240 quant=v.quant or "--",
241 size_gb=v.size_mb / 1024,
242 ref=v.hf_repo,
243 fit=None,
244 )
245 for v in variants
246 ]
249def _size_variant_label(v: ModelVariant) -> str:
250 """Render a compact label for a ModelVariant chip (e.g. '8B Q4_K_M')."""
251 pieces = [p for p in (v.param_count, v.quant) if p]
252 return " ".join(pieces) if pieces else "--"
255def variant_to_row(v: ModelVariant, f: ModelFamily, installed: bool) -> LocalCatalogRow:
256 """Convert a ModelVariant + family to a LocalCatalogRow."""
257 # Avoid duplicating the param count when the family name already ends with it.
258 if v.param_count and not f.name.endswith(v.param_count):
259 label = f"{f.name} {v.param_count}"
260 else:
261 label = f.name
262 params = v.param_count if _is_param_count(v.param_count) else "--"
263 return LocalCatalogRow(
264 name=label,
265 task=f.task,
266 params=params,
267 size=_format_size_mb(v.size_mb),
268 quant=v.quant or "--",
269 downloads="--",
270 featured=True,
271 installed=installed,
272 sort_downloads=0,
273 sort_size=v.size_mb / 1024,
274 ref=v.hf_repo,
275 backend=NATIVE_BACKEND,
276 variant=v,
277 family=f,
278 compat=v.compat,
279 )
282def catalog_to_row(m: CatalogModel, installed: bool) -> LocalCatalogRow:
283 """Convert a CatalogModel to a LocalCatalogRow."""
284 quant = extract_quant(m.gguf_filename)
285 return LocalCatalogRow(
286 name=m.display_name,
287 task=m.task,
288 params=parse_param_label(m.display_name),
289 size=format_size_gb(m.size_gb),
290 quant=quant or "--",
291 downloads=_format_downloads(m.downloads) if m.downloads > 0 else "--",
292 featured=m.featured,
293 installed=installed,
294 sort_downloads=m.downloads,
295 sort_size=m.size_gb,
296 ref=m.ref,
297 backend=NATIVE_BACKEND,
298 catalog_model=m,
299 # An installed model demonstrably runs, whatever the catalog probe said.
300 compat=ModelCompat.SUPPORTED if installed else m.compat,
301 )
304def remote_to_row(rm: RemoteModel) -> LocalCatalogRow:
305 """Convert a RemoteModel to a LocalCatalogRow.
307 ``ref`` is the canonical ``provider/name`` form so it round-trips
308 through ``Config.chat_model``'s validator without a per-call-site
309 fixup.
310 """
311 return LocalCatalogRow(
312 name=rm.name,
313 task=rm.task,
314 params=rm.parameter_size or "--",
315 size="--",
316 quant="--",
317 downloads="--",
318 featured=False,
319 installed=True,
320 sort_downloads=0,
321 sort_size=0.0,
322 ref=format_remote_ref(rm.name, rm.provider),
323 backend=rm.provider.lower(),
324 remote_model=rm,
325 # The model is live on the reporting server, so it demonstrably runs.
326 compat=ModelCompat.SUPPORTED,
327 )
330def frontier_row_from_remote(
331 rm: RemoteModel, *, provider_id: str, key_status: KeyStatus
332) -> FrontierCatalogRow:
333 """Convert a discovered cloud chat model to a FrontierCatalogRow.
335 ``ref`` is the canonical ``provider/name`` form so callers pass it
336 straight to ``Config.chat_model`` without re-prefixing.
337 """
338 return FrontierCatalogRow(
339 name=rm.name,
340 ref=format_remote_ref(rm.name, rm.provider),
341 task=rm.task,
342 provider=rm.provider,
343 provider_id=provider_id,
344 key_status=key_status,
345 )
348# Column sort key extractors. Local-only because every column except
349# Name reads a field FrontierCatalogRow doesn't carry, and the catalog
350# screen sorts local and frontier rows independently before concat.
351SORT_KEYS: dict[str, Callable[[LocalCatalogRow], Any]] = {
352 "Name": lambda r: r.name.lower(),
353 "Task": lambda r: r.task,
354 "Backend": lambda r: r.backend.lower(),
355 "Params": lambda r: _param_sort_value(r.params),
356 "Size": lambda r: r.sort_size,
357 "Quant": lambda r: r.quant,
358 "Downloads": lambda r: r.sort_downloads,
359}
362def _param_sort_value(params: str) -> float:
363 """Convert param label to sortable float (e.g. '8B' -> 8.0)."""
364 match = re.search(r"(\d+\.?\d*)", params)
365 return float(match.group(1)) if match else 0.0
368def row_delete_id(row: CatalogRow) -> str | None:
369 """Return the model_manager-compatible identifier for *row*.
371 Remote rows hand back the bare ``RemoteModel.name`` because the
372 Ollama HTTP API keys models by bare name, while ``ref`` carries the
373 canonical ``ollama/<name>`` chat_model form.
374 """
375 if row.kind == CatalogRowKind.FRONTIER:
376 return row.ref or None
377 if row.remote_model is not None:
378 return row.remote_model.name or None
379 return row.ref or None
382def matches_search(row: CatalogRow, search: str) -> bool:
383 """Return True if the row matches the search text (hyphen/underscore-insensitive).
385 Local rows match against name/task/params/quant/backend; frontier
386 rows match against name + provider so users can type "gemini" and
387 see every Gemini model regardless of suffix.
388 """
389 if not search:
390 return True
391 needle = _normalize_for_search(search)
392 if row.kind == CatalogRowKind.FRONTIER:
393 return any(
394 needle in _normalize_for_search(field)
395 for field in (row.name, row.provider, row.provider_id)
396 )
397 return any(
398 needle in _normalize_for_search(field)
399 for field in (row.name, row.task, row.params, row.quant, row.backend)
400 )
403def _normalize_for_search(value: str) -> str:
404 return value.lower().replace("-", " ").replace("_", " ")