Coverage for src/lilbee/cli/helpers.py: 100%
181 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""CLI-specific helpers: JSON formatter, Rich rendering, and CLI workflows."""
3from __future__ import annotations
5import asyncio
6import json
7import logging
8from collections.abc import Callable, Generator
9from pathlib import Path
10from typing import TYPE_CHECKING
12from rich.console import RenderableType
13from rich.table import Table
14from rich.text import Text
16from lilbee.app.ingest import RegisterResult, register_sources
17from lilbee.app.settings import SCANNED_PAGES_LABEL
18from lilbee.app.status import StatusResult
19from lilbee.cli import theme
20from lilbee.core.config import cfg
21from lilbee.core.text import strip_control_characters
22from lilbee.runtime.console import PlainConsole, styled
23from lilbee.runtime.lock import LibraryUnlockableError
25if TYPE_CHECKING:
26 from lilbee.cli.sync import SyncStatus
28log = logging.getLogger(__name__)
30_SYNC_SKIPPED = "Skipped the sync before the answer. %s"
33def json_output(data: dict) -> None:
34 """Print a JSON object to stdout."""
35 print(json.dumps(data))
38def announce_cold_start(role: object, model: str) -> PlainConsole | None:
39 """Print a "Starting <role> engine (loading <model>)..." stderr line if cold.
41 Returns a stderr console to print the matching "ready" line through when the
42 blocking call returns, or ``None`` when the role's server is already warm (no
43 status needed) or output is JSON (machine-readable, no chatter). The role
44 parameter is a ``WorkerRole``; typed as ``object`` to keep this CLI helper
45 free of a provider-layer import at module top.
46 """
47 from lilbee.app.services import get_services
48 from lilbee.providers.roles import WorkerRole
50 if cfg.json_mode or not isinstance(role, WorkerRole):
51 return None
52 if get_services().provider.role_ready(role):
53 return None
54 err = PlainConsole(stderr=True)
55 err.print(
56 Text.assemble(
57 (f"Starting {role.value} engine (loading ", theme.MUTED),
58 model,
59 (")...", theme.MUTED),
60 ),
61 soft_wrap=True,
62 )
63 return err
66def announce_ready(err: PlainConsole | None, role: object) -> None:
67 """Print the matching "<role> engine ready." stderr line, if cold-start announced.
69 A token arriving is not evidence the chat model came up: in RAG mode a grounded
70 refusal streams without it. When warm-up recorded a load failure, that reason is
71 printed instead of a readiness line.
72 """
73 from lilbee.providers.roles import WorkerRole
75 if err is None or not isinstance(role, WorkerRole):
76 return
77 failure = _chat_warm_error(role)
78 if failure is not None:
79 err.print(failure, style=theme.ERROR, soft_wrap=True)
80 return
81 err.print(f"{role.value} engine ready.", style=theme.MUTED)
84def announce_retrieval_query(query: str) -> None:
85 """Print the "Searching for: <query>" stderr line for a rewritten follow-up."""
86 line = SEARCHING_FOR.format(query=query)
87 PlainConsole(stderr=True).print(line, style=theme.MUTED, soft_wrap=True)
90def _chat_warm_error(role: object) -> str | None:
91 """The chat warm-up's recorded failure, or None when it did not fail.
93 Read from the warm tracker rather than re-probing readiness: llama-swap can
94 report a freshly loaded model as not-yet-running, which would turn a healthy
95 engine into a spurious failure line.
96 """
97 from lilbee.app.services import get_services
98 from lilbee.providers.roles import WorkerRole
99 from lilbee.providers.warm_progress import WarmPhase
101 if role is not WorkerRole.CHAT:
102 return None
103 snapshot = get_services().provider.warm_progress()
104 if snapshot is None or snapshot.phase is not WarmPhase.ERROR:
105 return None
106 return snapshot.error or "The chat model did not finish loading."
109_LABEL_WIDTH = len("Chat model:")
112def _label_line(label: str, value: object) -> Text:
113 """An aligned ``Label: value`` line with *value* as literal text."""
114 gap = " " * max(1, _LABEL_WIDTH - len(label))
115 return Text.assemble((f"{label}:", theme.LABEL), gap, str(value))
118def render_status_result(status: StatusResult) -> Generator[RenderableType, None, None]:
119 """Yield Rich renderables for a :class:`StatusResult`."""
120 yield _label_line("Documents", status.config.documents_dir)
121 yield _label_line("Database", status.config.data_dir)
122 yield _label_line("Chat model", status.config.chat_model)
123 yield _label_line("Embeddings", status.config.embedding_model)
124 if status.index is not None:
125 yield _label_line(
126 "Index built with",
127 f"{status.index.embedding_model} ({status.index.embedding_dim} dims)",
128 )
129 vision = status.config.vision_model or "(disabled)"
130 reranker = status.config.reranker_model or "(disabled)"
131 yield _label_line("Vision", vision)
132 yield _label_line("Reranker", reranker)
133 yield _label_line(SCANNED_PAGES_LABEL, status.ocr_note)
134 if status.entities is not None:
135 names = ", ".join(status.entities.types) or "schema pending (induced on next sync)"
136 yield _label_line("Entities", f"{status.entities.rows} entities extracted ({names})")
137 yield ""
139 if status.skipped:
140 held = Table(title="Held out of the index")
141 held.add_column("File", style=theme.ACCENT)
142 held.add_column("Reason", style=theme.MUTED)
143 for skipped in status.skipped:
144 held.add_row(Text(strip_control_characters(skipped.filename)), Text(skipped.reason))
145 yield held
146 hidden = status.skipped_total - len(status.skipped)
147 more = f" ({hidden} more not shown)" if hidden > 0 else ""
148 yield styled(
149 (str(status.skipped_total), theme.LABEL),
150 f" held out{more}; 'lilbee sync --retry-skipped' retries failed files",
151 )
152 yield ""
154 if not status.document_count:
155 yield (
156 "No documents indexed. Drop files into the documents directory and run 'lilbee sync'."
157 )
158 return
160 table = Table(title="Newest Documents")
161 table.add_column("File", style=theme.ACCENT)
162 table.add_column("Hash", style=theme.MUTED, max_width=12)
163 table.add_column("Chunks", justify="right")
164 table.add_column("Ingested", style=theme.MUTED)
165 for s in status.sources:
166 table.add_row(
167 Text(strip_control_characters(s.filename)),
168 s.file_hash,
169 str(s.chunk_count),
170 s.ingested_at,
171 )
172 yield table
173 yield styled(
174 "\n",
175 (str(status.document_count), theme.LABEL),
176 " documents, ",
177 (str(status.total_chunks), theme.LABEL),
178 " chunks",
179 )
180 if status.source_kinds:
181 kinds = ", ".join(f"{kind} {count}" for kind, count in sorted(status.source_kinds.items()))
182 yield _label_line("Kinds", kinds)
183 yield styled(
184 f"Showing the newest {len(status.sources)}. ",
185 "Run 'lilbee status --list-sources' to list every document.",
186 )
189def print_source_listing() -> None:
190 """Print every indexed document, one line each, holding one batch in memory at most."""
191 from lilbee.app.services import get_services
193 for head in get_services().store.iter_source_heads():
194 print(json.dumps(head) if cfg.json_mode else strip_control_characters(head["filename"]))
197def render_status(con: PlainConsole) -> None:
198 """Print status info (documents, paths, chunk counts)."""
199 from lilbee.app.status import gather_status
201 for renderable in render_status_result(gather_status()):
202 con.print(renderable)
205NAME_TAKEN_WARNING = "The name {name} is taken by another source (use --force to overwrite)."
206"""Said when a label belongs to a different source, the one case --force fixes.
208The TUI states the same thing in its own words (``messages.CMD_ADD_NAME_TAKEN``);
209the two surfaces do not share a string because ``cli.tui.messages`` pulls the
210fleet and wiki import chains that a plain CLI command has no reason to pay for.
211"""
212CONTAINS_SOURCE = "contains a source lilbee already indexes, not added: {names}"
213"""Said of a directory that is the parent of a registered source; its other files stay out."""
214NOW_INCLUDES = "{parent} now includes {children}"
215"""Said of a new source that took in the registered sources inside it."""
216SEARCHING_FOR = "Searching for: {query}"
217"""The stderr line ``ask`` prints when retrieval ran on a rewritten follow-up."""
220def print_prefixed(con: PlainConsole, prefix: str, detail: object, *, style: str) -> None:
221 """Print *prefix* in *style*, then *detail* as literal text (never markup)."""
222 con.print(Text.assemble((prefix, style), str(detail)), soft_wrap=True)
225def register_paths(
226 paths: list[Path],
227 con: PlainConsole,
228 *,
229 force: bool = False,
230) -> RegisterResult:
231 """Register *paths* as source roots, reporting what happened to each."""
232 result = register_sources(paths, force=force)
233 for name in result.name_taken:
234 warning = NAME_TAKEN_WARNING.format(name=name)
235 print_prefixed(con, "Warning: ", warning, style=theme.WARNING)
236 return result
239def describe_registration(result: RegisterResult) -> str:
240 """One line saying what ``add`` did with the paths it was given.
242 A bare count reads as a failure when the answer is "already tracked, and
243 the sync below covers it" -- which is what re-adding a source lilbee
244 already knows about does.
245 """
246 parts = []
247 if result.registered:
248 parts.append(f"Registered {len(result.registered)} source(s)")
249 if result.tracked:
250 parts.append(f"already tracked: {', '.join(result.tracked)}")
251 if result.overlapping_inside:
252 parts.append(f"overlaps a registered source: {', '.join(result.overlapping_inside)}")
253 if result.containing:
254 parts.append(CONTAINS_SOURCE.format(names=", ".join(result.containing)))
255 summary = ", ".join(parts) if parts else "Registered 0 source(s)"
256 took_in = [
257 NOW_INCLUDES.format(parent=parent, children=", ".join(children))
258 for parent, children in result.absorbed_into.items()
259 ]
260 return "; ".join([summary, *took_in])
263def add_paths(
264 paths: list[Path],
265 con: PlainConsole,
266 *,
267 force: bool = False,
268 background: bool = False,
269 chat_mode: bool = False,
270 sync_status: SyncStatus | None = None,
271 run_sync: Callable[[RegisterResult], object] | None = None,
272 sync_anyway: bool = False,
273) -> None:
274 """Register *paths* as source roots and sync (human output).
275 When *background* is True (chat ``/add``), sync runs in a background thread
276 and this function returns immediately after registering. *run_sync*
277 overrides the foreground sync call and receives the registration result
278 (the CLI passes a Ctrl+C-cancellable runner); it defaults to a plain
279 ``asyncio.run(sync(started_by="lilbee add"))``. The sync runs when a path reached the corpus, or
280 when *sync_anyway* says the caller has other new content to index.
281 """
282 registration = register_paths(paths, con, force=force)
283 summary = describe_registration(registration)
284 if chat_mode:
285 print(summary)
286 else:
287 con.print(Text(summary, style=theme.MUTED), soft_wrap=True)
288 if not (registration.reached_corpus or sync_anyway):
289 return
291 if background:
292 from lilbee.cli.sync import run_sync_background
294 run_sync_background(con, chat_mode=chat_mode, sync_status=sync_status)
295 return
297 result = run_sync(registration) if run_sync is not None else _run_foreground_sync()
298 con.print(result)
301def _run_foreground_sync() -> object:
302 """Run a blocking sync with no cancellation hook (default for non-CLI callers)."""
303 from lilbee.data.ingest import sync
305 return asyncio.run(sync(started_by="lilbee add"))
308def sync_result_to_json(result: object) -> dict:
309 """Convert a SyncResult to the JSON output envelope."""
310 from lilbee.data.ingest import SyncResult
312 if not isinstance(result, SyncResult):
313 raise TypeError(f"Expected SyncResult, got {type(result).__name__}")
314 return {"command": "sync", **result.model_dump()}
317def auto_sync(
318 con: PlainConsole, run_sync: Callable[[], object], *, background: bool = False
319) -> None:
320 """Run document sync before queries.
321 When *background* is True, sync runs in a background thread and this
322 function returns immediately (for chat/REPL). When False (default),
323 *run_sync* (the CLI's Ctrl+C-cancellable runner) blocks until the sync
324 completes (for ``lilbee ask``). A sync that cannot lock the library is
325 skipped with one warning in the log and one notice on stderr, and the
326 caller goes on to its query.
327 """
328 if background:
329 from lilbee.cli.sync import run_sync_background
331 run_sync_background(con)
332 return
334 from lilbee.cli.sync import _format_sync_summary
335 from lilbee.data.ingest import SyncResult
337 try:
338 result = run_sync()
339 except LibraryUnlockableError as exc:
340 # The question only reads, so it is answered from the library as it is.
341 log.warning(_SYNC_SKIPPED, exc)
342 print_prefixed(
343 PlainConsole(stderr=True), "Warning: ", _SYNC_SKIPPED % exc, style=theme.WARNING
344 )
345 return
346 except RuntimeError as exc:
347 print_prefixed(con, "Error: ", exc, style=theme.ERROR)
348 raise SystemExit(1) from None
349 # The sync runner is typed to return object; narrow it before reading counts.
350 if not isinstance(result, SyncResult):
351 raise TypeError(f"Expected SyncResult, got {type(result).__name__}")
352 summary = _format_sync_summary(
353 len(result.added),
354 len(result.updated),
355 len(result.removed),
356 len(result.failed),
357 len(result.skipped),
358 )
359 if summary:
360 con.print(f"Synced: {summary}", style=theme.MUTED)