Coverage for src/lilbee/cli/tui/screens/catalog_grouping.py: 100%
93 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +0000
1"""Row-grouping helpers and the GridSection container for CatalogScreen."""
3from __future__ import annotations
5from dataclasses import dataclass
6from typing import cast
8from lilbee.catalog.types import ModelCompat, ModelTask
9from lilbee.cli.tui import messages as msg
10from lilbee.cli.tui.screens.catalog_utils import (
11 CatalogRow,
12 CatalogRowKind,
13 FrontierCatalogRow,
14 LocalCatalogRow,
15)
16from lilbee.cli.tui.widgets.model_list import ModelListSection
17from lilbee.core.config import cfg
18from lilbee.runtime.hardware import FIT_RANK, FitLevel
21@dataclass
22class GridSection:
23 """A named group of rows for the grid view."""
25 heading: str
26 rows: list[CatalogRow]
29TASK_BUCKET_ORDER = (ModelTask.CHAT, ModelTask.EMBEDDING, ModelTask.VISION, ModelTask.RERANK)
30PICKS_SECTION_HEADING = "★ Picks"
33def row_cache_signature(row: CatalogRow) -> tuple[str, bool]:
34 """Pair (name, installed-flag) for the per-tab cache key.
36 Frontier rows don't carry an ``installed`` field; they're keyed as
37 if installed=False since each frontier entry is provider-managed
38 rather than on-disk.
39 """
40 if row.kind == CatalogRowKind.FRONTIER:
41 return (row.name, False)
42 return (row.name, row.installed)
45def _is_runnable_pick(row: LocalCatalogRow) -> bool:
46 """Whether the row is a recommendation: supported, runnable, and either clean or opted in."""
47 return (
48 row.compat is ModelCompat.SUPPORTED
49 and (cfg.include_uncensored or not row.safety_stripped)
50 and (row.fit is None or row.fit.level is not FitLevel.WONT_RUN)
51 )
54def _unknown_fit_rank(row: LocalCatalogRow) -> int:
55 """0 for a row with a fit chip, 1 for a row without one, so measured rows lead."""
56 return 0 if row.fit is not None else 1
59def _backfill_sort_key(row: LocalCatalogRow) -> tuple[int, int, str]:
60 """Rank a backfilled Discover pick: known fit first, most downloaded, then alphabetical."""
61 return (_unknown_fit_rank(row), -row.sort_downloads, row.name.lower())
64def _failed_probe_sort_key(row: LocalCatalogRow) -> tuple[float, str]:
65 """Rank an unmeasured Discover pick: smallest first, then alphabetical."""
66 return (row.sort_size, row.name.lower())
69def _pick_for_task(candidates: list[LocalCatalogRow]) -> LocalCatalogRow:
70 """One Discover pick for a role: smallest when nothing is measured, else featured-leads."""
71 if all(r.fit is None for r in candidates):
72 return min(candidates, key=_failed_probe_sort_key)
73 featured = [r for r in candidates if r.featured]
74 if featured:
75 return min(featured, key=for_you_sort_key)
76 return min(candidates, key=_backfill_sort_key)
79def for_you_by_role(rows: list[LocalCatalogRow]) -> list[LocalCatalogRow]:
80 """Runnable picks grouped by role: chat, embedding, vision, rerank.
82 Featured rows lead, known fit first. When no row carries a fit chip
83 the smallest row wins across featured and non-featured rows. A role
84 whose featured rows cannot run backfills with the most downloaded row
85 the host cannot rule out, so a card that does not fit is replaced
86 rather than dropped. A role with nothing runnable yields no pick.
87 """
88 runnable = [r for r in rows if _is_runnable_pick(r)]
89 out: list[LocalCatalogRow] = []
90 for task in TASK_BUCKET_ORDER:
91 candidates = [r for r in runnable if r.task == task]
92 if candidates:
93 out.append(_pick_for_task(candidates))
94 return out
97def for_you_sort_key(row: LocalCatalogRow) -> tuple[int, int, str]:
98 """Rank Discover 'For You' rows: known fit first, best fit next, then alphabetical."""
99 level_rank = 0 if row.fit is None else FIT_RANK[row.fit.level]
100 return (_unknown_fit_rank(row), level_rank, row.name.lower())
103def group_frontier_rows(
104 frontier_rows: list[FrontierCatalogRow],
105) -> list[ModelListSection]:
106 """Group frontier rows into provider-headed sections.
108 Section order follows :data:`PROVIDER_KEYS` (the canonical display
109 order); providers absent from PROVIDER_KEYS land at the tail in
110 alphabetical order. Rows within each section are alphabetical.
111 """
112 if not frontier_rows:
113 return []
114 from lilbee.providers.sdk_backend import PROVIDER_KEYS
116 per_provider: dict[str, list[FrontierCatalogRow]] = {}
117 for row in frontier_rows:
118 per_provider.setdefault(row.provider, []).append(row)
119 canonical_order = [label for _, _, _, label in PROVIDER_KEYS]
120 ordered = [p for p in canonical_order if p in per_provider]
121 extras = sorted(set(per_provider) - set(canonical_order))
122 sections: list[ModelListSection] = []
123 for provider in [*ordered, *extras]:
124 rows = sorted(per_provider[provider], key=lambda r: r.name.lower())
125 sections.append(ModelListSection(heading=provider, rows=list(rows)))
126 return sections
129def group_task_rows_with_picks(
130 task_rows: list[LocalCatalogRow], task_label: str
131) -> list[GridSection]:
132 """Per-tab grouping: ★ Picks pinned, then Installed, then the rest.
134 Lifts featured rows out of their task bucket into a dedicated pinned
135 section at the top of the tab. Today's behavior interleaved them at
136 the top of the task bucket; the redesign treats curation as its own
137 layer so the eye lands on Picks first instead of having to scan past
138 them to find non-featured rows.
140 Pre-condition: caller has already filtered ``task_rows`` to a single
141 task (the active per-task tab).
142 """
143 picks: list[CatalogRow] = []
144 installed: list[CatalogRow] = []
145 others: list[CatalogRow] = []
146 for row in task_rows:
147 if row.featured:
148 picks.append(row)
149 elif row.installed:
150 installed.append(row)
151 else:
152 others.append(row)
153 return [
154 GridSection(PICKS_SECTION_HEADING, picks),
155 GridSection(msg.HEADING_INSTALLED, installed),
156 GridSection(task_label, others),
157 ]
160def flatten_sections(sections: list[GridSection], heading: str) -> list[GridSection]:
161 """Collapse *sections* into a single section, preserving row order.
163 Used while a search filter is active: every mounted section costs a heading
164 plus a whole card row even when it holds one match.
165 """
166 rows = [row for section in sections for row in section.rows]
167 if not rows:
168 return []
169 return [GridSection(heading, rows)]
172def group_rows_for_grid(local_rows: list[LocalCatalogRow]) -> list[GridSection]:
173 """Group local rows into sections for the grid view.
175 Layout: Installed first, then one section per task. Featured rows live
176 at the top of their task section (recognizable by the ``pick`` pill);
177 no separate "Our picks" bucket so the catalog reads as a single
178 task-organized list.
179 """
180 installed: list[CatalogRow] = []
181 by_task: dict[str, list[CatalogRow]] = {task: [] for task in TASK_BUCKET_ORDER}
182 extras: dict[str, list[CatalogRow]] = {}
183 for row in local_rows:
184 if row.installed:
185 installed.append(row)
186 continue
187 bucket = by_task.get(row.task)
188 if bucket is not None:
189 bucket.append(row)
190 else:
191 extras.setdefault(row.task, []).append(row)
192 # Within each task bucket: featured first (preserving their input order),
193 # then the rest in their incoming order. Stable so HF rank from the API
194 # is preserved among non-featured rows.
195 for bucket in by_task.values():
196 bucket.sort(key=lambda r: not cast("LocalCatalogRow", r).featured)
197 for bucket in extras.values():
198 bucket.sort(key=lambda r: not cast("LocalCatalogRow", r).featured)
199 return [
200 GridSection(msg.HEADING_INSTALLED, installed),
201 *[GridSection(task.capitalize(), by_task[task]) for task in TASK_BUCKET_ORDER],
202 *[GridSection(task.capitalize(), extras[task]) for task in extras],
203 ]