Coverage for src / lilbee / cli / tui / screens / setup.py: 100%
156 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-04-29 19:16 +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 ClassVar
21from textual import on
22from textual.app import ComposeResult
23from textual.binding import Binding, BindingType
24from textual.containers import VerticalScroll
25from textual.screen import Screen
26from textual.widgets import Label, Static
28from lilbee.catalog import FEATURED_CHAT, FEATURED_EMBEDDING, CatalogModel
29from lilbee.cli.tui import messages as msg
30from lilbee.cli.tui.app import apply_active_model
31from lilbee.cli.tui.screens.catalog_utils import (
32 TableRow,
33 catalog_to_row,
34 parse_param_label,
35)
36from lilbee.cli.tui.widgets.grid_select import GridSelect
37from lilbee.cli.tui.widgets.model_card import ModelCard
38from lilbee.config import cfg
39from lilbee.models import ModelTask, get_system_ram_gb
40from lilbee.services import get_services, reset_services
42log = logging.getLogger(__name__)
44SETUP_CHAT_GRID_ID = "setup-chat-grid"
47def _scan_installed_models() -> tuple[list[str], list[str]]:
48 """List installed models from the registry, split into chat vs embedding."""
49 try:
50 from lilbee.registry import ModelRegistry
52 registry = ModelRegistry(cfg.models_dir)
53 chat: list[str] = []
54 embed: list[str] = []
55 for m in registry.list_installed():
56 if m.task == ModelTask.EMBEDDING:
57 embed.append(m.ref)
58 elif m.task == ModelTask.CHAT:
59 chat.append(m.ref)
60 return sorted(chat), sorted(embed)
61 except Exception:
62 return [], []
65def _installed_name_to_row(name: str, task: str) -> TableRow:
66 """Create a minimal TableRow for an already-installed model."""
67 return TableRow(
68 name=name,
69 task=task,
70 params=parse_param_label(name),
71 size="--",
72 quant="--",
73 downloads="--",
74 featured=False,
75 installed=True,
76 sort_downloads=0,
77 sort_size=0.0,
78 ref=name,
79 )
82def _pick_recommended(ram_gb: float) -> tuple[CatalogModel, CatalogModel]:
83 """Pick chat + embedding models appropriate for system RAM."""
84 eligible = [m for m in FEATURED_CHAT if m.min_ram_gb <= ram_gb]
85 chat = max(eligible, key=lambda m: m.size_gb) if eligible else FEATURED_CHAT[0]
86 embed = FEATURED_EMBEDDING[0]
87 return chat, embed
90def _pending_download(card: ModelCard | None) -> CatalogModel | None:
91 """Return the CatalogModel to download for a non-installed card, or None."""
92 if card and not card.row.installed:
93 return card.row.catalog_model
94 return None
97class SetupWizard(Screen[str | None]):
98 """First-run setup — pick chat + embedding, Enter installs, Esc exits.
100 Each card you press Enter on:
101 1. Becomes the saved selection for its task (chat or embedding).
102 2. Triggers a download via the app's ``TaskBarController`` unless
103 the card is already installed.
104 3. Leaves the wizard open so you can pick the other task next.
106 Selections are persisted to settings eagerly (not at dismiss time),
107 so Esc-ing out mid-wizard keeps your picks.
108 """
110 CSS_PATH = "setup.tcss"
112 BINDINGS: ClassVar[list[BindingType]] = [
113 Binding("escape", "cancel", "Done", show=True),
114 Binding("tab", "app.focus_next", "Next", show=False),
115 Binding("shift+tab", "app.focus_previous", "Prev", show=False),
116 ]
118 def __init__(self) -> None:
119 super().__init__()
120 self._selections: dict[str, tuple[str | None, ModelCard | None]] = {
121 ModelTask.CHAT: (None, None),
122 ModelTask.EMBEDDING: (None, None),
123 }
124 self._chat_installed, self._embed_installed = _scan_installed_models()
125 self._recommended_chat: CatalogModel | None = None
126 self._recommended_embed: CatalogModel | None = None
127 # Model refs already submitted to the controller (avoid duplicate
128 # start_download calls when a card is re-selected by arrow + Enter).
129 self._submitted: set[str] = set()
131 @property
132 def _selected_chat(self) -> str | None:
133 return self._selections[ModelTask.CHAT][0]
135 @property
136 def _selected_embed(self) -> str | None:
137 return self._selections[ModelTask.EMBEDDING][0]
139 def compose(self) -> ComposeResult:
140 from textual.widgets import Footer
142 from lilbee.cli.tui.widgets.bottom_bars import BottomBars
143 from lilbee.cli.tui.widgets.status_bar import ViewTabs
144 from lilbee.cli.tui.widgets.task_bar import TaskBar
145 from lilbee.cli.tui.widgets.top_bars import TopBars
147 with TopBars():
148 yield ViewTabs()
149 yield Static(msg.SETUP_WELCOME, id="setup-title")
150 yield Static(msg.SETUP_INTRO, id="setup-intro")
151 yield VerticalScroll(id="setup-grid-container")
152 with BottomBars():
153 yield Label(self._initial_hint_text(), id="setup-enter-hint")
154 yield TaskBar()
155 yield Footer()
157 def _initial_hint_text(self) -> str:
158 """Return SETUP_RETURN_HINT when both roles already resolve, else SETUP_ENTER_HINT."""
159 if self._chat_installed and self._embed_installed:
160 return msg.SETUP_RETURN_HINT
161 return msg.SETUP_ENTER_HINT
163 def on_mount(self) -> None:
164 self._build_grid()
165 # Focus the chat-model grid so arrow keys / Enter work without a mouse.
166 with contextlib.suppress(Exception):
167 self.query_one(f"#{SETUP_CHAT_GRID_ID}", GridSelect).focus()
169 def _build_section(
170 self,
171 heading: str,
172 models: tuple[CatalogModel, ...],
173 installed_refs: set[str],
174 widgets_out: list[Static | GridSelect],
175 grid_id: str | None = None,
176 ) -> list[ModelCard]:
177 """Build a heading + GridSelect for a list of catalog models."""
178 widgets_out.append(Static(heading, classes="section-heading"))
179 cards = [ModelCard(catalog_to_row(m, installed=m.ref in installed_refs)) for m in models]
180 widgets_out.append(GridSelect(*cards, min_column_width=30, max_column_width=50, id=grid_id))
181 return cards
183 def _build_grid(self) -> None:
184 """Build all model sections and pre-select recommended combo."""
185 ram_gb = get_system_ram_gb()
186 rec_chat, rec_embed = _pick_recommended(ram_gb)
187 self._recommended_chat = rec_chat
188 self._recommended_embed = rec_embed
190 container = self.query_one("#setup-grid-container", VerticalScroll)
191 widgets_to_mount: list[Static | GridSelect] = []
192 installed_refs = set(self._chat_installed) | set(self._embed_installed)
194 if self._chat_installed or self._embed_installed:
195 widgets_to_mount.append(Static(msg.HEADING_INSTALLED, classes="section-heading"))
196 installed_cards = [
197 ModelCard(_installed_name_to_row(n, ModelTask.CHAT)) for n in self._chat_installed
198 ] + [
199 ModelCard(_installed_name_to_row(n, ModelTask.EMBEDDING))
200 for n in self._embed_installed
201 ]
202 widgets_to_mount.append(
203 GridSelect(*installed_cards, min_column_width=30, max_column_width=50)
204 )
206 chat_cards = self._build_section(
207 msg.SETUP_HEADING_CHAT,
208 FEATURED_CHAT,
209 installed_refs,
210 widgets_to_mount,
211 grid_id=SETUP_CHAT_GRID_ID,
212 )
213 embed_cards = self._build_section(
214 msg.SETUP_HEADING_EMBED, FEATURED_EMBEDDING, installed_refs, widgets_to_mount
215 )
217 container.mount_all(widgets_to_mount)
218 self._preselect_recommended(chat_cards, embed_cards)
220 def _preselect_recommended(
221 self, chat_cards: list[ModelCard], embed_cards: list[ModelCard]
222 ) -> None:
223 """Pre-select the RAM-appropriate recommended models (without installing)."""
224 for cards, recommended in [
225 (chat_cards, self._recommended_chat),
226 (embed_cards, self._recommended_embed),
227 ]:
228 if not recommended:
229 continue
230 for card in cards:
231 cm = card.row.catalog_model
232 if cm and cm.ref == recommended.ref:
233 self._mark_selection(card, card.row.task)
234 break
236 def _mark_selection(self, card: ModelCard, task: str) -> None:
237 """Record a selection and repaint its card. No download yet."""
238 _ref, prev_card = self._selections[task]
239 if prev_card is not None and prev_card is not card:
240 prev_card.selected = False
241 ref = card.row.ref or card.row.name
242 card.selected = True
243 self._selections[task] = (ref, card)
245 def _commit_selection(self, card: ModelCard, task: str) -> None:
246 """Persist the selection to settings and submit a download if pending.
248 Called when the user presses Enter on a card. Saves the config
249 fragment eagerly so Esc mid-wizard doesn't lose the pick.
250 """
251 from lilbee.cli.tui.app import LilbeeApp
253 self._mark_selection(card, task)
254 ref = self._selections[task][0]
255 if ref is None:
256 return
257 if task == ModelTask.CHAT:
258 apply_active_model(self.app, "chat_model", ref)
259 elif task == ModelTask.EMBEDDING:
260 # Pin a legacy store's identity to the OLD model BEFORE the cfg
261 # mutation so the gate in store.search/add_chunks correctly detects
262 # drift on the next op. See bb-x1qa.
263 get_services().store.initialize_meta_if_legacy()
264 apply_active_model(self.app, "embedding_model", ref)
266 pending = _pending_download(card)
267 if (
268 pending is not None
269 and pending.ref not in self._submitted
270 and isinstance(self.app, LilbeeApp)
271 ):
272 self._submitted.add(pending.ref)
273 self.app.task_bar.start_download(pending)
275 @on(GridSelect.Selected)
276 def _on_grid_selected(self, event: GridSelect.Selected) -> None:
277 """Enter on a card installs it (or records selection if already installed)."""
278 if not isinstance(event.widget, ModelCard):
279 return
280 card = event.widget
281 task = card.row.task
282 if task in self._selections:
283 self._commit_selection(card, task)
285 @on(GridSelect.LeaveDown)
286 def _on_grid_leave_down(self, event: GridSelect.LeaveDown) -> None:
287 """Arrow-down past the last card walks to the next focusable widget."""
288 self.focus_next()
290 @on(GridSelect.LeaveUp)
291 def _on_grid_leave_up(self, event: GridSelect.LeaveUp) -> None:
292 """Arrow-up past the first card walks to the previous focusable widget."""
293 self.focus_previous()
295 def action_cancel(self) -> None:
296 """Escape dismisses the wizard; any submitted downloads keep running.
298 Selections are saved eagerly in ``_commit_selection``; we reset
299 services here so the next screen pulls the updated config.
300 """
301 if self._selected_chat or self._selected_embed:
302 reset_services()
303 self.dismiss("completed")
304 else:
305 self.dismiss("skipped")