Coverage for src / lilbee / cli / tui / screens / setup.py: 100%
167 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-05-15 20:55 +0000
1"""First-run setup: single-screen model picker with RAM-based recommendations.
3The wizard mirrors the catalog's grid aesthetic: one ``GridSelect`` per
4section (chat, embed), pressing Enter on a card installs that model
5immediately via ``TaskBarController.start_download``. No separate
6Install & Go button, no Browse, no Skip: pick what you want, press
7Esc when done. Downloads continue under the app-level controller, so
8dismissing the wizard while they're in flight is fine.
10Scope: chat and embedding only. Vision and reranker roles are optional,
11so they are configured post-setup via the catalog screen rather than
12gating the first-run path on them.
13"""
15from __future__ import annotations
17import contextlib
18import logging
19from typing import TYPE_CHECKING, ClassVar
21if TYPE_CHECKING:
22 from lilbee.cli.tui.app import LilbeeApp
24from textual import on
25from textual.app import ComposeResult
26from textual.binding import Binding, BindingType
27from textual.containers import VerticalScroll
28from textual.screen import Screen
29from textual.widgets import Label, Static
31from lilbee.app.services import get_services, reset_services
32from lilbee.catalog import (
33 FEATURED_CHAT,
34 FEATURED_EMBEDDING,
35 CatalogModel,
36 display_label_for_ref,
37 extract_quant,
38)
39from lilbee.catalog.types import ModelTask
40from lilbee.cli.tui import messages as msg
41from lilbee.cli.tui.app import apply_active_model
42from lilbee.cli.tui.screens.catalog_utils import (
43 CatalogRowKind,
44 LocalCatalogRow,
45 catalog_to_row,
46 parse_param_label,
47)
48from lilbee.cli.tui.widgets.grid_select import GridSelect
49from lilbee.cli.tui.widgets.model_card import ModelCard
50from lilbee.core.config import cfg
51from lilbee.modelhub.models import get_system_ram_gb
53log = logging.getLogger(__name__)
55SETUP_CHAT_GRID_ID = "setup-chat-grid"
58def _scan_installed_models() -> tuple[list[str], list[str]]:
59 """List installed models from the registry, split into chat vs embedding."""
60 try:
61 from lilbee.modelhub.model_manager.discovery import reclassify_by_name
62 from lilbee.modelhub.registry import ModelRegistry
64 registry = ModelRegistry(cfg.models_dir)
65 chat: list[str] = []
66 embed: list[str] = []
67 for m in registry.list_installed():
68 task = reclassify_by_name(m.ref, m.task)
69 if task == ModelTask.EMBEDDING:
70 embed.append(m.ref)
71 elif task == ModelTask.CHAT:
72 chat.append(m.ref)
73 return sorted(chat), sorted(embed)
74 except Exception:
75 return [], []
78def _installed_name_to_row(name: str, task: str) -> LocalCatalogRow:
79 """Build a LocalCatalogRow for an installed registry ref with a cleaned label."""
80 label = display_label_for_ref(name)
81 quant = extract_quant(name.rsplit("/", 1)[-1]) or "--"
82 return LocalCatalogRow(
83 name=label,
84 task=task,
85 params=parse_param_label(label),
86 size="--",
87 quant=quant,
88 downloads="--",
89 featured=False,
90 installed=True,
91 sort_downloads=0,
92 sort_size=0.0,
93 ref=name,
94 )
97def _pick_recommended(ram_gb: float) -> tuple[CatalogModel, CatalogModel]:
98 """Pick chat + embedding models appropriate for system RAM."""
99 eligible = [m for m in FEATURED_CHAT if m.min_ram_gb <= ram_gb]
100 chat = max(eligible, key=lambda m: m.size_gb) if eligible else FEATURED_CHAT[0]
101 embed = FEATURED_EMBEDDING[0]
102 return chat, embed
105def _pending_download(card: ModelCard | None) -> CatalogModel | None:
106 """Return the CatalogModel to download for a non-installed local card, or None."""
107 if card is None:
108 return None
109 row = card.row
110 if row.kind != CatalogRowKind.LOCAL: # setup never shows frontier rows
111 return None
112 if row.installed:
113 return None
114 return row.catalog_model
117class SetupWizard(Screen[str | None]):
118 """First-run setup: pick chat + embedding, Enter installs, Esc exits.
120 Each card you press Enter on:
121 1. Becomes the saved selection for its task (chat or embedding).
122 2. Triggers a download via the app's ``TaskBarController`` unless
123 the card is already installed.
124 3. Leaves the wizard open so you can pick the other task next.
126 Selections are persisted to settings eagerly (not at dismiss time),
127 so Esc-ing out mid-wizard keeps your picks.
128 """
130 app: LilbeeApp # type: ignore[assignment]
132 CSS_PATH = "setup.tcss"
134 BINDINGS: ClassVar[list[BindingType]] = [
135 Binding("escape", "cancel", "Done", show=True),
136 Binding("tab", "app.focus_next", "Next", show=False),
137 Binding("shift+tab", "app.focus_previous", "Prev", show=False),
138 ]
140 def __init__(self) -> None:
141 super().__init__()
142 self._selections: dict[str, tuple[str | None, ModelCard | None]] = {
143 ModelTask.CHAT: (None, None),
144 ModelTask.EMBEDDING: (None, None),
145 }
146 self._chat_installed, self._embed_installed = _scan_installed_models()
147 self._recommended_chat: CatalogModel | None = None
148 self._recommended_embed: CatalogModel | None = None
149 # Model refs already submitted to the controller (avoid duplicate
150 # start_download calls when a card is re-selected by arrow + Enter).
151 self._submitted: set[str] = set()
153 @property
154 def _selected_chat(self) -> str | None:
155 return self._selections[ModelTask.CHAT][0]
157 @property
158 def _selected_embed(self) -> str | None:
159 return self._selections[ModelTask.EMBEDDING][0]
161 def compose(self) -> ComposeResult:
162 from textual.widgets import Footer
164 from lilbee.cli.tui.widgets.bottom_bars import BottomBars
165 from lilbee.cli.tui.widgets.status_bar import ViewTabs
166 from lilbee.cli.tui.widgets.task_bar import TaskBar
167 from lilbee.cli.tui.widgets.top_bars import TopBars
169 with TopBars():
170 yield ViewTabs()
171 yield Static(msg.SETUP_WELCOME, id="setup-title")
172 yield Static(msg.SETUP_INTRO, id="setup-intro")
173 yield VerticalScroll(id="setup-grid-container")
174 with BottomBars():
175 yield Label(self._initial_hint_text(), id="setup-enter-hint")
176 yield TaskBar()
177 yield Footer()
179 def _initial_hint_text(self) -> str:
180 """Return SETUP_RETURN_HINT when both roles already resolve, else SETUP_ENTER_HINT."""
181 if self._chat_installed and self._embed_installed:
182 return msg.SETUP_RETURN_HINT
183 return msg.SETUP_ENTER_HINT
185 def on_mount(self) -> None:
186 self._build_grid()
187 # Focus the chat-model grid so arrow keys / Enter work without a mouse.
188 with contextlib.suppress(Exception):
189 self.query_one(f"#{SETUP_CHAT_GRID_ID}", GridSelect).focus()
191 def _build_section(
192 self,
193 heading: str,
194 models: tuple[CatalogModel, ...],
195 installed_refs: set[str],
196 widgets_out: list[Static | GridSelect],
197 grid_id: str | None = None,
198 ) -> list[ModelCard]:
199 """Build a heading + GridSelect for a list of catalog models."""
200 widgets_out.append(Static(heading, classes="section-heading"))
201 cards = [ModelCard(catalog_to_row(m, installed=m.ref in installed_refs)) for m in models]
202 widgets_out.append(GridSelect(*cards, min_column_width=30, max_column_width=50, id=grid_id))
203 return cards
205 def _build_grid(self) -> None:
206 """Build all model sections and pre-select recommended combo."""
207 ram_gb = get_system_ram_gb()
208 rec_chat, rec_embed = _pick_recommended(ram_gb)
209 self._recommended_chat = rec_chat
210 self._recommended_embed = rec_embed
212 container = self.query_one("#setup-grid-container", VerticalScroll)
213 widgets_to_mount: list[Static | GridSelect] = []
214 installed_refs = set(self._chat_installed) | set(self._embed_installed)
216 if self._chat_installed or self._embed_installed:
217 widgets_to_mount.append(Static(msg.HEADING_INSTALLED, classes="section-heading"))
218 installed_cards = [
219 ModelCard(_installed_name_to_row(n, ModelTask.CHAT)) for n in self._chat_installed
220 ] + [
221 ModelCard(_installed_name_to_row(n, ModelTask.EMBEDDING))
222 for n in self._embed_installed
223 ]
224 widgets_to_mount.append(
225 GridSelect(*installed_cards, min_column_width=30, max_column_width=50)
226 )
228 chat_cards = self._build_section(
229 msg.SETUP_HEADING_CHAT,
230 FEATURED_CHAT,
231 installed_refs,
232 widgets_to_mount,
233 grid_id=SETUP_CHAT_GRID_ID,
234 )
235 embed_cards = self._build_section(
236 msg.SETUP_HEADING_EMBED, FEATURED_EMBEDDING, installed_refs, widgets_to_mount
237 )
239 container.mount_all(widgets_to_mount)
240 self._preselect_recommended(chat_cards, embed_cards)
242 def _preselect_recommended(
243 self, chat_cards: list[ModelCard], embed_cards: list[ModelCard]
244 ) -> None:
245 """Pre-select the RAM-appropriate recommended models (without installing)."""
246 for cards, recommended in [
247 (chat_cards, self._recommended_chat),
248 (embed_cards, self._recommended_embed),
249 ]:
250 if not recommended:
251 continue
252 for card in cards:
253 row = card.row
254 if (
255 row.kind != CatalogRowKind.LOCAL
256 ): # pragma: no cover - setup never mounts frontier
257 continue
258 cm = row.catalog_model
259 if cm and cm.ref == recommended.ref:
260 self._mark_selection(card, row.task)
261 break
263 def _mark_selection(self, card: ModelCard, task: str) -> None:
264 """Record a selection and repaint its card. No download yet."""
265 _ref, prev_card = self._selections[task]
266 if prev_card is not None and prev_card is not card:
267 prev_card.selected = False
268 ref = card.row.ref or card.row.name
269 card.selected = True
270 self._selections[task] = (ref, card)
272 def _commit_selection(self, card: ModelCard, task: str) -> None:
273 """Persist the selection to settings and submit a download if pending.
275 Called when the user presses Enter on a card. Saves the config
276 fragment eagerly so Esc mid-wizard doesn't lose the pick.
277 """
278 self._mark_selection(card, task)
279 ref = self._selections[task][0]
280 if ref is None:
281 return
282 if task == ModelTask.CHAT:
283 apply_active_model(self.app, "chat_model", ref)
284 elif task == ModelTask.EMBEDDING:
285 # Pin a legacy store's identity to the OLD model BEFORE the cfg
286 # mutation so the gate in store.search/add_chunks correctly detects
287 # drift on the next op. See bb-x1qa.
288 get_services().store.initialize_meta_if_legacy()
289 apply_active_model(self.app, "embedding_model", ref)
291 pending = _pending_download(card)
292 if pending is not None and pending.ref not in self._submitted:
293 self._submitted.add(pending.ref)
294 self.app.task_bar.start_download(pending)
296 @on(GridSelect.Selected)
297 def _on_grid_selected(self, event: GridSelect.Selected) -> None:
298 """Enter on a card installs it (or records selection if already installed)."""
299 if not isinstance(event.widget, ModelCard):
300 return
301 card = event.widget
302 task = card.row.task
303 if task in self._selections:
304 self._commit_selection(card, task)
306 @on(GridSelect.LeaveDown)
307 def _on_grid_leave_down(self, event: GridSelect.LeaveDown) -> None:
308 """Arrow-down past the last card walks to the next focusable widget."""
309 self.focus_next()
311 @on(GridSelect.LeaveUp)
312 def _on_grid_leave_up(self, event: GridSelect.LeaveUp) -> None:
313 """Arrow-up past the first card walks to the previous focusable widget."""
314 self.focus_previous()
316 def action_cancel(self) -> None:
317 """Escape dismisses the wizard; any submitted downloads keep running.
319 Selections are saved eagerly in ``_commit_selection``; we reset
320 services here so the next screen pulls the updated config.
321 """
322 if self._selected_chat or self._selected_embed:
323 reset_services()
324 self.dismiss("completed")
325 else:
326 self.dismiss("skipped")