Coverage for src/lilbee/cli/helpers.py: 100%

1"""CLI-specific helpers: JSON formatter, Rich rendering, and CLI workflows.""" 

2 

3from __future__ import annotations 

4 

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 

13 

14from rich.console import Console, RenderableType 

15from rich.markup import escape 

16from rich.table import Table 

17 

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 

22 

23if TYPE_CHECKING: 

24 from lilbee.cli.sync import SyncStatus 

25 

26 

27def json_output(data: dict) -> None: 

28 """Print a JSON object to stdout.""" 

29 print(json.dumps(data)) 

30 

31 

32def announce_cold_start(role: object, model: str) -> Console | None: 

33 """Print a "Starting <role> engine (loading <model>)..." stderr line if cold. 

34 

35 Returns a stderr console to print the matching "ready" line through when the 

36 blocking call returns, or ``None`` when the role's server is already warm (no 

37 status needed) or output is JSON (machine-readable, no chatter). The role 

38 parameter is a ``WorkerRole``; typed as ``object`` to keep this CLI helper 

39 free of a provider-layer import at module top. 

40 """ 

41 from lilbee.app.services import get_services 

42 from lilbee.providers.roles import WorkerRole 

43 

44 if cfg.json_mode or not isinstance(role, WorkerRole): 

45 return None 

46 if get_services().provider.role_ready(role): 

47 return None 

48 err = Console(stderr=True) 

49 err.print(f"[{theme.MUTED}]Starting {role.value} engine (loading {model})...[/{theme.MUTED}]") 

50 return err 

51 

52 

53def announce_ready(err: Console | None, role: object) -> None: 

54 """Print the matching "<role> engine ready." stderr line, if cold-start announced. 

55 

56 A token arriving is not evidence the chat model came up: in RAG mode a grounded 

57 refusal streams without it. When warm-up recorded a load failure, that reason is 

58 printed instead of a readiness line. 

59 """ 

60 from lilbee.providers.roles import WorkerRole 

61 

62 if err is None or not isinstance(role, WorkerRole): 

63 return 

64 failure = _chat_warm_error(role) 

65 if failure is not None: 

66 err.print(f"[{theme.ERROR}]{failure}[/{theme.ERROR}]") 

67 return 

68 err.print(f"[{theme.MUTED}]{role.value} engine ready.[/{theme.MUTED}]") 

69 

70 

71def announce_retrieval_query(query: str) -> None: 

72 """Print the "Searching for: <query>" stderr line for a rewritten follow-up.""" 

73 line = SEARCHING_FOR.format(query=escape(query)) 

74 Console(stderr=True).print(f"[{theme.MUTED}]{line}[/{theme.MUTED}]") 

75 

76 

77def _chat_warm_error(role: object) -> str | None: 

78 """The chat warm-up's recorded failure, or None when it did not fail. 

79 

80 Read from the warm tracker rather than re-probing readiness: llama-swap can 

81 report a freshly loaded model as not-yet-running, which would turn a healthy 

82 engine into a spurious failure line. 

83 """ 

84 from lilbee.app.services import get_services 

85 from lilbee.providers.roles import WorkerRole 

86 from lilbee.providers.warm_progress import WarmPhase 

87 

88 if role is not WorkerRole.CHAT: 

89 return None 

90 snapshot = get_services().provider.warm_progress() 

91 if snapshot is None or snapshot.phase is not WarmPhase.ERROR: 

92 return None 

93 return snapshot.error or "The chat model did not finish loading." 

94 

95 

96def render_status_result(status: StatusResult) -> Generator[RenderableType, None, None]: 

97 """Yield Rich renderables for a :class:`StatusResult`.""" 

98 yield f"[{theme.LABEL}]Documents:[/{theme.LABEL}] {status.config.documents_dir}" 

99 yield f"[{theme.LABEL}]Database:[/{theme.LABEL}] {status.config.data_dir}" 

100 yield f"[{theme.LABEL}]Chat model:[/{theme.LABEL}] {status.config.chat_model}" 

101 yield f"[{theme.LABEL}]Embeddings:[/{theme.LABEL}] {status.config.embedding_model}" 

102 vision = status.config.vision_model or "(disabled)" 

103 reranker = status.config.reranker_model or "(disabled)" 

104 yield f"[{theme.LABEL}]Vision:[/{theme.LABEL}] {vision}" 

105 yield f"[{theme.LABEL}]Reranker:[/{theme.LABEL}] {reranker}" 

106 if status.config.enable_ocr is not None: 

107 ocr_label = "enabled" if status.config.enable_ocr else "disabled" 

108 yield f"[{theme.LABEL}]Vision OCR:[/{theme.LABEL}] {ocr_label}" 

109 if status.entities is not None: 

110 names = ", ".join(status.entities.types) or "schema pending (induced on next sync)" 

111 yield ( 

112 f"[{theme.LABEL}]Entities:[/{theme.LABEL}] " 

113 f"{status.entities.rows} entities extracted ({names})" 

114 ) 

115 yield "" 

116 

117 if status.skipped: 

118 held = Table(title="Held out of the index") 

119 held.add_column("File", style=theme.ACCENT) 

120 held.add_column("Reason", style=theme.MUTED) 

121 for skipped in status.skipped: 

122 held.add_row(escape(skipped.filename), escape(skipped.reason)) 

123 yield held 

124 b = theme.LABEL 

125 hidden = status.skipped_total - len(status.skipped) 

126 more = f" ({hidden} more not shown)" if hidden > 0 else "" 

127 yield ( 

128 f"[{b}]{status.skipped_total}[/{b}] held out{more}; " 

129 "run 'lilbee sync --retry-skipped' to try them again" 

130 ) 

131 yield "" 

132 

133 if not status.sources: 

134 yield ( 

135 "No documents indexed. Drop files into the documents directory and run 'lilbee sync'." 

136 ) 

137 return 

138 

139 table = Table(title="Indexed Documents") 

140 table.add_column("File", style=theme.ACCENT) 

141 table.add_column("Hash", style=theme.MUTED, max_width=12) 

142 table.add_column("Chunks", justify="right") 

143 table.add_column("Ingested", style=theme.MUTED) 

144 for s in status.sources: 

145 table.add_row(s.filename, s.file_hash, str(s.chunk_count), s.ingested_at) 

146 yield table 

147 b = theme.LABEL 

148 yield f"\n[{b}]{len(status.sources)}[/{b}] documents, [{b}]{status.total_chunks}[/{b}] chunks" 

149 

150 

151def render_status(con: Console) -> None: 

152 """Print status info (documents, paths, chunk counts).""" 

153 from lilbee.app.status import gather_status 

154 

155 for renderable in render_status_result(gather_status()): 

156 con.print(renderable) 

157 

158 

159NAME_TAKEN_WARNING = "The name {name} is taken by another source (use --force to overwrite)." 

160"""Said when a label belongs to a different source, the one case --force fixes. 

161 

162The TUI states the same thing in its own words (``messages.CMD_ADD_NAME_TAKEN``); 

163the two surfaces do not share a string because ``cli.tui.messages`` pulls the 

164fleet and wiki import chains that a plain CLI command has no reason to pay for. 

165""" 

166SEARCHING_FOR = "Searching for: {query}" 

167"""The stderr line ``ask`` prints when retrieval ran on a rewritten follow-up.""" 

168 

169 

170def register_paths(paths: list[Path], con: Console, *, force: bool = False) -> RegisterResult: 

171 """Register *paths* as source roots, reporting what happened to each.""" 

172 result = register_sources(paths, force=force) 

173 for name in result.skipped: 

174 warning = NAME_TAKEN_WARNING.format(name=name) 

175 con.print(f"[{theme.WARNING}]Warning:[/{theme.WARNING}] {warning}") 

176 return result 

177 

178 

179def describe_registration(result: RegisterResult) -> str: 

180 """One line saying what ``add`` did with the paths it was given. 

181 

182 A bare count reads as a failure when the answer is "already tracked, and 

183 the sync below covers it" -- which is what re-adding a source lilbee 

184 already knows about does. 

185 """ 

186 parts = [] 

187 if result.registered: 

188 parts.append(f"Registered {len(result.registered)} source(s)") 

189 if result.tracked: 

190 parts.append(f"already tracked: {', '.join(result.tracked)}") 

191 return ", ".join(parts) if parts else "Registered 0 source(s)" 

192 

193 

194def add_paths( 

195 paths: list[Path], 

196 con: Console, 

197 *, 

198 force: bool = False, 

199 background: bool = False, 

200 chat_mode: bool = False, 

201 sync_status: SyncStatus | None = None, 

202 run_sync: Callable[[], object] | None = None, 

203) -> None: 

204 """Register *paths* as source roots and sync (human output). 

205 When *background* is True (chat ``/add``), sync runs in a background thread 

206 and this function returns immediately after registering. *run_sync* 

207 overrides the foreground sync call (the CLI passes a Ctrl+C-cancellable 

208 runner); it defaults to a plain ``asyncio.run(sync())``. 

209 """ 

210 summary = describe_registration(register_paths(paths, con, force=force)) 

211 if chat_mode: 

212 print(summary) 

213 else: 

214 con.print(f"[{theme.MUTED}]{summary}[/{theme.MUTED}]") 

215 

216 if background: 

217 from lilbee.cli.sync import run_sync_background 

218 

219 run_sync_background(con, chat_mode=chat_mode, sync_status=sync_status) 

220 return 

221 

222 result = run_sync() if run_sync is not None else _run_foreground_sync() 

223 con.print(result) 

224 

225 

226def _run_foreground_sync() -> object: 

227 """Run a blocking sync with no cancellation hook (default for non-CLI callers).""" 

228 from lilbee.data.ingest import sync 

229 

230 return asyncio.run(sync()) 

231 

232 

233def sync_result_to_json(result: object) -> dict: 

234 """Convert a SyncResult to the JSON output envelope.""" 

235 from lilbee.data.ingest import SyncResult 

236 

237 if not isinstance(result, SyncResult): 

238 raise TypeError(f"Expected SyncResult, got {type(result).__name__}") 

239 return {"command": "sync", **result.model_dump()} 

240 

241 

242def auto_sync(con: Console, *, background: bool = False) -> None: 

243 """Run document sync before queries. 

244 When *background* is True, sync runs in a background thread and this 

245 function returns immediately (for chat/REPL). When False (default), 

246 sync blocks until complete (for ``lilbee ask``). 

247 """ 

248 if background: 

249 from lilbee.cli.sync import run_sync_background 

250 

251 run_sync_background(con) 

252 return 

253 

254 from lilbee.cli.sync import _format_sync_summary 

255 from lilbee.data.ingest import sync 

256 

257 try: 

258 result = asyncio.run(sync()) 

259 except RuntimeError as exc: 

260 con.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}") 

261 raise SystemExit(1) from None 

262 summary = _format_sync_summary( 

263 len(result.added), 

264 len(result.updated), 

265 len(result.removed), 

266 len(result.failed), 

267 len(result.skipped), 

268 ) 

269 if summary: 

270 con.print(f"[{theme.MUTED}]Synced: {summary}[/{theme.MUTED}]") 

271 

272 

273@contextmanager 

274def sigint_cancel() -> Iterator[threading.Event]: 

275 """Turn Ctrl-C into a token the wiki pass polls, not a mid-page abort. 

276 

277 A build runs for hours and writes pages as it goes, so the default 

278 KeyboardInterrupt drops it wherever the interpreter happened to be. Setting 

279 a token instead lets it stop at a source boundary with what it wrote intact. 

280 The previous handler is restored as soon as it fires, so a second Ctrl-C 

281 still hard-exits a pass that is not checking the token. 

282 

283 signal.signal only works on the main thread; off it (pytest-xdist workers) 

284 the token is simply never set and Ctrl-C keeps its default behaviour. 

285 """ 

286 token = threading.Event() 

287 if threading.current_thread() is not threading.main_thread(): 

288 yield token 

289 return 

290 previous = signal.getsignal(signal.SIGINT) 

291 

292 def _on_sigint(_signum: int, _frame: object) -> None: 

293 signal.signal(signal.SIGINT, previous) 

294 token.set() 

295 

296 signal.signal(signal.SIGINT, _on_sigint) 

297 try: 

298 yield token 

299 finally: 

300 signal.signal(signal.SIGINT, previous)