Coverage for src/lilbee/cli/tui/screens/status.py: 100%
233 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-04 17:08 +0000
1"""Status screen: knowledge base info with collapsible sections."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import logging
8from dataclasses import dataclass, field
9from pathlib import Path
10from typing import TYPE_CHECKING, ClassVar
12if TYPE_CHECKING:
13 from lilbee.cli.tui.app import LilbeeApp
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
24from lilbee.app.services import get_services
25from lilbee.app.status import held_out_sources
26from lilbee.cli.tui import messages as msg
27from lilbee.cli.tui.browse_bindings import BROWSE_LIST_BINDINGS, browse_back_bindings
28from lilbee.cli.tui.pill import pill
29from lilbee.core.config import cfg
30from lilbee.data.store import SourceRecord
31from lilbee.data.types import SkippedSource
32from lilbee.modelhub.model_info import ModelArchInfo, get_model_architecture
34log = logging.getLogger(__name__)
36# Rows appended to the Documents table per refresh tick. Each ``add_row`` runs
37# on the UI thread, so dumping a whole ingested wiki in one loop froze the
38# screen for seconds; rendering a bounded batch per ``call_after_refresh`` lets
39# the user scroll and interact while the rest streams in.
40_DOC_RENDER_BATCH = 100
43@dataclass
44class _DocsResult:
45 """Outcome of the background sources read.
47 ``load_failed`` distinguishes "store opened but empty" (``sources == []``)
48 from "the read raised" so the UI shows the right placeholder.
49 """
51 sources: list[SourceRecord]
52 load_failed: bool
53 held_out: list[SkippedSource] = field(default_factory=list)
54 held_out_total: int = 0
57def _model_pill(name: str) -> Content:
58 """Return a green 'loaded' pill if name is set, red 'not set' otherwise."""
59 if name:
60 return pill("loaded", "$success", "$text")
61 return pill("not set", "$error", "$text")
64# Label-column width used across the status sections so keys line up
65# when scanned vertically. Values past this column render bold.
66_KV_LABEL_WIDTH = 14
69def _kv_line(label: str, value: str | Content, status: Content | None = None) -> Content:
70 """Assemble one key/value row: dim padded label, bold value, optional pill."""
71 padded = label.ljust(_KV_LABEL_WIDTH)
72 parts: list[Content] = [Content.styled(padded, "$text-muted")]
73 if isinstance(value, Content):
74 parts.append(value)
75 else:
76 parts.append(Content.styled(value, "bold"))
77 if status is not None:
78 parts.append(Content(" "))
79 parts.append(status)
80 return Content.assemble(*parts)
83def _collapse_home(path: Path | str) -> str:
84 """Replace the user's home prefix with '~' so long paths stay scannable."""
85 text = str(path)
86 home = str(Path.home())
87 return text.replace(home, "~", 1) if text.startswith(home) else text
90def _ocr_label() -> str:
91 """Return a human-readable OCR status string."""
92 if cfg.enable_ocr is True:
93 return "enabled"
94 if cfg.enable_ocr is False:
95 return "disabled"
96 return "auto"
99def _ocr_pill() -> Content:
100 """Return a pill reflecting OCR status."""
101 if cfg.enable_ocr is True:
102 return pill("on", "$success", "$text")
103 if cfg.enable_ocr is False:
104 return pill("off", "$warning", "$text")
105 return pill("auto", "$accent", "$text")
108def _data_dir_pill() -> Content:
109 """Return a pill based on whether the data directory exists."""
110 if Path(cfg.data_dir).exists():
111 return pill("exists", "$success", "$text")
112 return pill("missing", "$error", "$text")
115def _build_config_content() -> Content:
116 """Build the configuration section content."""
117 lines = [
118 _kv_line("Data dir", _collapse_home(cfg.data_dir), _data_dir_pill()),
119 _kv_line("Chat model", cfg.chat_model or "(disabled)", _model_pill(cfg.chat_model)),
120 _kv_line(
121 "Embed model", cfg.embedding_model or "(disabled)", _model_pill(cfg.embedding_model)
122 ),
123 _kv_line("Vision model", cfg.vision_model or "(disabled)", _model_pill(cfg.vision_model)),
124 _kv_line("Reranker", cfg.reranker_model or "(disabled)", _model_pill(cfg.reranker_model)),
125 _kv_line("OCR", _ocr_label(), _ocr_pill()),
126 ]
127 return Content("\n").join(lines)
130def _build_storage_content(doc_count: int) -> Content:
131 """Build the storage section content."""
132 lines = [
133 _kv_line("Documents", str(doc_count)),
134 _kv_line("Data dir", _collapse_home(cfg.data_dir)),
135 _kv_line("Models dir", _collapse_home(cfg.models_dir)),
136 ]
137 if cfg.entity_extraction:
138 from lilbee.app.status import entity_status
140 section = entity_status()
141 if section is not None:
142 names = ", ".join(section.types) or "schema pending"
143 lines.append(_kv_line("Entities", f"{section.rows} extracted ({names})"))
144 return Content("\n").join(lines)
147def _build_arch_content(info: ModelArchInfo) -> Content:
148 """Build the model architecture section from GGUF metadata."""
149 lines = [
150 _kv_line("Chat arch", info.chat_arch),
151 _kv_line("Embed arch", info.embed_arch),
152 _kv_line("Handler", pill(info.active_handler, "$accent", "$text")),
153 ]
154 if info.vision_projector:
155 lines.append(_kv_line("Vision proj", info.vision_projector))
156 return Content("\n").join(lines)
159class StatusScreen(Screen[None]):
160 """Knowledge base status view with collapsible sections."""
162 app: LilbeeApp # type: ignore[assignment]
164 CSS_PATH = "status.tcss"
165 AUTO_FOCUS = "CollapsibleTitle"
166 HELP = (
167 "Knowledge base status.\n\n"
168 "View configuration, documents, model architecture, and storage info."
169 )
171 BINDINGS: ClassVar[list[BindingType]] = [
172 *browse_back_bindings(),
173 Binding("tab", "app.focus_next", "Next section", show=False),
174 Binding(
175 "shift+tab",
176 "app.focus_previous",
177 "Prev section",
178 show=False,
179 ),
180 *BROWSE_LIST_BINDINGS,
181 ]
183 def __init__(self) -> None:
184 super().__init__()
185 self._sections_mounted: bool = False
186 self._pending_docs: _DocsResult | None = None
187 self._pending_arch: ModelArchInfo | None = None
188 # Sources still waiting to be appended to the table by the batched
189 # renderer. Bumped each time a new docs result arrives so a stale
190 # render chain (from a previous load) stops itself.
191 self._docs_render_queue: list[SourceRecord] = []
192 self._docs_render_gen: int = 0
194 def compose(self) -> ComposeResult:
195 from textual.widgets import Footer
197 from lilbee.cli.tui.widgets.bottom_bars import BottomBars
198 from lilbee.cli.tui.widgets.status_bar import ViewTabs
199 from lilbee.cli.tui.widgets.task_bar import TaskBar
200 from lilbee.cli.tui.widgets.top_bars import TopBars
202 with TopBars():
203 yield ViewTabs()
204 # Mount only the first (Configuration) collapsible up front so the
205 # screen paints fast on push. Documents/arch/storage hydrate via
206 # ``call_after_refresh`` once the screen is visible -- their
207 # backing widgets are still cheap to mount, but the synchronous
208 # cost of mounting all four under a single VerticalScroll spiked
209 # screen-switch latency to ~1s on cold caches.
210 yield VerticalScroll(
211 Collapsible(
212 Static(id="config-info"),
213 title="Configuration",
214 id="config-section",
215 collapsed=False,
216 ),
217 id="status-scroll",
218 )
219 with BottomBars():
220 yield TaskBar()
221 yield Footer()
223 def on_mount(self) -> None:
224 # ``cfg`` reads are in-memory and cheap. Anything that touches
225 # disk runs in a worker so the screen paints instantly.
226 # ``get_model_architecture`` opens up to three GGUF files and
227 # parses their headers (~hundreds of ms each cold); ``get_sources``
228 # reads LanceDB (seconds on cold caches).
229 self._load_config()
230 self.call_after_refresh(self._mount_remaining_sections)
231 self._fetch_sources_worker()
232 self._fetch_arch_worker()
234 async def _mount_remaining_sections(self) -> None:
235 """Mount Documents/Architecture/Storage once the screen is visible."""
236 if not self.is_mounted:
237 return
238 scroll = self.query_one("#status-scroll", VerticalScroll)
239 await scroll.mount_all(
240 [
241 Collapsible(
242 DataTable(id="docs-table"),
243 title=msg.STATUS_DOCS_TITLE,
244 id="docs-section",
245 collapsed=False,
246 ),
247 Collapsible(
248 DataTable(id="held-out-table"),
249 title=msg.STATUS_HELD_OUT_TITLE,
250 id="held-out-section",
251 collapsed=True,
252 ),
253 Collapsible(
254 Static(id="arch-info"),
255 title="Model Architecture",
256 id="arch-section",
257 collapsed=False,
258 ),
259 Collapsible(
260 Static(id="storage-info"),
261 title="Storage",
262 id="storage-section",
263 collapsed=False,
264 ),
265 ]
266 )
267 self._sections_mounted = True
268 # Yield once so Textual gets a chance to compose the freshly-
269 # mounted Collapsibles' children. Without this, querying
270 # #docs-table immediately after mount_all races on Windows.
271 await asyncio.sleep(0)
272 self._show_loading_placeholders()
273 # Replay any worker callbacks that arrived before the deferred
274 # mount completed.
275 if self._pending_docs is not None:
276 self._apply_docs(self._pending_docs)
277 self._pending_docs = None
278 if self._pending_arch is not None:
279 self._load_arch(self._pending_arch)
280 self._pending_arch = None
282 def _show_loading_placeholders(self) -> None:
283 """Surface a 'Loading…' marker for sections backed by workers.
285 Wrapped in NoMatches suppression because Collapsible children
286 compose on the next refresh tick, which on Windows can outlast
287 the synchronous return from mount_all. Worker callbacks repaint
288 the same widgets when they arrive, so a missed placeholder is
289 only a brief cosmetic gap.
290 """
291 from textual.css.query import NoMatches
293 with contextlib.suppress(NoMatches):
294 table = self.query_one("#docs-table", DataTable)
295 table.add_columns("Document", "Chunks")
296 table.cursor_type = "row"
297 table.add_row("Loading...", "")
298 with contextlib.suppress(NoMatches):
299 self.query_one("#storage-info", Static).update(
300 Content.styled("Loading...", "$text-muted")
301 )
302 with contextlib.suppress(NoMatches):
303 self.query_one("#arch-info", Static).update(Content.styled("Loading...", "$text-muted"))
305 @work(thread=True, name="status_fetch_sources", exit_on_error=False)
306 def _fetch_sources_worker(self) -> _DocsResult:
307 """Read the full source list off the UI thread.
309 ``load_failed`` is True only when the store read actually raised;
310 an empty store with zero documents is the routine first-run state.
311 Rendering the (potentially large) list happens back on the UI thread
312 in bounded batches via :meth:`_render_doc_batch`.
313 """
314 try:
315 held_out, held_out_total = held_out_sources()
316 return _DocsResult(
317 sources=get_services().store.get_sources(),
318 load_failed=False,
319 held_out=held_out,
320 held_out_total=held_out_total,
321 )
322 except Exception:
323 log.debug("Failed to read store for status screen", exc_info=True)
324 return _DocsResult(sources=[], load_failed=True)
326 @work(thread=True, name="status_fetch_arch", exit_on_error=False)
327 def _fetch_arch_worker(self) -> ModelArchInfo:
328 try:
329 return get_model_architecture()
330 except Exception:
331 log.debug("Failed to read model architecture for status", exc_info=True)
332 return ModelArchInfo()
334 def on_worker_state_changed(self, event: Worker.StateChanged) -> None:
335 if event.state != WorkerState.SUCCESS:
336 return
337 if event.worker.name == "status_fetch_sources":
338 result = event.worker.result
339 docs = result if isinstance(result, _DocsResult) else _DocsResult([], True)
340 if self._sections_mounted:
341 self._apply_docs(docs)
342 else:
343 self._pending_docs = docs
344 elif event.worker.name == "status_fetch_arch":
345 arch = event.worker.result
346 if isinstance(arch, ModelArchInfo):
347 if self._sections_mounted:
348 self._load_arch(arch)
349 else:
350 self._pending_arch = arch
352 def _apply_docs(self, docs: _DocsResult) -> None:
353 """Render *docs* into the Documents table (batched) + storage section."""
354 self._load_documents(docs)
355 self._load_held_out(docs)
356 self._load_storage(len(docs.sources))
358 def _load_held_out(self, docs: _DocsResult) -> None:
359 """Fill the held-out table; the section opens only when something is held out."""
360 from textual.css.query import NoMatches
362 with contextlib.suppress(NoMatches):
363 table = self.query_one("#held-out-table", DataTable)
364 table.clear(columns=True)
365 table.add_columns("File", "Reason")
366 if not docs.held_out:
367 table.add_row(msg.STATUS_HELD_OUT_EMPTY, "")
368 return
369 for held in docs.held_out:
370 table.add_row(held.filename, held.reason)
371 hidden = docs.held_out_total - len(docs.held_out)
372 if hidden > 0:
373 table.add_row(msg.STATUS_HELD_OUT_MORE.format(count=hidden), "")
374 self.query_one("#held-out-section", Collapsible).collapsed = False
376 def _load_arch(self, info: ModelArchInfo) -> None:
377 """Populate the model architecture section from worker result."""
378 from textual.css.query import NoMatches
380 with contextlib.suppress(NoMatches):
381 self.query_one("#arch-info", Static).update(_build_arch_content(info))
383 def _load_config(self) -> None:
384 """Populate the configuration section."""
385 self.query_one("#config-info", Static).update(_build_config_content())
387 def _load_documents(self, docs: _DocsResult) -> None:
388 """Clear the table, then stream rows in batches over successive refreshes.
390 Streaming via ``call_after_refresh`` keeps the screen responsive even
391 with a whole ingested wiki in the store: each batch is small, and the
392 user can scroll/interact between batches. Suppresses NoMatches because
393 the deferred Collapsible composes its inner DataTable a refresh tick
394 after ``mount_all`` returns.
395 """
396 from textual.css.query import NoMatches
398 with contextlib.suppress(NoMatches):
399 table = self.query_one("#docs-table", DataTable)
400 table.clear()
401 if not docs.sources:
402 placeholder = (
403 msg.STATUS_DOCS_LOAD_FAILED if docs.load_failed else msg.STATUS_DOCS_EMPTY
404 )
405 table.add_row(placeholder, "")
406 self._docs_render_queue = []
407 return
408 # Bump the generation so any in-flight render chain from a previous
409 # load stops itself, then kick off a fresh one.
410 self._docs_render_gen += 1
411 self._docs_render_queue = list(docs.sources)
412 self._render_doc_batch(self._docs_render_gen)
414 def _render_doc_batch(self, generation: int) -> None:
415 """Append up to ``_DOC_RENDER_BATCH`` rows; reschedule itself if more remain."""
416 if generation != self._docs_render_gen or not self._docs_render_queue:
417 return
418 from textual.css.query import NoMatches
420 with contextlib.suppress(NoMatches):
421 table = self.query_one("#docs-table", DataTable)
422 batch = self._docs_render_queue[:_DOC_RENDER_BATCH]
423 del self._docs_render_queue[:_DOC_RENDER_BATCH]
424 for src in batch:
425 table.add_row(src.get("filename", "?"), str(src.get("chunk_count", 0)))
426 if self._docs_render_queue:
427 self.call_after_refresh(self._render_doc_batch, generation)
429 def _load_storage(self, doc_count: int) -> None:
430 """Populate the storage section."""
431 from textual.css.query import NoMatches
433 with contextlib.suppress(NoMatches):
434 self.query_one("#storage-info", Static).update(_build_storage_content(doc_count))
436 def action_go_back(self) -> None:
437 self.app.go_back()
439 def action_cursor_down(self) -> None:
440 self.query_one("#status-scroll", VerticalScroll).scroll_down()
442 def action_cursor_up(self) -> None:
443 self.query_one("#status-scroll", VerticalScroll).scroll_up()
445 def action_jump_top(self) -> None:
446 self.query_one("#status-scroll", VerticalScroll).scroll_home(animate=False)
448 def action_jump_bottom(self) -> None:
449 self.query_one("#status-scroll", VerticalScroll).scroll_end(animate=False)