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