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.table import Table 

16 

17from lilbee.app.ingest import RegisterResult, register_sources 

18from lilbee.app.status import StatusResult 

19from lilbee.cli import theme 

20from lilbee.core.config import cfg 

21 

22if TYPE_CHECKING: 

23 from lilbee.cli.sync import SyncStatus 

24 

25 

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

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

28 print(json.dumps(data)) 

29 

30 

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

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

33 

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

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

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

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

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

39 """ 

40 from lilbee.app.services import get_services 

41 from lilbee.providers.roles import WorkerRole 

42 

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

44 return None 

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

46 return None 

47 err = Console(stderr=True) 

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

49 return err 

50 

51 

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

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

54 

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

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

57 printed instead of a readiness line. 

58 """ 

59 from lilbee.providers.roles import WorkerRole 

60 

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

62 return 

63 failure = _chat_warm_error(role) 

64 if failure is not None: 

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

66 return 

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

68 

69 

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

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

72 

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

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

75 engine into a spurious failure line. 

76 """ 

77 from lilbee.app.services import get_services 

78 from lilbee.providers.roles import WorkerRole 

79 from lilbee.providers.warm_progress import WarmPhase 

80 

81 if role is not WorkerRole.CHAT: 

82 return None 

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

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

85 return None 

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

87 

88 

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

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

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

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

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

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

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

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

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

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

99 if status.config.enable_ocr is not None: 

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

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

102 if status.entities is not None: 

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

104 yield ( 

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

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

107 ) 

108 yield "" 

109 

110 if not status.sources: 

111 yield ( 

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

113 ) 

114 return 

115 

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

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

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

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

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

121 for s in status.sources: 

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

123 yield table 

124 b = theme.LABEL 

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

126 

127 

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

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

130 from lilbee.app.status import gather_status 

131 

132 for renderable in render_status_result(gather_status()): 

133 con.print(renderable) 

134 

135 

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

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

138 

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

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

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

142""" 

143 

144 

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

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

147 result = register_sources(paths, force=force) 

148 for name in result.skipped: 

149 warning = NAME_TAKEN_WARNING.format(name=name) 

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

151 return result 

152 

153 

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

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

156 

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

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

159 already knows about does. 

160 """ 

161 parts = [] 

162 if result.registered: 

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

164 if result.tracked: 

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

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

167 

168 

169def add_paths( 

170 paths: list[Path], 

171 con: Console, 

172 *, 

173 force: bool = False, 

174 background: bool = False, 

175 chat_mode: bool = False, 

176 sync_status: SyncStatus | None = None, 

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

178) -> None: 

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

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

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

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

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

184 """ 

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

186 if chat_mode: 

187 print(summary) 

188 else: 

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

190 

191 if background: 

192 from lilbee.cli.sync import run_sync_background 

193 

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

195 return 

196 

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

198 con.print(result) 

199 

200 

201def _run_foreground_sync() -> object: 

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

203 from lilbee.data.ingest import sync 

204 

205 return asyncio.run(sync()) 

206 

207 

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

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

210 from lilbee.data.ingest import SyncResult 

211 

212 if not isinstance(result, SyncResult): 

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

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

215 

216 

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

218 """Run document sync before queries. 

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

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

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

222 """ 

223 if background: 

224 from lilbee.cli.sync import run_sync_background 

225 

226 run_sync_background(con) 

227 return 

228 

229 from lilbee.cli.sync import _format_sync_summary 

230 from lilbee.data.ingest import sync 

231 

232 try: 

233 result = asyncio.run(sync()) 

234 except RuntimeError as exc: 

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

236 raise SystemExit(1) from None 

237 summary = _format_sync_summary( 

238 len(result.added), 

239 len(result.updated), 

240 len(result.removed), 

241 len(result.failed), 

242 len(result.skipped), 

243 ) 

244 if summary: 

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

246 

247 

248@contextmanager 

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

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

251 

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

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

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

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

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

257 

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

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

260 """ 

261 token = threading.Event() 

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

263 yield token 

264 return 

265 previous = signal.getsignal(signal.SIGINT) 

266 

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

268 signal.signal(signal.SIGINT, previous) 

269 token.set() 

270 

271 signal.signal(signal.SIGINT, _on_sigint) 

272 try: 

273 yield token 

274 finally: 

275 signal.signal(signal.SIGINT, previous)