Coverage for src/lilbee/cli/tui/screens/chat_helpers.py: 100%
128 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Module-level helpers used by ChatScreen: progress callbacks, file cleanup, stream close."""
3from __future__ import annotations
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
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 ProgressEvent,
30 SyncDoneEvent,
31)
33if TYPE_CHECKING:
34 from lilbee.data.types import SyncResult
36log = logging.getLogger(__name__)
38_ADD_EMBED_THROTTLE_SECONDS = 0.15
39"""Throttle EMBED reporter updates to avoid TaskBar update storms.
41The embed worker fires one EmbedEvent per sub-batch, which on a fast
42laptop can be dozens per second. The Task Center only repaints at 10 Hz
43anyway, so we coalesce here at the same cadence.
44"""
47def close_stream(stream: Any) -> None:
48 """Close a streaming iterator if it satisfies the ClosableIterator protocol."""
49 if isinstance(stream, ClosableIterator):
50 with contextlib.suppress(Exception):
51 stream.close()
54def _opener_argv(platform: str) -> list[str] | None:
55 """The platform's open-with-default-app command, or None to use the browser."""
56 if platform == "darwin":
57 return ["open"]
58 if platform.startswith("linux"):
59 return ["xdg-open"]
60 return None
63def open_local_file(href: str) -> None:
64 """Open a ``file:`` URL with the OS opener so it lands in the default app
65 for its type (an editor for markdown, a viewer for PDF), not a browser
66 rendering raw text. Platforms without a known opener fall back to the
67 webbrowser module."""
68 argv = _opener_argv(sys.platform)
69 if argv is None:
70 webbrowser.open(href)
71 return
72 path = url2pathname(urlparse(href).path)
73 try:
74 subprocess.run([*argv, path], check=False, timeout=10) # noqa: S603 - fixed opener command; path comes from lilbee's own store
75 except (OSError, subprocess.TimeoutExpired):
76 log.warning("Could not open source file: %s", path)
79def detail_for_batch_progress(data: BatchProgressEvent, in_flight: list[str]) -> str:
80 """Pick the user-facing detail label for a BATCH_PROGRESS tick.
82 Per-page rasterization (vision OCR) is the only producer that uses
83 BatchStatus.RASTERIZING; it emits an absolute path in data.file
84 which never matches the relative source name kept in in_flight, so
85 identity-based detection would never fire. Status-based dispatch is
86 the reliable discriminator between per-page and per-file ticks.
87 """
88 if data.status == BatchStatus.RASTERIZING:
89 return msg.ADD_PAGE_PROGRESS.format(
90 status=data.status.capitalize(), current=data.current, total=data.total
91 )
92 if in_flight:
93 return msg.ADD_SYNCING_FILE.format(file=in_flight[0])
94 return msg.ADD_FILE_DONE.format(file=data.file)
97_PREFERENCE_PREFIX = "pref:"
100@dataclass(frozen=True)
101class RememberOutcome:
102 """A /remember result: the toast message plus the notify severity to use."""
104 message: str
105 severity: str = "information"
108def remember_from_input(raw: str) -> RememberOutcome:
109 """Parse, gate, and store a ``/remember`` command; return the toast outcome.
111 Pure orchestration so the ``@work`` worker body stays a single call and the
112 parse/gate/store path is testable without a running TUI. A leading
113 ``pref:`` marks the text as an always-recalled preference; anything else is
114 stored as a fact.
115 """
116 from lilbee.app.memory import MEMORY_DISABLED_HINT, memory_enabled, remember
117 from lilbee.app.services import get_services
118 from lilbee.data.store import MemoryKind
120 if not memory_enabled():
121 return RememberOutcome(MEMORY_DISABLED_HINT, "warning")
123 text = raw.strip()
124 kind = MemoryKind.FACT
125 if text[: len(_PREFERENCE_PREFIX)].lower() == _PREFERENCE_PREFIX:
126 kind = MemoryKind.PREFERENCE
127 text = text[len(_PREFERENCE_PREFIX) :].strip()
128 if not text:
129 return RememberOutcome(msg.CMD_REMEMBER_USAGE, "warning")
131 if not get_services().embedder.embedding_available():
132 return RememberOutcome(msg.CMD_REMEMBER_NO_EMBED, "warning")
134 remember(text, kind=kind)
135 return RememberOutcome(msg.CMD_REMEMBER_SUCCESS.format(kind=kind.value))
138def unregister_added_roots(labels: list[str]) -> None:
139 """Un-register roots a /add invocation created, for cancel/failure cleanup.
141 Called on cancel or failure of the add task so a cancelled source is not
142 re-found on the next sync. Only the registry entries this invocation added are
143 dropped; the source bytes on disk and files the user owns are never touched.
144 """
145 from lilbee.app.ingest import unregister_roots
147 if labels:
148 unregister_roots(labels)
151def add_indexed_anything(registered: list[str], result: SyncResult) -> bool:
152 """Whether any file under this add's registered roots reached the index.
154 Sync is global, so its added/updated/relocated lists can name files from
155 other sources; only names keyed under a registered label count. A directory
156 root keys files as ``label/relpath``; a single-file root keys as ``label``.
157 """
158 indexed = (*result.added, *result.updated, *result.relocated)
159 return any(
160 name == root or name.startswith(f"{root}/") for root in registered for name in indexed
161 )
164def _throttled_embed_tick(reporter: ProgressReporter) -> Callable[[EmbedEvent], None]:
165 """Return the throttled EMBED tick shared by the add/sync/import callbacks."""
166 last_update = 0.0
168 def _tick(data: EmbedEvent) -> None:
169 nonlocal last_update
170 now = time.monotonic()
171 if now - last_update < _ADD_EMBED_THROTTLE_SECONDS:
172 return
173 last_update = now
174 pct = int(data.chunk * 100 / data.total_chunks) if data.total_chunks else 0
175 reporter.update(pct, msg.SYNC_EMBEDDING.format(file=data.file), indeterminate=False)
177 return _tick
180def build_add_progress_callback(reporter: ProgressReporter) -> DetailedProgressCallback:
181 """Build the on_progress callback used by /add.
183 Tracks files in flight in start order so the displayed filename pins
184 to the oldest unfinished file (the pipeline runs files concurrently;
185 without pinning the label flips around the queue). EXTRACT surfaces
186 "extracted N pages" once per file so a 44MB scanned PDF doesn't read
187 as a hang; EMBED ticks per chunk, throttled to a steady cadence.
188 """
189 in_flight: list[str] = []
190 embed_tick = _throttled_embed_tick(reporter)
192 def on_progress(event_type: EventType, data: ProgressEvent) -> None:
193 reporter.check_cancelled()
194 if event_type == EventType.FILE_START and isinstance(data, FileStartEvent):
195 in_flight.append(data.file)
196 reporter.update(0, msg.ADD_SYNCING_FILE.format(file=in_flight[0]), indeterminate=True)
197 elif event_type == EventType.FILE_DONE and isinstance(data, FileDoneEvent):
198 with contextlib.suppress(ValueError):
199 in_flight.remove(data.file)
200 elif event_type == EventType.BATCH_PROGRESS and isinstance(data, BatchProgressEvent):
201 pct = (data.current / data.total * 100.0) if data.total else 0.0
202 reporter.update(pct, detail_for_batch_progress(data, in_flight), indeterminate=False)
203 elif event_type == EventType.EXTRACT and isinstance(data, ExtractEvent):
204 reporter.update(
205 0,
206 msg.SYNC_FILE_PROGRESS.format(
207 current=data.page, total=data.total_pages, file=data.file
208 ),
209 indeterminate=True,
210 )
211 elif event_type == EventType.EMBED and isinstance(data, EmbedEvent):
212 embed_tick(data)
214 return on_progress
217def build_sync_progress_callback(
218 reporter: ProgressReporter,
219) -> Callable[[EventType, ProgressEvent], None]:
220 """Return the on_progress shim used by ``_do_sync``.
222 EXTRACT mirrors the /add path: a 44MB scanned PDF needs a per-page
223 tick or the row reads as frozen.
224 """
225 embed_tick = _throttled_embed_tick(reporter)
227 def on_progress(event_type: EventType, data: ProgressEvent) -> None:
228 # Mirror /add: explicit cancel check on every event so a SYNC task
229 # cancelled mid-batch stops at the next progress tick instead of
230 # finishing the current file. update() also checks, but events
231 # without a reporter.update call (e.g. BATCH_PROGRESS in the
232 # ingest_stream path) would otherwise miss the cooperative checkpoint.
233 reporter.check_cancelled()
234 if event_type == EventType.FILE_START and isinstance(data, FileStartEvent):
235 pct = int((data.current_file - 1) * 100 / data.total_files)
236 status = msg.SYNC_FILE_PROGRESS.format(
237 current=data.current_file, total=data.total_files, file=data.file
238 )
239 reporter.update(pct, status, indeterminate=False)
240 elif event_type == EventType.FILE_DONE and isinstance(data, FileDoneEvent):
241 reporter.update(0, msg.SYNC_FILE_DONE.format(file=data.file), indeterminate=False)
242 elif event_type == EventType.EXTRACT and isinstance(data, ExtractEvent):
243 reporter.update(
244 0,
245 msg.SYNC_FILE_PROGRESS.format(
246 current=data.page, total=data.total_pages, file=data.file
247 ),
248 indeterminate=True,
249 )
250 elif event_type == EventType.EMBED and isinstance(data, EmbedEvent):
251 embed_tick(data)
252 elif event_type == EventType.DONE and isinstance(data, SyncDoneEvent):
253 total = data.added + data.updated + data.removed
254 reporter.update(100, msg.SYNC_STATUS_DONE.format(count=total), indeterminate=False)
256 return on_progress
259def build_import_progress_callback(reporter: ProgressReporter) -> DetailedProgressCallback:
260 """Build the on_progress callback used by /import (EMBED events only)."""
261 embed_tick = _throttled_embed_tick(reporter)
263 def on_progress(event_type: EventType, data: ProgressEvent) -> None:
264 reporter.check_cancelled()
265 if event_type == EventType.EMBED and isinstance(data, EmbedEvent):
266 embed_tick(data)
268 return on_progress