Coverage for src/lilbee/cli/helpers.py: 100%
176 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +0000
1"""CLI-specific helpers: JSON formatter, Rich rendering, and CLI workflows."""
3from __future__ import annotations
5import asyncio
6import json
7import signal
8import threading
9from collections.abc import Callable, Generator, Iterator
10from contextlib import contextmanager
11from pathlib import Path
12from typing import TYPE_CHECKING
14from rich.console import RenderableType
15from rich.table import Table
16from rich.text import Text
18from lilbee.app.ingest import RegisterResult, register_sources
19from lilbee.app.status import StatusResult
20from lilbee.cli import theme
21from lilbee.core.config import cfg
22from lilbee.runtime.console import PlainConsole, styled
24if TYPE_CHECKING:
25 from lilbee.cli.sync import SyncStatus
28def json_output(data: dict) -> None:
29 """Print a JSON object to stdout."""
30 print(json.dumps(data))
33def announce_cold_start(role: object, model: str) -> PlainConsole | None:
34 """Print a "Starting <role> engine (loading <model>)..." stderr line if cold.
36 Returns a stderr console to print the matching "ready" line through when the
37 blocking call returns, or ``None`` when the role's server is already warm (no
38 status needed) or output is JSON (machine-readable, no chatter). The role
39 parameter is a ``WorkerRole``; typed as ``object`` to keep this CLI helper
40 free of a provider-layer import at module top.
41 """
42 from lilbee.app.services import get_services
43 from lilbee.providers.roles import WorkerRole
45 if cfg.json_mode or not isinstance(role, WorkerRole):
46 return None
47 if get_services().provider.role_ready(role):
48 return None
49 err = PlainConsole(stderr=True)
50 err.print(
51 Text.assemble(
52 (f"Starting {role.value} engine (loading ", theme.MUTED),
53 model,
54 (")...", theme.MUTED),
55 ),
56 soft_wrap=True,
57 )
58 return err
61def announce_ready(err: PlainConsole | None, role: object) -> None:
62 """Print the matching "<role> engine ready." stderr line, if cold-start announced.
64 A token arriving is not evidence the chat model came up: in RAG mode a grounded
65 refusal streams without it. When warm-up recorded a load failure, that reason is
66 printed instead of a readiness line.
67 """
68 from lilbee.providers.roles import WorkerRole
70 if err is None or not isinstance(role, WorkerRole):
71 return
72 failure = _chat_warm_error(role)
73 if failure is not None:
74 err.print(failure, style=theme.ERROR, soft_wrap=True)
75 return
76 err.print(f"{role.value} engine ready.", style=theme.MUTED)
79def announce_retrieval_query(query: str) -> None:
80 """Print the "Searching for: <query>" stderr line for a rewritten follow-up."""
81 line = SEARCHING_FOR.format(query=query)
82 PlainConsole(stderr=True).print(line, style=theme.MUTED, soft_wrap=True)
85def _chat_warm_error(role: object) -> str | None:
86 """The chat warm-up's recorded failure, or None when it did not fail.
88 Read from the warm tracker rather than re-probing readiness: llama-swap can
89 report a freshly loaded model as not-yet-running, which would turn a healthy
90 engine into a spurious failure line.
91 """
92 from lilbee.app.services import get_services
93 from lilbee.providers.roles import WorkerRole
94 from lilbee.providers.warm_progress import WarmPhase
96 if role is not WorkerRole.CHAT:
97 return None
98 snapshot = get_services().provider.warm_progress()
99 if snapshot is None or snapshot.phase is not WarmPhase.ERROR:
100 return None
101 return snapshot.error or "The chat model did not finish loading."
104_LABEL_WIDTH = len("Chat model:")
107def _label_line(label: str, value: object) -> Text:
108 """An aligned ``Label: value`` line with *value* as literal text."""
109 gap = " " * max(1, _LABEL_WIDTH - len(label))
110 return Text.assemble((f"{label}:", theme.LABEL), gap, str(value))
113def render_status_result(status: StatusResult) -> Generator[RenderableType, None, None]:
114 """Yield Rich renderables for a :class:`StatusResult`."""
115 yield _label_line("Documents", status.config.documents_dir)
116 yield _label_line("Database", status.config.data_dir)
117 yield _label_line("Chat model", status.config.chat_model)
118 yield _label_line("Embeddings", status.config.embedding_model)
119 if status.index is not None:
120 yield _label_line(
121 "Index built with",
122 f"{status.index.embedding_model} ({status.index.embedding_dim} dims)",
123 )
124 vision = status.config.vision_model or "(disabled)"
125 reranker = status.config.reranker_model or "(disabled)"
126 yield _label_line("Vision", vision)
127 yield _label_line("Reranker", reranker)
128 if status.config.enable_ocr is not None:
129 ocr_label = "enabled" if status.config.enable_ocr else "disabled"
130 yield _label_line("OCR", ocr_label)
131 if status.ocr_note is not None:
132 yield _label_line("OCR engine", status.ocr_note)
133 if status.ocr_warning is not None:
134 yield styled((status.ocr_warning, theme.WARNING))
135 if status.entities is not None:
136 names = ", ".join(status.entities.types) or "schema pending (induced on next sync)"
137 yield _label_line("Entities", f"{status.entities.rows} entities extracted ({names})")
138 yield ""
140 if status.skipped:
141 held = Table(title="Held out of the index")
142 held.add_column("File", style=theme.ACCENT)
143 held.add_column("Reason", style=theme.MUTED)
144 for skipped in status.skipped:
145 held.add_row(Text(skipped.filename), Text(skipped.reason))
146 yield held
147 hidden = status.skipped_total - len(status.skipped)
148 more = f" ({hidden} more not shown)" if hidden > 0 else ""
149 yield styled(
150 (str(status.skipped_total), theme.LABEL),
151 f" held out{more}; 'lilbee sync --retry-skipped' retries failed files",
152 )
153 yield ""
155 if not status.sources:
156 yield (
157 "No documents indexed. Drop files into the documents directory and run 'lilbee sync'."
158 )
159 return
161 table = Table(title="Indexed Documents")
162 table.add_column("File", style=theme.ACCENT)
163 table.add_column("Hash", style=theme.MUTED, max_width=12)
164 table.add_column("Chunks", justify="right")
165 table.add_column("Ingested", style=theme.MUTED)
166 for s in status.sources:
167 table.add_row(Text(s.filename), s.file_hash, str(s.chunk_count), s.ingested_at)
168 yield table
169 yield styled(
170 "\n",
171 (str(len(status.sources)), theme.LABEL),
172 " documents, ",
173 (str(status.total_chunks), theme.LABEL),
174 " chunks",
175 )
178def render_status(con: PlainConsole) -> None:
179 """Print status info (documents, paths, chunk counts)."""
180 from lilbee.app.status import gather_status
182 for renderable in render_status_result(gather_status()):
183 con.print(renderable)
186NAME_TAKEN_WARNING = "The name {name} is taken by another source (use --force to overwrite)."
187"""Said when a label belongs to a different source, the one case --force fixes.
189The TUI states the same thing in its own words (``messages.CMD_ADD_NAME_TAKEN``);
190the two surfaces do not share a string because ``cli.tui.messages`` pulls the
191fleet and wiki import chains that a plain CLI command has no reason to pay for.
192"""
193SEARCHING_FOR = "Searching for: {query}"
194"""The stderr line ``ask`` prints when retrieval ran on a rewritten follow-up."""
197def print_prefixed(con: PlainConsole, prefix: str, detail: object, *, style: str) -> None:
198 """Print *prefix* in *style*, then *detail* as literal text (never markup)."""
199 con.print(Text.assemble((prefix, style), str(detail)), soft_wrap=True)
202def register_paths(paths: list[Path], con: PlainConsole, *, force: bool = False) -> RegisterResult:
203 """Register *paths* as source roots, reporting what happened to each."""
204 result = register_sources(paths, force=force)
205 for name in result.name_taken:
206 warning = NAME_TAKEN_WARNING.format(name=name)
207 print_prefixed(con, "Warning: ", warning, style=theme.WARNING)
208 return result
211def describe_registration(result: RegisterResult) -> str:
212 """One line saying what ``add`` did with the paths it was given.
214 A bare count reads as a failure when the answer is "already tracked, and
215 the sync below covers it" -- which is what re-adding a source lilbee
216 already knows about does.
217 """
218 parts = []
219 if result.registered:
220 parts.append(f"Registered {len(result.registered)} source(s)")
221 if result.tracked:
222 parts.append(f"already tracked: {', '.join(result.tracked)}")
223 if result.overlapping:
224 parts.append(f"overlaps a registered source: {', '.join(result.overlapping)}")
225 return ", ".join(parts) if parts else "Registered 0 source(s)"
228def add_paths(
229 paths: list[Path],
230 con: PlainConsole,
231 *,
232 force: bool = False,
233 background: bool = False,
234 chat_mode: bool = False,
235 sync_status: SyncStatus | None = None,
236 run_sync: Callable[[], object] | None = None,
237) -> None:
238 """Register *paths* as source roots and sync (human output).
239 When *background* is True (chat ``/add``), sync runs in a background thread
240 and this function returns immediately after registering. *run_sync*
241 overrides the foreground sync call (the CLI passes a Ctrl+C-cancellable
242 runner); it defaults to a plain ``asyncio.run(sync())``.
243 """
244 registration = register_paths(paths, con, force=force)
245 summary = describe_registration(registration)
246 if chat_mode:
247 print(summary)
248 else:
249 con.print(Text(summary, style=theme.MUTED), soft_wrap=True)
250 if not registration.reached_corpus:
251 return
253 if background:
254 from lilbee.cli.sync import run_sync_background
256 run_sync_background(con, chat_mode=chat_mode, sync_status=sync_status)
257 return
259 result = run_sync() if run_sync is not None else _run_foreground_sync()
260 con.print(result)
263def _run_foreground_sync() -> object:
264 """Run a blocking sync with no cancellation hook (default for non-CLI callers)."""
265 from lilbee.data.ingest import sync
267 return asyncio.run(sync())
270def sync_result_to_json(result: object) -> dict:
271 """Convert a SyncResult to the JSON output envelope."""
272 from lilbee.data.ingest import SyncResult
274 if not isinstance(result, SyncResult):
275 raise TypeError(f"Expected SyncResult, got {type(result).__name__}")
276 return {"command": "sync", **result.model_dump()}
279def auto_sync(con: PlainConsole, *, background: bool = False) -> None:
280 """Run document sync before queries.
281 When *background* is True, sync runs in a background thread and this
282 function returns immediately (for chat/REPL). When False (default),
283 sync blocks until complete (for ``lilbee ask``).
284 """
285 if background:
286 from lilbee.cli.sync import run_sync_background
288 run_sync_background(con)
289 return
291 from lilbee.cli.sync import _format_sync_summary
292 from lilbee.data.ingest import sync
294 try:
295 result = asyncio.run(sync())
296 except RuntimeError as exc:
297 print_prefixed(con, "Error: ", exc, style=theme.ERROR)
298 raise SystemExit(1) from None
299 summary = _format_sync_summary(
300 len(result.added),
301 len(result.updated),
302 len(result.removed),
303 len(result.failed),
304 len(result.skipped),
305 )
306 if summary:
307 con.print(f"Synced: {summary}", style=theme.MUTED)
310@contextmanager
311def sigint_cancel() -> Iterator[threading.Event]:
312 """Turn Ctrl-C into a token the wiki pass polls, not a mid-page abort.
314 A build runs for hours and writes pages as it goes, so the default
315 KeyboardInterrupt drops it wherever the interpreter happened to be. Setting
316 a token instead lets it stop at a source boundary with what it wrote intact.
317 The previous handler is restored as soon as it fires, so a second Ctrl-C
318 still hard-exits a pass that is not checking the token.
320 signal.signal only works on the main thread; off it (pytest-xdist workers)
321 the token is simply never set and Ctrl-C keeps its default behaviour.
322 """
323 token = threading.Event()
324 if threading.current_thread() is not threading.main_thread():
325 yield token
326 return
327 previous = signal.getsignal(signal.SIGINT)
329 def _on_sigint(_signum: int, _frame: object) -> None:
330 signal.signal(signal.SIGINT, previous)
331 token.set()
333 signal.signal(signal.SIGINT, _on_sigint)
334 try:
335 yield token
336 finally:
337 signal.signal(signal.SIGINT, previous)