Coverage for src / lilbee / cli / tui / screens / setup.py: 100%
172 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-06-28 01:01 +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 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.thread_safe import call_from_thread
49from lilbee.cli.tui.widgets.grid_select import GridSelect
50from lilbee.cli.tui.widgets.model_card import ModelCard
51from lilbee.core.config import cfg
52from lilbee.modelhub.models import get_system_ram_gb
54log = logging.getLogger(__name__)
56SETUP_CHAT_GRID_ID = "setup-chat-grid"
59def _scan_installed_models() -> tuple[list[str], list[str]]:
60 """List installed models from the registry, split into chat vs embedding."""
61 try:
62 from lilbee.modelhub.model_manager.discovery import reclassify_by_name
63 from lilbee.modelhub.registry import ModelRegistry
65 registry = ModelRegistry(cfg.models_dir)
66 chat: list[str] = []
67 embed: list[str] = []
68 for m in registry.list_installed():
69 task = reclassify_by_name(m.ref, m.task)
70 if task == ModelTask.EMBEDDING:
71 embed.append(m.ref)
72 elif task == ModelTask.CHAT:
73 chat.append(m.ref)
74 return sorted(chat), sorted(embed)
75 except Exception:
76 return [], []
79def _installed_name_to_row(name: str, task: str) -> LocalCatalogRow:
80 """Build a LocalCatalogRow for an installed registry ref with a cleaned label."""
81 label = display_label_for_ref(name)
82 quant = extract_quant(name.rsplit("/", 1)[-1]) or "--"
83 return LocalCatalogRow(
84 name=label,
85 task=task,
86 params=parse_param_label(label),
87 size="--",
88 quant=quant,
89 downloads="--",
90 featured=False,
91 installed=True,
92 sort_downloads=0,
93 sort_size=0.0,
94 ref=name,
95 )
98def _pick_recommended(ram_gb: float) -> tuple[CatalogModel, CatalogModel]:
99 """Pick chat + embedding models appropriate for system RAM."""
100 eligible = [m for m in FEATURED_CHAT if m.min_ram_gb <= ram_gb]
101 chat = max(eligible, key=lambda m: m.size_gb) if eligible else FEATURED_CHAT[0]
102 embed = FEATURED_EMBEDDING[0]
103 return chat, embed
106def _pending_download(card: ModelCard | None) -> CatalogModel | None:
107 """Return the CatalogModel to download for a non-installed local card, or None."""
108 if card is None:
109 return None
110 row = card.row
111 if row.kind != CatalogRowKind.LOCAL: # setup never shows frontier rows
112 return None
113 if row.installed:
114 return None
115 return row.catalog_model
118class SetupWizard(Screen[str | None]):
119 """First-run setup: pick chat + embedding, Enter installs, Esc exits.
121 Each card you press Enter on:
122 1. Becomes the saved selection for its task (chat or embedding).
123 2. Triggers a download via the app's ``TaskBarController`` unless
124 the card is already installed.
125 3. Leaves the wizard open so you can pick the other task next.
127 The active-model write is chained to the download's on_success hook
128 so chat can never see an active-but-missing ref. Esc-ing the wizard
129 while a download is running is fine: the controller owns the task,
130 and the deferred apply runs against the live app when the file
131 lands on disk.
132 """
134 app: LilbeeApp # type: ignore[assignment]
136 CSS_PATH = "setup.tcss"
138 BINDINGS: ClassVar[list[BindingType]] = [
139 Binding("escape", "cancel", "Done", show=True),
140 Binding("tab", "app.focus_next", "Next", show=False),
141 Binding("shift+tab", "app.focus_previous", "Prev", show=False),
142 ]
144 def __init__(self) -> None:
145 super().__init__()
146 self._selections: dict[str, tuple[str | None, ModelCard | None]] = {
147 ModelTask.CHAT: (None, None),
148 ModelTask.EMBEDDING: (None, None),
149 }
150 self._chat_installed, self._embed_installed = _scan_installed_models()
151 self._recommended_chat: CatalogModel | None = None
152 self._recommended_embed: CatalogModel | None = None
153 # Model refs already submitted to the controller (avoid duplicate
154 # start_download calls when a card is re-selected by arrow + Enter).
155 self._submitted: set[str] = set()
157 @property
158 def _selected_chat(self) -> str | None:
159 return self._selections[ModelTask.CHAT][0]
161 @property
162 def _selected_embed(self) -> str | None:
163 return self._selections[ModelTask.EMBEDDING][0]
165 def compose(self) -> ComposeResult:
166 from textual.widgets import Footer
168 from lilbee.cli.tui.widgets.bottom_bars import BottomBars
169 from lilbee.cli.tui.widgets.status_bar import ViewTabs
170 from lilbee.cli.tui.widgets.task_bar import TaskBar
171 from lilbee.cli.tui.widgets.top_bars import TopBars
173 with TopBars():
174 yield ViewTabs()
175 yield Static(msg.SETUP_WELCOME, id="setup-title")
176 yield Static(msg.SETUP_INTRO, id="setup-intro")
177 yield VerticalScroll(id="setup-grid-container")
178 with BottomBars():
179 yield Label(self._initial_hint_text(), id="setup-enter-hint")
180 yield TaskBar()
181 yield Footer()
183 def _initial_hint_text(self) -> str:
184 """Return SETUP_RETURN_HINT when both roles already resolve, else SETUP_ENTER_HINT."""
185 if self._chat_installed and self._embed_installed:
186 return msg.SETUP_RETURN_HINT
187 return msg.SETUP_ENTER_HINT
189 def on_mount(self) -> None:
190 self._build_grid()
191 # Focus the chat-model grid so arrow keys / Enter work without a mouse.
192 with contextlib.suppress(Exception):
193 self.query_one(f"#{SETUP_CHAT_GRID_ID}", GridSelect).focus()
195 def _build_section(
196 self,
197 heading: str,
198 models: tuple[CatalogModel, ...],
199 installed_refs: set[str],
200 widgets_out: list[Static | GridSelect],
201 grid_id: str | None = None,
202 ) -> list[ModelCard]:
203 """Build a heading + GridSelect for a list of catalog models."""
204 widgets_out.append(Static(heading, classes="section-heading"))
205 cards = [ModelCard(catalog_to_row(m, installed=m.ref in installed_refs)) for m in models]
206 widgets_out.append(GridSelect(*cards, min_column_width=30, max_column_width=50, id=grid_id))
207 return cards
209 def _build_grid(self) -> None:
210 """Build all model sections and pre-select recommended combo."""
211 ram_gb = get_system_ram_gb()
212 rec_chat, rec_embed = _pick_recommended(ram_gb)
213 self._recommended_chat = rec_chat
214 self._recommended_embed = rec_embed
216 container = self.query_one("#setup-grid-container", VerticalScroll)
217 widgets_to_mount: list[Static | GridSelect] = []
218 installed_refs = set(self._chat_installed) | set(self._embed_installed)
220 if self._chat_installed or self._embed_installed:
221 widgets_to_mount.append(Static(msg.HEADING_INSTALLED, classes="section-heading"))
222 installed_cards = [
223 ModelCard(_installed_name_to_row(n, ModelTask.CHAT)) for n in self._chat_installed
224 ] + [
225 ModelCard(_installed_name_to_row(n, ModelTask.EMBEDDING))
226 for n in self._embed_installed
227 ]
228 widgets_to_mount.append(
229 GridSelect(*installed_cards, min_column_width=30, max_column_width=50)
230 )
232 chat_cards = self._build_section(
233 msg.SETUP_HEADING_CHAT,
234 FEATURED_CHAT,
235 installed_refs,
236 widgets_to_mount,
237 grid_id=SETUP_CHAT_GRID_ID,
238 )
239 embed_cards = self._build_section(
240 msg.SETUP_HEADING_EMBED, FEATURED_EMBEDDING, installed_refs, widgets_to_mount
241 )
243 container.mount_all(widgets_to_mount)
244 self._preselect_recommended(chat_cards, embed_cards)
246 def _preselect_recommended(
247 self, chat_cards: list[ModelCard], embed_cards: list[ModelCard]
248 ) -> None:
249 """Pre-select the RAM-appropriate recommended models (without installing)."""
250 for cards, recommended in [
251 (chat_cards, self._recommended_chat),
252 (embed_cards, self._recommended_embed),
253 ]:
254 if not recommended:
255 continue
256 for card in cards:
257 row = card.row
258 if (
259 row.kind != CatalogRowKind.LOCAL
260 ): # pragma: no cover - setup never mounts frontier
261 continue
262 cm = row.catalog_model
263 if cm and cm.ref == recommended.ref:
264 self._mark_selection(card, row.task)
265 break
267 def _mark_selection(self, card: ModelCard, task: str) -> None:
268 """Record a selection and repaint its card. No download yet."""
269 _ref, prev_card = self._selections[task]
270 if prev_card is not None and prev_card is not card:
271 prev_card.selected = False
272 ref = card.row.ref or card.row.name
273 card.selected = True
274 self._selections[task] = (ref, card)
276 def _commit_selection(self, card: ModelCard, task: str) -> None:
277 """Persist the selection: install if needed, set active once the file is on disk."""
278 self._mark_selection(card, task)
279 ref = self._selections[task][0]
280 if ref is None:
281 return
282 pending = _pending_download(card)
283 if pending is None:
284 self._apply_selection(task, ref)
285 return
286 if pending.ref in self._submitted:
287 return
288 self._submitted.add(pending.ref)
289 self.app.task_bar.start_download(
290 pending,
291 on_success=lambda: call_from_thread(self, self._apply_selection, task, ref),
292 )
294 def _apply_selection(self, task: str, ref: str) -> None:
295 """Write the pick into config. Runs on the UI thread after install."""
296 if task == ModelTask.CHAT:
297 apply_active_model(self.app, "chat_model", ref)
298 elif task == ModelTask.EMBEDDING:
299 apply_active_model(self.app, "embedding_model", ref)
301 @on(GridSelect.Selected)
302 def _on_grid_selected(self, event: GridSelect.Selected) -> None:
303 """Enter on a card installs it (or records selection if already installed)."""
304 if not isinstance(event.widget, ModelCard):
305 return
306 card = event.widget
307 task = card.row.task
308 if task in self._selections:
309 self._commit_selection(card, task)
311 @on(GridSelect.LeaveDown)
312 def _on_grid_leave_down(self, event: GridSelect.LeaveDown) -> None:
313 """Arrow-down past the last card walks to the next focusable widget."""
314 self.focus_next()
316 @on(GridSelect.LeaveUp)
317 def _on_grid_leave_up(self, event: GridSelect.LeaveUp) -> None:
318 """Arrow-up past the first card walks to the previous focusable widget."""
319 self.focus_previous()
321 def action_cancel(self) -> None:
322 """Escape dismisses the wizard; any submitted downloads keep running.
324 Selections are saved eagerly in ``_commit_selection``; we reset
325 services here so the next screen pulls the updated config.
326 """
327 if self._selected_chat or self._selected_embed:
328 reset_services()
329 self.dismiss("completed")
330 else:
331 self.dismiss("skipped")