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, for cancel/failure cleanup. 

141 

142 Called on cancel or failure of the add task so a cancelled source is not 

143 re-found on the next sync. Only the registry entries this invocation added are 

144 dropped, with the skip records its sync wrote under them; the source bytes on 

145 disk and files the user owns are never touched. 

146 """ 

147 from lilbee.app.ingest import forget_roots 

148 

149 if labels: 

150 forget_roots(labels) 

151 

152 

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

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

155 

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

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

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

159 """ 

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

161 return any( 

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

163 ) 

164 

165 

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

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

168 last_update = 0.0 

169 

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

171 nonlocal last_update 

172 now = time.monotonic() 

173 if now - last_update < _ADD_EMBED_THROTTLE_SECONDS: 

174 return 

175 last_update = now 

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

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

178 

179 return _tick 

180 

181 

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

183 """Show that Tesseract is OCRing a file; it reports no page count while it runs.""" 

184 reporter.update( 

185 0, 

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

187 indeterminate=True, 

188 ) 

189 

190 

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

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

193 

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

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

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

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

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

199 """ 

200 in_flight: list[str] = [] 

201 embed_tick = _throttled_embed_tick(reporter) 

202 

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

204 reporter.check_cancelled() 

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

206 in_flight.append(data.file) 

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

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

209 with contextlib.suppress(ValueError): 

210 in_flight.remove(data.file) 

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

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

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

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

215 reporter.update( 

216 0, 

217 msg.SYNC_FILE_PROGRESS.format( 

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

219 ), 

220 indeterminate=True, 

221 ) 

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

223 _show_tesseract_ocr(reporter, data) 

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

225 embed_tick(data) 

226 

227 return on_progress 

228 

229 

230def build_sync_progress_callback( 

231 reporter: ProgressReporter, 

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

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

234 

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

236 tick or the row reads as frozen. 

237 """ 

238 embed_tick = _throttled_embed_tick(reporter) 

239 

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

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

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

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

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

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

246 reporter.check_cancelled() 

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

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

249 status = msg.SYNC_FILE_PROGRESS.format( 

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

251 ) 

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

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

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

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

256 reporter.update( 

257 0, 

258 msg.SYNC_FILE_PROGRESS.format( 

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

260 ), 

261 indeterminate=True, 

262 ) 

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

264 _show_tesseract_ocr(reporter, data) 

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

266 embed_tick(data) 

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

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

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

270 

271 return on_progress 

272 

273 

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

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

276 embed_tick = _throttled_embed_tick(reporter) 

277 

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

279 reporter.check_cancelled() 

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

281 embed_tick(data) 

282 

283 return on_progress