Coverage for src/lilbee/cli/tui/screens/chat_helpers.py: 100%

1"""Module-level helpers used by ChatScreen: progress callbacks, file cleanup, stream close.""" 

2 

3from __future__ import annotations 

4 

5import contextlib 

6import logging 

7import subprocess 

8import sys 

9import time 

10import webbrowser 

11from collections.abc import Callable 

12from dataclasses import dataclass 

13from typing import TYPE_CHECKING, Any 

14from urllib.parse import urlparse 

15from urllib.request import url2pathname 

16 

17from lilbee.cli.tui import messages as msg 

18from lilbee.cli.tui.widgets.task_bar_controller import ProgressReporter 

19from lilbee.providers.base import ClosableIterator 

20from lilbee.runtime.progress import ( 

21 BatchProgressEvent, 

22 BatchStatus, 

23 DetailedProgressCallback, 

24 EmbedEvent, 

25 EventType, 

26 ExtractEvent, 

27 FileDoneEvent, 

28 FileStartEvent, 

29 OcrStartEvent, 

30 ProgressEvent, 

31 SyncDoneEvent, 

32) 

33 

34if TYPE_CHECKING: 

35 from lilbee.data.types import SyncResult 

36 

37log = logging.getLogger(__name__) 

38 

39_ADD_EMBED_THROTTLE_SECONDS = 0.15 

40"""Throttle EMBED reporter updates to avoid TaskBar update storms. 

41 

42The embed worker fires one EmbedEvent per sub-batch, which on a fast 

43laptop can be dozens per second. The Task Center only repaints at 10 Hz 

44anyway, so we coalesce here at the same cadence. 

45""" 

46 

47 

48def close_stream(stream: Any) -> None: 

49 """Close a streaming iterator if it satisfies the ClosableIterator protocol.""" 

50 if isinstance(stream, ClosableIterator): 

51 with contextlib.suppress(Exception): 

52 stream.close() 

53 

54 

55def _opener_argv(platform: str) -> list[str] | None: 

56 """The platform's open-with-default-app command, or None to use the browser.""" 

57 if platform == "darwin": 

58 return ["open"] 

59 if platform.startswith("linux"): 

60 return ["xdg-open"] 

61 return None 

62 

63 

64def open_local_file(href: str) -> None: 

65 """Open a ``file:`` URL with the OS opener so it lands in the default app 

66 for its type (an editor for markdown, a viewer for PDF), not a browser 

67 rendering raw text. Platforms without a known opener fall back to the 

68 webbrowser module.""" 

69 argv = _opener_argv(sys.platform) 

70 if argv is None: 

71 webbrowser.open(href) 

72 return 

73 path = url2pathname(urlparse(href).path) 

74 try: 

75 subprocess.run([*argv, path], check=False, timeout=10) # noqa: S603 - fixed opener command; path comes from lilbee's own store 

76 except (OSError, subprocess.TimeoutExpired): 

77 log.warning("Could not open source file: %s", path) 

78 

79 

80def detail_for_batch_progress(data: BatchProgressEvent, in_flight: list[str]) -> str: 

81 """Pick the user-facing detail label for a BATCH_PROGRESS tick. 

82 

83 Per-page rasterization (vision OCR) is the only producer that uses 

84 BatchStatus.RASTERIZING; it emits an absolute path in data.file 

85 which never matches the relative source name kept in in_flight, so 

86 identity-based detection would never fire. Status-based dispatch is 

87 the reliable discriminator between per-page and per-file ticks. 

88 """ 

89 if data.status == BatchStatus.RASTERIZING: 

90 return msg.ADD_PAGE_PROGRESS.format( 

91 status=data.status.capitalize(), current=data.current, total=data.total 

92 ) 

93 if in_flight: 

94 return msg.ADD_SYNCING_FILE.format(file=in_flight[0]) 

95 return msg.ADD_FILE_DONE.format(file=data.file) 

96 

97 

98_PREFERENCE_PREFIX = "pref:" 

99 

100 

101@dataclass(frozen=True) 

102class RememberOutcome: 

103 """A /remember result: the toast message plus the notify severity to use.""" 

104 

105 message: str 

106 severity: str = "information" 

107 

108 

109def remember_from_input(raw: str) -> RememberOutcome: 

110 """Parse, gate, and store a ``/remember`` command; return the toast outcome. 

111 

112 Pure orchestration so the ``@work`` worker body stays a single call and the 

113 parse/gate/store path is testable without a running TUI. A leading 

114 ``pref:`` marks the text as an always-recalled preference; anything else is 

115 stored as a fact. 

116 """ 

117 from lilbee.app.memory import MEMORY_DISABLED_HINT, memory_enabled, remember 

118 from lilbee.app.services import get_services 

119 from lilbee.data.store import MemoryKind 

120 

121 if not memory_enabled(): 

122 return RememberOutcome(MEMORY_DISABLED_HINT, "warning") 

123 

124 text = raw.strip() 

125 kind = MemoryKind.FACT 

126 if text[: len(_PREFERENCE_PREFIX)].lower() == _PREFERENCE_PREFIX: 

127 kind = MemoryKind.PREFERENCE 

128 text = text[len(_PREFERENCE_PREFIX) :].strip() 

129 if not text: 

130 return RememberOutcome(msg.CMD_REMEMBER_USAGE, "warning") 

131 

132 if not get_services().embedder.embedding_available(): 

133 return RememberOutcome(msg.CMD_REMEMBER_NO_EMBED, "warning") 

134 

135 remember(text, kind=kind) 

136 return RememberOutcome(msg.CMD_REMEMBER_SUCCESS.format(kind=kind.value)) 

137 

138 

139def unregister_added_roots(labels: list[str]) -> None: 

140 """Un-register roots a /add invocation created, when its finished sync failed them. 

141 

142 Only the registry entries this invocation added are dropped, with the skip 

143 records its sync wrote under them; the source bytes on disk and files the 

144 user owns are never touched. 

145 """ 

146 from lilbee.app.ingest import forget_roots 

147 

148 if labels: 

149 forget_roots(labels) 

150 

151 

152def add_indexed_anything(registered: list[str], result: SyncResult) -> bool: 

153 """Whether any file under this add's registered roots reached the index. 

154 

155 Sync is global, so its added/updated/relocated lists can name files from 

156 other sources; only names keyed under a registered label count. A directory 

157 root keys files as ``label/relpath``; a single-file root keys as ``label``. 

158 """ 

159 indexed = (*result.added, *result.updated, *result.relocated) 

160 return any( 

161 name == root or name.startswith(f"{root}/") for root in registered for name in indexed 

162 ) 

163 

164 

165def _throttled_embed_tick(reporter: ProgressReporter) -> Callable[[EmbedEvent], None]: 

166 """Return the throttled EMBED tick shared by the add/sync/import callbacks.""" 

167 last_update = 0.0 

168 

169 def _tick(data: EmbedEvent) -> None: 

170 nonlocal last_update 

171 now = time.monotonic() 

172 if now - last_update < _ADD_EMBED_THROTTLE_SECONDS: 

173 return 

174 last_update = now 

175 pct = int(data.chunk * 100 / data.total_chunks) if data.total_chunks else 0 

176 reporter.update(pct, msg.SYNC_EMBEDDING.format(file=data.file), indeterminate=False) 

177 

178 return _tick 

179 

180 

181def _show_tesseract_ocr(reporter: ProgressReporter, data: OcrStartEvent) -> None: 

182 """Show that Tesseract is OCRing a file, until its first page completes.""" 

183 reporter.update( 

184 0, 

185 msg.SYNC_TESSERACT_OCR.format(total=data.total_pages, file=data.file), 

186 indeterminate=True, 

187 ) 

188 

189 

190def build_add_progress_callback(reporter: ProgressReporter) -> DetailedProgressCallback: 

191 """Build the on_progress callback used by /add. 

192 

193 Tracks files in flight in start order so the displayed filename pins 

194 to the oldest unfinished file (the pipeline runs files concurrently; 

195 without pinning the label flips around the queue). EXTRACT surfaces 

196 "extracted N pages" once per file so a 44MB scanned PDF doesn't read 

197 as a hang; EMBED ticks per chunk, throttled to a steady cadence. 

198 """ 

199 in_flight: list[str] = [] 

200 embed_tick = _throttled_embed_tick(reporter) 

201 

202 def on_progress(event_type: EventType, data: ProgressEvent) -> None: 

203 reporter.check_cancelled() 

204 if event_type == EventType.FILE_START and isinstance(data, FileStartEvent): 

205 in_flight.append(data.file) 

206 reporter.update(0, msg.ADD_SYNCING_FILE.format(file=in_flight[0]), indeterminate=True) 

207 elif event_type == EventType.FILE_DONE and isinstance(data, FileDoneEvent): 

208 with contextlib.suppress(ValueError): 

209 in_flight.remove(data.file) 

210 elif event_type == EventType.BATCH_PROGRESS and isinstance(data, BatchProgressEvent): 

211 pct = (data.current / data.total * 100.0) if data.total else 0.0 

212 reporter.update(pct, detail_for_batch_progress(data, in_flight), indeterminate=False) 

213 elif event_type == EventType.EXTRACT and isinstance(data, ExtractEvent): 

214 reporter.update( 

215 0, 

216 msg.SYNC_FILE_PROGRESS.format( 

217 current=data.page, total=data.total_pages, file=data.file 

218 ), 

219 indeterminate=True, 

220 ) 

221 elif event_type == EventType.OCR_START and isinstance(data, OcrStartEvent): 

222 _show_tesseract_ocr(reporter, data) 

223 elif event_type == EventType.EMBED and isinstance(data, EmbedEvent): 

224 embed_tick(data) 

225 elif event_type == EventType.SYNC_WAIT: 

226 _show_sync_wait(reporter) 

227 

228 return on_progress 

229 

230 

231def _show_sync_wait(reporter: ProgressReporter) -> None: 

232 """Show on the task bar that this sync waits for another; the sync logs the full notice.""" 

233 reporter.update(0, msg.SYNC_WAITING, indeterminate=True) 

234 

235 

236def build_sync_progress_callback( 

237 reporter: ProgressReporter, 

238) -> Callable[[EventType, ProgressEvent], None]: 

239 """Return the on_progress shim used by ``_do_sync``. 

240 

241 EXTRACT mirrors the /add path: a 44MB scanned PDF needs a per-page 

242 tick or the row reads as frozen. 

243 """ 

244 embed_tick = _throttled_embed_tick(reporter) 

245 

246 def on_progress(event_type: EventType, data: ProgressEvent) -> None: 

247 # Mirror /add: explicit cancel check on every event so a SYNC task 

248 # cancelled mid-batch stops at the next progress tick instead of 

249 # finishing the current file. update() also checks, but events 

250 # without a reporter.update call (e.g. BATCH_PROGRESS in the 

251 # ingest_stream path) would otherwise miss the cooperative checkpoint. 

252 reporter.check_cancelled() 

253 if event_type == EventType.FILE_START and isinstance(data, FileStartEvent): 

254 pct = int((data.current_file - 1) * 100 / data.total_files) 

255 status = msg.SYNC_FILE_PROGRESS.format( 

256 current=data.current_file, total=data.total_files, file=data.file 

257 ) 

258 reporter.update(pct, status, indeterminate=False) 

259 elif event_type == EventType.FILE_DONE and isinstance(data, FileDoneEvent): 

260 reporter.update(0, msg.SYNC_FILE_DONE.format(file=data.file), indeterminate=False) 

261 elif event_type == EventType.EXTRACT and isinstance(data, ExtractEvent): 

262 reporter.update( 

263 0, 

264 msg.SYNC_FILE_PROGRESS.format( 

265 current=data.page, total=data.total_pages, file=data.file 

266 ), 

267 indeterminate=True, 

268 ) 

269 elif event_type == EventType.OCR_START and isinstance(data, OcrStartEvent): 

270 _show_tesseract_ocr(reporter, data) 

271 elif event_type == EventType.EMBED and isinstance(data, EmbedEvent): 

272 embed_tick(data) 

273 elif event_type == EventType.SYNC_WAIT: 

274 _show_sync_wait(reporter) 

275 elif event_type == EventType.SYNC_DONE and isinstance(data, SyncDoneEvent): 

276 total = data.added + data.updated + data.removed 

277 reporter.update(100, msg.SYNC_STATUS_DONE.format(count=total), indeterminate=False) 

278 

279 return on_progress 

280 

281 

282def build_import_progress_callback(reporter: ProgressReporter) -> DetailedProgressCallback: 

283 """Build the on_progress callback used by /import (EMBED events only).""" 

284 embed_tick = _throttled_embed_tick(reporter) 

285 

286 def on_progress(event_type: EventType, data: ProgressEvent) -> None: 

287 reporter.check_cancelled() 

288 if event_type == EventType.EMBED and isinstance(data, EmbedEvent): 

289 embed_tick(data) 

290 

291 return on_progress