Coverage for src/lilbee/cli/tui/screens/catalog_grouping.py: 100%
85 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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 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
19@dataclass
20class GridSection:
21 """A named group of rows for the grid view."""
23 heading: str
24 rows: list[CatalogRow]
27TASK_BUCKET_ORDER = (ModelTask.CHAT, ModelTask.EMBEDDING, ModelTask.VISION, ModelTask.RERANK)
28PICKS_SECTION_HEADING = "★ Picks"
31def row_cache_signature(row: CatalogRow) -> tuple[str, bool]:
32 """Pair (name, installed-flag) for the per-tab cache key.
34 Frontier rows don't carry an ``installed`` field; they're keyed as
35 if installed=False since each frontier entry is provider-managed
36 rather than on-disk.
37 """
38 if row.kind == CatalogRowKind.FRONTIER:
39 return (row.name, False)
40 return (row.name, row.installed)
43def for_you_by_role(rows: list[LocalCatalogRow]) -> list[LocalCatalogRow]:
44 """Runnable picks grouped by role: chat, embedding, vision, rerank.
46 A row qualifies only when the engine supports its architecture and the
47 machine can hold it, so every card in the rail is installable as-is.
48 Rows with no fit chip are excluded: an unknown size cannot be promised.
49 """
50 from lilbee.catalog.types import ModelCompat, ModelTask
51 from lilbee.runtime.hardware import FitLevel
53 runnable = [
54 r
55 for r in rows
56 if r.featured
57 and r.compat is ModelCompat.SUPPORTED
58 and r.fit is not None
59 and r.fit.level is not FitLevel.WONT_RUN
60 ]
61 out: list[LocalCatalogRow] = []
62 for task in (ModelTask.CHAT, ModelTask.EMBEDDING, ModelTask.VISION, ModelTask.RERANK):
63 for row in sorted((r for r in runnable if r.task == task), key=for_you_sort_key):
64 out.append(row)
65 break
66 return out
69def for_you_sort_key(row: LocalCatalogRow) -> tuple[int, str]:
70 """Rank Discover 'For You' rows: best fit first, then alphabetical.
72 Fit rank: FITS=0, TIGHT=1, WONT_RUN=2, no chip=3. Featured-only
73 callers already filtered, so featured isn't in the key.
74 """
75 from lilbee.runtime.hardware import FitLevel
77 if row.fit is None:
78 rank = 3
79 elif row.fit.level is FitLevel.FITS:
80 rank = 0
81 elif row.fit.level is FitLevel.TIGHT:
82 rank = 1
83 else:
84 rank = 2
85 return (rank, row.name.lower())
88def group_frontier_rows(
89 frontier_rows: list[FrontierCatalogRow],
90) -> list[ModelListSection]:
91 """Group frontier rows into provider-headed sections.
93 Section order follows :data:`PROVIDER_KEYS` (the canonical display
94 order); providers absent from PROVIDER_KEYS land at the tail in
95 alphabetical order. Rows within each section are alphabetical.
96 """
97 if not frontier_rows:
98 return []
99 from lilbee.providers.sdk_backend import PROVIDER_KEYS
101 per_provider: dict[str, list[FrontierCatalogRow]] = {}
102 for row in frontier_rows:
103 per_provider.setdefault(row.provider, []).append(row)
104 canonical_order = [label for _, _, _, label in PROVIDER_KEYS]
105 ordered = [p for p in canonical_order if p in per_provider]
106 extras = sorted(set(per_provider) - set(canonical_order))
107 sections: list[ModelListSection] = []
108 for provider in [*ordered, *extras]:
109 rows = sorted(per_provider[provider], key=lambda r: r.name.lower())
110 sections.append(ModelListSection(heading=provider, rows=list(rows)))
111 return sections
114def group_task_rows_with_picks(
115 task_rows: list[LocalCatalogRow], task_label: str
116) -> list[GridSection]:
117 """Per-tab grouping: ★ Picks pinned, then Installed, then the rest.
119 Lifts featured rows out of their task bucket into a dedicated pinned
120 section at the top of the tab. Today's behavior interleaved them at
121 the top of the task bucket; the redesign treats curation as its own
122 layer so the eye lands on Picks first instead of having to scan past
123 them to find non-featured rows.
125 Pre-condition: caller has already filtered ``task_rows`` to a single
126 task (the active per-task tab).
127 """
128 picks: list[CatalogRow] = []
129 installed: list[CatalogRow] = []
130 others: list[CatalogRow] = []
131 for row in task_rows:
132 if row.featured:
133 picks.append(row)
134 elif row.installed:
135 installed.append(row)
136 else:
137 others.append(row)
138 return [
139 GridSection(PICKS_SECTION_HEADING, picks),
140 GridSection(msg.HEADING_INSTALLED, installed),
141 GridSection(task_label, others),
142 ]
145def flatten_sections(sections: list[GridSection], heading: str) -> list[GridSection]:
146 """Collapse *sections* into a single section, preserving row order.
148 Used while a search filter is active: every mounted section costs a heading
149 plus a whole card row even when it holds one match.
150 """
151 rows = [row for section in sections for row in section.rows]
152 if not rows:
153 return []
154 return [GridSection(heading, rows)]
157def group_rows_for_grid(local_rows: list[LocalCatalogRow]) -> list[GridSection]:
158 """Group local rows into sections for the grid view.
160 Layout: Installed first, then one section per task. Featured rows live
161 at the top of their task section (recognizable by the ``pick`` pill);
162 no separate "Our picks" bucket so the catalog reads as a single
163 task-organized list.
164 """
165 installed: list[CatalogRow] = []
166 by_task: dict[str, list[CatalogRow]] = {task: [] for task in TASK_BUCKET_ORDER}
167 extras: dict[str, list[CatalogRow]] = {}
168 for row in local_rows:
169 if row.installed:
170 installed.append(row)
171 continue
172 bucket = by_task.get(row.task)
173 if bucket is not None:
174 bucket.append(row)
175 else:
176 extras.setdefault(row.task, []).append(row)
177 # Within each task bucket: featured first (preserving their input order),
178 # then the rest in their incoming order. Stable so HF rank from the API
179 # is preserved among non-featured rows.
180 for bucket in by_task.values():
181 bucket.sort(key=lambda r: not cast("LocalCatalogRow", r).featured)
182 for bucket in extras.values():
183 bucket.sort(key=lambda r: not cast("LocalCatalogRow", r).featured)
184 return [
185 GridSection(msg.HEADING_INSTALLED, installed),
186 *[GridSection(task.capitalize(), by_task[task]) for task in TASK_BUCKET_ORDER],
187 *[GridSection(task.capitalize(), extras[task]) for task in extras],
188 ]