Coverage for src/lilbee/cli/tui/screens/status.py: 100%

212 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-08-14 11:46 +0000

1"""Status screen: knowledge base info with collapsible sections.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import contextlib 

7import logging 

8from dataclasses import dataclass 

9from pathlib import Path 

10from typing import TYPE_CHECKING, ClassVar 

11 

12if TYPE_CHECKING: 

13 from lilbee.cli.tui.app import LilbeeApp 

14 

15from textual import work 

16from textual.app import ComposeResult 

17from textual.binding import Binding, BindingType 

18from textual.containers import VerticalScroll 

19from textual.content import Content 

20from textual.screen import Screen 

21from textual.widgets import Collapsible, DataTable, Static 

22from textual.worker import Worker, WorkerState 

23 

24from lilbee.app.services import get_services 

25from lilbee.cli.tui import messages as msg 

26from lilbee.cli.tui.browse_bindings import BROWSE_LIST_BINDINGS, browse_back_bindings 

27from lilbee.cli.tui.pill import pill 

28from lilbee.core.config import cfg 

29from lilbee.data.store import SourceRecord 

30from lilbee.modelhub.model_info import ModelArchInfo, get_model_architecture 

31 

32log = logging.getLogger(__name__) 

33 

34# Rows appended to the Documents table per refresh tick. Each ``add_row`` runs 

35# on the UI thread, so dumping a whole ingested wiki in one loop froze the 

36# screen for seconds; rendering a bounded batch per ``call_after_refresh`` lets 

37# the user scroll and interact while the rest streams in. 

38_DOC_RENDER_BATCH = 100 

39 

40 

41@dataclass 

42class _DocsResult: 

43 """Outcome of the background sources read. 

44 

45 ``load_failed`` distinguishes "store opened but empty" (``sources == []``) 

46 from "the read raised" so the UI shows the right placeholder. 

47 """ 

48 

49 sources: list[SourceRecord] 

50 load_failed: bool 

51 

52 

53def _model_pill(name: str) -> Content: 

54 """Return a green 'loaded' pill if name is set, red 'not set' otherwise.""" 

55 if name: 

56 return pill("loaded", "$success", "$text") 

57 return pill("not set", "$error", "$text") 

58 

59 

60# Label-column width used across the status sections so keys line up 

61# when scanned vertically. Values past this column render bold. 

62_KV_LABEL_WIDTH = 14 

63 

64 

65def _kv_line(label: str, value: str | Content, status: Content | None = None) -> Content: 

66 """Assemble one key/value row: dim padded label, bold value, optional pill.""" 

67 padded = label.ljust(_KV_LABEL_WIDTH) 

68 parts: list[Content] = [Content.styled(padded, "$text-muted")] 

69 if isinstance(value, Content): 

70 parts.append(value) 

71 else: 

72 parts.append(Content.styled(value, "bold")) 

73 if status is not None: 

74 parts.append(Content(" ")) 

75 parts.append(status) 

76 return Content.assemble(*parts) 

77 

78 

79def _collapse_home(path: Path | str) -> str: 

80 """Replace the user's home prefix with '~' so long paths stay scannable.""" 

81 text = str(path) 

82 home = str(Path.home()) 

83 return text.replace(home, "~", 1) if text.startswith(home) else text 

84 

85 

86def _ocr_label() -> str: 

87 """Return a human-readable OCR status string.""" 

88 if cfg.enable_ocr is True: 

89 return "enabled" 

90 if cfg.enable_ocr is False: 

91 return "disabled" 

92 return "auto" 

93 

94 

95def _ocr_pill() -> Content: 

96 """Return a pill reflecting OCR status.""" 

97 if cfg.enable_ocr is True: 

98 return pill("on", "$success", "$text") 

99 if cfg.enable_ocr is False: 

100 return pill("off", "$warning", "$text") 

101 return pill("auto", "$accent", "$text") 

102 

103 

104def _data_dir_pill() -> Content: 

105 """Return a pill based on whether the data directory exists.""" 

106 if Path(cfg.data_dir).exists(): 

107 return pill("exists", "$success", "$text") 

108 return pill("missing", "$error", "$text") 

109 

110 

111def _build_config_content() -> Content: 

112 """Build the configuration section content.""" 

113 lines = [ 

114 _kv_line("Data dir", _collapse_home(cfg.data_dir), _data_dir_pill()), 

115 _kv_line("Chat model", cfg.chat_model or "(disabled)", _model_pill(cfg.chat_model)), 

116 _kv_line( 

117 "Embed model", cfg.embedding_model or "(disabled)", _model_pill(cfg.embedding_model) 

118 ), 

119 _kv_line("Vision model", cfg.vision_model or "(disabled)", _model_pill(cfg.vision_model)), 

120 _kv_line("Reranker", cfg.reranker_model or "(disabled)", _model_pill(cfg.reranker_model)), 

121 _kv_line("OCR", _ocr_label(), _ocr_pill()), 

122 ] 

123 return Content("\n").join(lines) 

124 

125 

126def _build_storage_content(doc_count: int) -> Content: 

127 """Build the storage section content.""" 

128 lines = [ 

129 _kv_line("Documents", str(doc_count)), 

130 _kv_line("Data dir", _collapse_home(cfg.data_dir)), 

131 _kv_line("Models dir", _collapse_home(cfg.models_dir)), 

132 ] 

133 if cfg.entity_extraction: 

134 from lilbee.app.status import entity_status 

135 

136 section = entity_status() 

137 if section is not None: 

138 names = ", ".join(section.types) or "schema pending" 

139 lines.append(_kv_line("Entities", f"{section.rows} extracted ({names})")) 

140 return Content("\n").join(lines) 

141 

142 

143def _build_arch_content(info: ModelArchInfo) -> Content: 

144 """Build the model architecture section from GGUF metadata.""" 

145 lines = [ 

146 _kv_line("Chat arch", info.chat_arch), 

147 _kv_line("Embed arch", info.embed_arch), 

148 _kv_line("Handler", pill(info.active_handler, "$accent", "$text")), 

149 ] 

150 if info.vision_projector: 

151 lines.append(_kv_line("Vision proj", info.vision_projector)) 

152 return Content("\n").join(lines) 

153 

154 

155class StatusScreen(Screen[None]): 

156 """Knowledge base status view with collapsible sections.""" 

157 

158 app: LilbeeApp # type: ignore[assignment] 

159 

160 CSS_PATH = "status.tcss" 

161 AUTO_FOCUS = "CollapsibleTitle" 

162 HELP = ( 

163 "Knowledge base status.\n\n" 

164 "View configuration, documents, model architecture, and storage info." 

165 ) 

166 

167 BINDINGS: ClassVar[list[BindingType]] = [ 

168 *browse_back_bindings(), 

169 Binding("tab", "app.focus_next", "Next section", show=False), 

170 Binding( 

171 "shift+tab", 

172 "app.focus_previous", 

173 "Prev section", 

174 show=False, 

175 ), 

176 *BROWSE_LIST_BINDINGS, 

177 ] 

178 

179 def __init__(self) -> None: 

180 super().__init__() 

181 self._sections_mounted: bool = False 

182 self._pending_docs: _DocsResult | None = None 

183 self._pending_arch: ModelArchInfo | None = None 

184 # Sources still waiting to be appended to the table by the batched 

185 # renderer. Bumped each time a new docs result arrives so a stale 

186 # render chain (from a previous load) stops itself. 

187 self._docs_render_queue: list[SourceRecord] = [] 

188 self._docs_render_gen: int = 0 

189 

190 def compose(self) -> ComposeResult: 

191 from textual.widgets import Footer 

192 

193 from lilbee.cli.tui.widgets.bottom_bars import BottomBars 

194 from lilbee.cli.tui.widgets.status_bar import ViewTabs 

195 from lilbee.cli.tui.widgets.task_bar import TaskBar 

196 from lilbee.cli.tui.widgets.top_bars import TopBars 

197 

198 with TopBars(): 

199 yield ViewTabs() 

200 # Mount only the first (Configuration) collapsible up front so the 

201 # screen paints fast on push. Documents/arch/storage hydrate via 

202 # ``call_after_refresh`` once the screen is visible -- their 

203 # backing widgets are still cheap to mount, but the synchronous 

204 # cost of mounting all four under a single VerticalScroll spiked 

205 # screen-switch latency to ~1s on cold caches. 

206 yield VerticalScroll( 

207 Collapsible( 

208 Static(id="config-info"), 

209 title="Configuration", 

210 id="config-section", 

211 collapsed=False, 

212 ), 

213 id="status-scroll", 

214 ) 

215 with BottomBars(): 

216 yield TaskBar() 

217 yield Footer() 

218 

219 def on_mount(self) -> None: 

220 # ``cfg`` reads are in-memory and cheap. Anything that touches 

221 # disk runs in a worker so the screen paints instantly. 

222 # ``get_model_architecture`` opens up to three GGUF files and 

223 # parses their headers (~hundreds of ms each cold); ``get_sources`` 

224 # reads LanceDB (seconds on cold caches). 

225 self._load_config() 

226 self.call_after_refresh(self._mount_remaining_sections) 

227 self._fetch_sources_worker() 

228 self._fetch_arch_worker() 

229 

230 async def _mount_remaining_sections(self) -> None: 

231 """Mount Documents/Architecture/Storage once the screen is visible.""" 

232 if not self.is_mounted: 

233 return 

234 scroll = self.query_one("#status-scroll", VerticalScroll) 

235 await scroll.mount_all( 

236 [ 

237 Collapsible( 

238 DataTable(id="docs-table"), 

239 title=msg.STATUS_DOCS_TITLE, 

240 id="docs-section", 

241 collapsed=False, 

242 ), 

243 Collapsible( 

244 Static(id="arch-info"), 

245 title="Model Architecture", 

246 id="arch-section", 

247 collapsed=False, 

248 ), 

249 Collapsible( 

250 Static(id="storage-info"), 

251 title="Storage", 

252 id="storage-section", 

253 collapsed=False, 

254 ), 

255 ] 

256 ) 

257 self._sections_mounted = True 

258 # Yield once so Textual gets a chance to compose the freshly- 

259 # mounted Collapsibles' children. Without this, querying 

260 # #docs-table immediately after mount_all races on Windows. 

261 await asyncio.sleep(0) 

262 self._show_loading_placeholders() 

263 # Replay any worker callbacks that arrived before the deferred 

264 # mount completed. 

265 if self._pending_docs is not None: 

266 self._apply_docs(self._pending_docs) 

267 self._pending_docs = None 

268 if self._pending_arch is not None: 

269 self._load_arch(self._pending_arch) 

270 self._pending_arch = None 

271 

272 def _show_loading_placeholders(self) -> None: 

273 """Surface a 'Loading…' marker for sections backed by workers. 

274 

275 Wrapped in NoMatches suppression because Collapsible children 

276 compose on the next refresh tick, which on Windows can outlast 

277 the synchronous return from mount_all. Worker callbacks repaint 

278 the same widgets when they arrive, so a missed placeholder is 

279 only a brief cosmetic gap. 

280 """ 

281 from textual.css.query import NoMatches 

282 

283 with contextlib.suppress(NoMatches): 

284 table = self.query_one("#docs-table", DataTable) 

285 table.add_columns("Document", "Chunks") 

286 table.cursor_type = "row" 

287 table.add_row("Loading...", "") 

288 with contextlib.suppress(NoMatches): 

289 self.query_one("#storage-info", Static).update( 

290 Content.styled("Loading...", "$text-muted") 

291 ) 

292 with contextlib.suppress(NoMatches): 

293 self.query_one("#arch-info", Static).update(Content.styled("Loading...", "$text-muted")) 

294 

295 @work(thread=True, name="status_fetch_sources", exit_on_error=False) 

296 def _fetch_sources_worker(self) -> _DocsResult: 

297 """Read the full source list off the UI thread. 

298 

299 ``load_failed`` is True only when the store read actually raised; 

300 an empty store with zero documents is the routine first-run state. 

301 Rendering the (potentially large) list happens back on the UI thread 

302 in bounded batches via :meth:`_render_doc_batch`. 

303 """ 

304 try: 

305 return _DocsResult(sources=get_services().store.get_sources(), load_failed=False) 

306 except Exception: 

307 log.debug("Failed to read store for status screen", exc_info=True) 

308 return _DocsResult(sources=[], load_failed=True) 

309 

310 @work(thread=True, name="status_fetch_arch", exit_on_error=False) 

311 def _fetch_arch_worker(self) -> ModelArchInfo: 

312 try: 

313 return get_model_architecture() 

314 except Exception: 

315 log.debug("Failed to read model architecture for status", exc_info=True) 

316 return ModelArchInfo() 

317 

318 def on_worker_state_changed(self, event: Worker.StateChanged) -> None: 

319 if event.state != WorkerState.SUCCESS: 

320 return 

321 if event.worker.name == "status_fetch_sources": 

322 result = event.worker.result 

323 docs = result if isinstance(result, _DocsResult) else _DocsResult([], True) 

324 if self._sections_mounted: 

325 self._apply_docs(docs) 

326 else: 

327 self._pending_docs = docs 

328 elif event.worker.name == "status_fetch_arch": 

329 arch = event.worker.result 

330 if isinstance(arch, ModelArchInfo): 

331 if self._sections_mounted: 

332 self._load_arch(arch) 

333 else: 

334 self._pending_arch = arch 

335 

336 def _apply_docs(self, docs: _DocsResult) -> None: 

337 """Render *docs* into the Documents table (batched) + storage section.""" 

338 self._load_documents(docs) 

339 self._load_storage(len(docs.sources)) 

340 

341 def _load_arch(self, info: ModelArchInfo) -> None: 

342 """Populate the model architecture section from worker result.""" 

343 from textual.css.query import NoMatches 

344 

345 with contextlib.suppress(NoMatches): 

346 self.query_one("#arch-info", Static).update(_build_arch_content(info)) 

347 

348 def _load_config(self) -> None: 

349 """Populate the configuration section.""" 

350 self.query_one("#config-info", Static).update(_build_config_content()) 

351 

352 def _load_documents(self, docs: _DocsResult) -> None: 

353 """Clear the table, then stream rows in batches over successive refreshes. 

354 

355 Streaming via ``call_after_refresh`` keeps the screen responsive even 

356 with a whole ingested wiki in the store: each batch is small, and the 

357 user can scroll/interact between batches. Suppresses NoMatches because 

358 the deferred Collapsible composes its inner DataTable a refresh tick 

359 after ``mount_all`` returns. 

360 """ 

361 from textual.css.query import NoMatches 

362 

363 with contextlib.suppress(NoMatches): 

364 table = self.query_one("#docs-table", DataTable) 

365 table.clear() 

366 if not docs.sources: 

367 placeholder = ( 

368 msg.STATUS_DOCS_LOAD_FAILED if docs.load_failed else msg.STATUS_DOCS_EMPTY 

369 ) 

370 table.add_row(placeholder, "") 

371 self._docs_render_queue = [] 

372 return 

373 # Bump the generation so any in-flight render chain from a previous 

374 # load stops itself, then kick off a fresh one. 

375 self._docs_render_gen += 1 

376 self._docs_render_queue = list(docs.sources) 

377 self._render_doc_batch(self._docs_render_gen) 

378 

379 def _render_doc_batch(self, generation: int) -> None: 

380 """Append up to ``_DOC_RENDER_BATCH`` rows; reschedule itself if more remain.""" 

381 if generation != self._docs_render_gen or not self._docs_render_queue: 

382 return 

383 from textual.css.query import NoMatches 

384 

385 with contextlib.suppress(NoMatches): 

386 table = self.query_one("#docs-table", DataTable) 

387 batch = self._docs_render_queue[:_DOC_RENDER_BATCH] 

388 del self._docs_render_queue[:_DOC_RENDER_BATCH] 

389 for src in batch: 

390 table.add_row(src.get("filename", "?"), str(src.get("chunk_count", 0))) 

391 if self._docs_render_queue: 

392 self.call_after_refresh(self._render_doc_batch, generation) 

393 

394 def _load_storage(self, doc_count: int) -> None: 

395 """Populate the storage section.""" 

396 from textual.css.query import NoMatches 

397 

398 with contextlib.suppress(NoMatches): 

399 self.query_one("#storage-info", Static).update(_build_storage_content(doc_count)) 

400 

401 def action_go_back(self) -> None: 

402 self.app.go_back() 

403 

404 def action_cursor_down(self) -> None: 

405 self.query_one("#status-scroll", VerticalScroll).scroll_down() 

406 

407 def action_cursor_up(self) -> None: 

408 self.query_one("#status-scroll", VerticalScroll).scroll_up() 

409 

410 def action_jump_top(self) -> None: 

411 self.query_one("#status-scroll", VerticalScroll).scroll_home(animate=False) 

412 

413 def action_jump_bottom(self) -> None: 

414 self.query_one("#status-scroll", VerticalScroll).scroll_end(animate=False)