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
« 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."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import logging
8from dataclasses import dataclass
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.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
32log = logging.getLogger(__name__)
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
41@dataclass
42class _DocsResult:
43 """Outcome of the background sources read.
45 ``load_failed`` distinguishes "store opened but empty" (``sources == []``)
46 from "the read raised" so the UI shows the right placeholder.
47 """
49 sources: list[SourceRecord]
50 load_failed: bool
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")
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
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)
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
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"
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")
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")
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)
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
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)
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)
155class StatusScreen(Screen[None]):
156 """Knowledge base status view with collapsible sections."""
158 app: LilbeeApp # type: ignore[assignment]
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 )
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 ]
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
190 def compose(self) -> ComposeResult:
191 from textual.widgets import Footer
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
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()
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()
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
272 def _show_loading_placeholders(self) -> None:
273 """Surface a 'Loading…' marker for sections backed by workers.
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
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"))
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.
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)
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()
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
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))
341 def _load_arch(self, info: ModelArchInfo) -> None:
342 """Populate the model architecture section from worker result."""
343 from textual.css.query import NoMatches
345 with contextlib.suppress(NoMatches):
346 self.query_one("#arch-info", Static).update(_build_arch_content(info))
348 def _load_config(self) -> None:
349 """Populate the configuration section."""
350 self.query_one("#config-info", Static).update(_build_config_content())
352 def _load_documents(self, docs: _DocsResult) -> None:
353 """Clear the table, then stream rows in batches over successive refreshes.
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
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)
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
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)
394 def _load_storage(self, doc_count: int) -> None:
395 """Populate the storage section."""
396 from textual.css.query import NoMatches
398 with contextlib.suppress(NoMatches):
399 self.query_one("#storage-info", Static).update(_build_storage_content(doc_count))
401 def action_go_back(self) -> None:
402 self.app.go_back()
404 def action_cursor_down(self) -> None:
405 self.query_one("#status-scroll", VerticalScroll).scroll_down()
407 def action_cursor_up(self) -> None:
408 self.query_one("#status-scroll", VerticalScroll).scroll_up()
410 def action_jump_top(self) -> None:
411 self.query_one("#status-scroll", VerticalScroll).scroll_home(animate=False)
413 def action_jump_bottom(self) -> None:
414 self.query_one("#status-scroll", VerticalScroll).scroll_end(animate=False)