Coverage for src / lilbee / cli / tui / screens / status.py: 100%
205 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-05-15 20:55 +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.pill import pill
27from lilbee.core.config import cfg
28from lilbee.data.store import SourceRecord
29from lilbee.modelhub.model_info import ModelArchInfo, get_model_architecture
31log = logging.getLogger(__name__)
33# Rows appended to the Documents table per refresh tick. Each ``add_row`` runs
34# on the UI thread, so dumping a whole ingested wiki in one loop froze the
35# screen for seconds; rendering a bounded batch per ``call_after_refresh`` lets
36# the user scroll and interact while the rest streams in.
37_DOC_RENDER_BATCH = 100
40@dataclass
41class _DocsResult:
42 """Outcome of the background sources read.
44 ``load_failed`` distinguishes "store opened but empty" (``sources == []``)
45 from "the read raised" so the UI shows the right placeholder.
46 """
48 sources: list[SourceRecord]
49 load_failed: bool
52def _model_pill(name: str) -> Content:
53 """Return a green 'loaded' pill if name is set, red 'not set' otherwise."""
54 if name:
55 return pill("loaded", "$success", "$text")
56 return pill("not set", "$error", "$text")
59# Label-column width used across the status sections so keys line up
60# when scanned vertically. Values past this column render bold.
61_KV_LABEL_WIDTH = 14
64def _kv_line(label: str, value: str | Content, status: Content | None = None) -> Content:
65 """Assemble one key/value row: dim padded label, bold value, optional pill."""
66 padded = label.ljust(_KV_LABEL_WIDTH)
67 parts: list[Content] = [Content.styled(padded, "$text-muted")]
68 if isinstance(value, Content):
69 parts.append(value)
70 else:
71 parts.append(Content.styled(value, "bold"))
72 if status is not None:
73 parts.append(Content(" "))
74 parts.append(status)
75 return Content.assemble(*parts)
78def _collapse_home(path: Path | str) -> str:
79 """Replace the user's home prefix with '~' so long paths stay scannable."""
80 text = str(path)
81 home = str(Path.home())
82 return text.replace(home, "~", 1) if text.startswith(home) else text
85def _ocr_label() -> str:
86 """Return a human-readable OCR status string."""
87 if cfg.enable_ocr is True:
88 return "enabled"
89 if cfg.enable_ocr is False:
90 return "disabled"
91 return "auto"
94def _ocr_pill() -> Content:
95 """Return a pill reflecting OCR status."""
96 if cfg.enable_ocr is True:
97 return pill("on", "$success", "$text")
98 if cfg.enable_ocr is False:
99 return pill("off", "$warning", "$text")
100 return pill("auto", "$accent", "$text")
103def _data_dir_pill() -> Content:
104 """Return a pill based on whether the data directory exists."""
105 if Path(cfg.data_dir).exists():
106 return pill("exists", "$success", "$text")
107 return pill("missing", "$error", "$text")
110def _build_config_content() -> Content:
111 """Build the configuration section content."""
112 lines = [
113 _kv_line("Data dir", _collapse_home(cfg.data_dir), _data_dir_pill()),
114 _kv_line("Chat model", cfg.chat_model or "(disabled)", _model_pill(cfg.chat_model)),
115 _kv_line(
116 "Embed model", cfg.embedding_model or "(disabled)", _model_pill(cfg.embedding_model)
117 ),
118 _kv_line("Vision model", cfg.vision_model or "(disabled)", _model_pill(cfg.vision_model)),
119 _kv_line("Reranker", cfg.reranker_model or "(disabled)", _model_pill(cfg.reranker_model)),
120 _kv_line("OCR", _ocr_label(), _ocr_pill()),
121 ]
122 return Content("\n").join(lines)
125def _build_storage_content(doc_count: int) -> Content:
126 """Build the storage section content."""
127 lines = [
128 _kv_line("Documents", str(doc_count)),
129 _kv_line("Data dir", _collapse_home(cfg.data_dir)),
130 _kv_line("Models dir", _collapse_home(cfg.models_dir)),
131 ]
132 return Content("\n").join(lines)
135def _build_arch_content(info: ModelArchInfo) -> Content:
136 """Build the model architecture section from GGUF metadata."""
137 lines = [
138 _kv_line("Chat arch", info.chat_arch),
139 _kv_line("Embed arch", info.embed_arch),
140 _kv_line("Handler", pill(info.active_handler, "$accent", "$text")),
141 ]
142 if info.vision_projector:
143 lines.append(_kv_line("Vision proj", info.vision_projector))
144 return Content("\n").join(lines)
147class StatusScreen(Screen[None]):
148 """Knowledge base status view with collapsible sections."""
150 app: LilbeeApp # type: ignore[assignment]
152 CSS_PATH = "status.tcss"
153 AUTO_FOCUS = "CollapsibleTitle"
154 HELP = (
155 "Knowledge base status.\n\n"
156 "View configuration, documents, model architecture, and storage info."
157 )
159 BINDINGS: ClassVar[list[BindingType]] = [
160 Binding("q", "go_back", "Back", show=True),
161 Binding("escape", "go_back", "Back", show=False),
162 Binding("tab", "app.focus_next", "Next section", show=True),
163 Binding("shift+tab", "app.focus_previous", "Prev section", show=True),
164 Binding("j", "cursor_down", "Nav", show=False),
165 Binding("k", "cursor_up", "Nav", show=False),
166 Binding("g", "jump_top", "Top", show=False),
167 Binding("G", "jump_bottom", "End", show=False),
168 ]
170 def __init__(self) -> None:
171 super().__init__()
172 self._sections_mounted: bool = False
173 self._pending_docs: _DocsResult | None = None
174 self._pending_arch: ModelArchInfo | None = None
175 # Sources still waiting to be appended to the table by the batched
176 # renderer. Bumped each time a new docs result arrives so a stale
177 # render chain (from a previous load) stops itself.
178 self._docs_render_queue: list[SourceRecord] = []
179 self._docs_render_gen: int = 0
181 def compose(self) -> ComposeResult:
182 from textual.widgets import Footer
184 from lilbee.cli.tui.widgets.bottom_bars import BottomBars
185 from lilbee.cli.tui.widgets.status_bar import ViewTabs
186 from lilbee.cli.tui.widgets.task_bar import TaskBar
187 from lilbee.cli.tui.widgets.top_bars import TopBars
189 with TopBars():
190 yield ViewTabs()
191 # Mount only the first (Configuration) collapsible up front so the
192 # screen paints fast on push. Documents/arch/storage hydrate via
193 # ``call_after_refresh`` once the screen is visible -- their
194 # backing widgets are still cheap to mount, but the synchronous
195 # cost of mounting all four under a single VerticalScroll spiked
196 # screen-switch latency to ~1s on cold caches.
197 yield VerticalScroll(
198 Collapsible(Static(id="config-info"), title="Configuration", id="config-section"),
199 id="status-scroll",
200 )
201 with BottomBars():
202 yield TaskBar()
203 yield Footer()
205 def on_mount(self) -> None:
206 # ``cfg`` reads are in-memory and cheap. Anything that touches
207 # disk runs in a worker so the screen paints instantly.
208 # ``get_model_architecture`` opens up to three GGUF files and
209 # parses their headers (~hundreds of ms each cold); ``get_sources``
210 # reads LanceDB (seconds on cold caches).
211 self._load_config()
212 self.call_after_refresh(self._mount_remaining_sections)
213 self._fetch_sources_worker()
214 self._fetch_arch_worker()
216 async def _mount_remaining_sections(self) -> None:
217 """Mount Documents/Architecture/Storage once the screen is visible."""
218 if not self.is_mounted:
219 return
220 scroll = self.query_one("#status-scroll", VerticalScroll)
221 await scroll.mount_all(
222 [
223 Collapsible(
224 DataTable(id="docs-table"), title=msg.STATUS_DOCS_TITLE, id="docs-section"
225 ),
226 Collapsible(Static(id="arch-info"), title="Model Architecture", id="arch-section"),
227 Collapsible(Static(id="storage-info"), title="Storage", id="storage-section"),
228 ]
229 )
230 self._sections_mounted = True
231 # Yield once so Textual gets a chance to compose the freshly-
232 # mounted Collapsibles' children. Without this, querying
233 # #docs-table immediately after mount_all races on Windows.
234 await asyncio.sleep(0)
235 self._show_loading_placeholders()
236 # Replay any worker callbacks that arrived before the deferred
237 # mount completed.
238 if self._pending_docs is not None:
239 self._apply_docs(self._pending_docs)
240 self._pending_docs = None
241 if self._pending_arch is not None:
242 self._load_arch(self._pending_arch)
243 self._pending_arch = None
245 def _show_loading_placeholders(self) -> None:
246 """Surface a 'Loading…' marker for sections backed by workers.
248 Wrapped in NoMatches suppression because Collapsible children
249 compose on the next refresh tick, which on Windows can outlast
250 the synchronous return from mount_all. Worker callbacks repaint
251 the same widgets when they arrive, so a missed placeholder is
252 only a brief cosmetic gap.
253 """
254 from textual.css.query import NoMatches
256 with contextlib.suppress(NoMatches):
257 table = self.query_one("#docs-table", DataTable)
258 table.add_columns("Document", "Chunks")
259 table.cursor_type = "row"
260 table.add_row("Loading...", "")
261 with contextlib.suppress(NoMatches):
262 self.query_one("#storage-info", Static).update(
263 Content.styled("Loading...", "$text-muted")
264 )
265 with contextlib.suppress(NoMatches):
266 self.query_one("#arch-info", Static).update(Content.styled("Loading...", "$text-muted"))
268 @work(thread=True, name="status_fetch_sources", exit_on_error=False)
269 def _fetch_sources_worker(self) -> _DocsResult:
270 """Read the full source list off the UI thread.
272 ``load_failed`` is True only when the store read actually raised;
273 an empty store with zero documents is the routine first-run state.
274 Rendering the (potentially large) list happens back on the UI thread
275 in bounded batches via :meth:`_render_doc_batch`.
276 """
277 try:
278 return _DocsResult(sources=get_services().store.get_sources(), load_failed=False)
279 except Exception:
280 log.debug("Failed to read store for status screen", exc_info=True)
281 return _DocsResult(sources=[], load_failed=True)
283 @work(thread=True, name="status_fetch_arch", exit_on_error=False)
284 def _fetch_arch_worker(self) -> ModelArchInfo:
285 try:
286 return get_model_architecture()
287 except Exception:
288 log.debug("Failed to read model architecture for status", exc_info=True)
289 return ModelArchInfo()
291 def on_worker_state_changed(self, event: Worker.StateChanged) -> None:
292 if event.state != WorkerState.SUCCESS:
293 return
294 if event.worker.name == "status_fetch_sources":
295 result = event.worker.result
296 docs = result if isinstance(result, _DocsResult) else _DocsResult([], True)
297 if self._sections_mounted:
298 self._apply_docs(docs)
299 else:
300 self._pending_docs = docs
301 elif event.worker.name == "status_fetch_arch":
302 arch = event.worker.result
303 if isinstance(arch, ModelArchInfo):
304 if self._sections_mounted:
305 self._load_arch(arch)
306 else:
307 self._pending_arch = arch
309 def _apply_docs(self, docs: _DocsResult) -> None:
310 """Render *docs* into the Documents table (batched) + storage section."""
311 self._load_documents(docs)
312 self._load_storage(len(docs.sources))
314 def _load_arch(self, info: ModelArchInfo) -> None:
315 """Populate the model architecture section from worker result."""
316 from textual.css.query import NoMatches
318 with contextlib.suppress(NoMatches):
319 self.query_one("#arch-info", Static).update(_build_arch_content(info))
321 def _load_config(self) -> None:
322 """Populate the configuration section."""
323 self.query_one("#config-info", Static).update(_build_config_content())
325 def _load_documents(self, docs: _DocsResult) -> None:
326 """Clear the table, then stream rows in batches over successive refreshes.
328 Streaming via ``call_after_refresh`` keeps the screen responsive even
329 with a whole ingested wiki in the store: each batch is small, and the
330 user can scroll/interact between batches. Suppresses NoMatches because
331 the deferred Collapsible composes its inner DataTable a refresh tick
332 after ``mount_all`` returns.
333 """
334 from textual.css.query import NoMatches
336 with contextlib.suppress(NoMatches):
337 table = self.query_one("#docs-table", DataTable)
338 table.clear()
339 if not docs.sources:
340 placeholder = (
341 msg.STATUS_DOCS_LOAD_FAILED if docs.load_failed else msg.STATUS_DOCS_EMPTY
342 )
343 table.add_row(placeholder, "")
344 self._docs_render_queue = []
345 return
346 # Bump the generation so any in-flight render chain from a previous
347 # load stops itself, then kick off a fresh one.
348 self._docs_render_gen += 1
349 self._docs_render_queue = list(docs.sources)
350 self._render_doc_batch(self._docs_render_gen)
352 def _render_doc_batch(self, generation: int) -> None:
353 """Append up to ``_DOC_RENDER_BATCH`` rows; reschedule itself if more remain."""
354 if generation != self._docs_render_gen or not self._docs_render_queue:
355 return
356 from textual.css.query import NoMatches
358 with contextlib.suppress(NoMatches):
359 table = self.query_one("#docs-table", DataTable)
360 batch = self._docs_render_queue[:_DOC_RENDER_BATCH]
361 del self._docs_render_queue[:_DOC_RENDER_BATCH]
362 for src in batch:
363 table.add_row(src.get("filename", "?"), str(src.get("chunk_count", 0)))
364 if self._docs_render_queue:
365 self.call_after_refresh(self._render_doc_batch, generation)
367 def _load_storage(self, doc_count: int) -> None:
368 """Populate the storage section."""
369 from textual.css.query import NoMatches
371 with contextlib.suppress(NoMatches):
372 self.query_one("#storage-info", Static).update(_build_storage_content(doc_count))
374 def action_go_back(self) -> None:
375 self.app.switch_view("Chat")
377 def action_cursor_down(self) -> None:
378 self.query_one("#status-scroll", VerticalScroll).scroll_down()
380 def action_cursor_up(self) -> None:
381 self.query_one("#status-scroll", VerticalScroll).scroll_up()
383 def action_jump_top(self) -> None:
384 self.query_one("#status-scroll", VerticalScroll).scroll_home()
386 def action_jump_bottom(self) -> None:
387 self.query_one("#status-scroll", VerticalScroll).scroll_end()