Coverage for src/lilbee/runtime/progress/types.py: 100%
171 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Event type enums and Pydantic models for the progress protocol."""
3from enum import StrEnum
5from pydantic import BaseModel
7from lilbee.core.config.enums import OcrMode
10class EventType(StrEnum):
11 """Progress event types emitted during sync/ingest.
13 No value may repeat one in :class:`SseEvent`. The HTTP layer writes a
14 progress event's value straight into the SSE ``event:`` field, so a shared
15 value reaches a client as one event name carrying two payload shapes.
16 """
18 FILE_START = "file_start"
19 FILE_DONE = "file_done"
20 BATCH_PROGRESS = "batch_progress"
21 SYNC_DONE = "sync_done"
22 SYNC_WAIT = "sync_wait"
23 EMBED = "embed"
24 EXTRACT = "extract"
25 OCR_START = "ocr_start"
26 CRAWL_START = "crawl_start"
27 CRAWL_PAGE = "crawl_page"
28 CRAWL_PAGE_FAILED = "crawl_page_failed"
29 CRAWL_DONE = "crawl_done"
30 SETUP_START = "setup_start"
31 SETUP_PROGRESS = "setup_progress"
32 SETUP_DONE = "setup_done"
33 WIKI_PHASE = "wiki_phase"
34 WIKI_PAGE = "wiki_page"
35 ANALYZE = "analyze"
36 LIBRARY_UPGRADE = "library_upgrade"
39class SseEvent(StrEnum):
40 """SSE event names used in the HTTP streaming protocol.
42 ``DONE`` is terminal and a stream emits it at most once.
43 """
45 TOKEN = "token" # noqa: S105 -- SSE event name, not a credential
46 REASONING = "reasoning"
47 SOURCES = "sources"
48 RETRIEVAL_QUERY = "retrieval_query"
49 ERROR = "error"
50 DONE = "done"
51 PROGRESS = "progress"
52 HEARTBEAT = "heartbeat"
53 ALREADY_INGESTING = "already_ingesting"
54 WARMING = "warming"
55 WARM = "warm"
56 COMPACTING = "compacting"
57 COMPACTION = "compaction"
58 MEMORY_EXTRACTED = "memory_extracted"
59 GPU_STATS = "gpu_stats"
62class SseErrorCode(StrEnum):
63 """Stable ``code`` values on SSE error events for clients to branch on."""
65 MODEL_TOO_LARGE = "model_too_large"
66 MODEL_NOT_INSTALLED = "model_not_installed"
67 INDEX_EMBEDDER_MISMATCH = "index_embedder_mismatch"
70class FileStartEvent(BaseModel):
71 """Emitted when a file begins ingestion."""
73 file: str
74 total_files: int
75 current_file: int
78class FileDoneEvent(BaseModel):
79 """Emitted when a file finishes ingestion (success or error)."""
81 file: str
82 status: str
83 chunks: int
86class BatchStatus(StrEnum):
87 """Status values for BatchProgressEvent.status."""
89 INGESTED = "ingested"
90 SKIPPED = "skipped"
91 FAILED = "failed"
92 RASTERIZING = "rasterizing"
95class BatchProgressEvent(BaseModel):
96 """Emitted after each file completes during batch ingestion."""
98 file: str
99 status: BatchStatus
100 current: int
101 total: int
104class OcrBackendUsed(StrEnum):
105 """The OCR backend an extraction ran with: none (OCR off), Tesseract, or the vision model."""
107 NONE = "none"
108 TESSERACT = "tesseract"
109 VISION = "vision"
111 @classmethod
112 def chosen(cls, ocr: OcrMode, vision_model: str) -> "OcrBackendUsed":
113 """The backend a configuration picks: none when OCR is off, else vision or Tesseract."""
114 if ocr is OcrMode.OFF:
115 return cls.NONE
116 return cls.VISION if vision_model else cls.TESSERACT
119_EXTRACT_STEP_NAMES: dict[OcrBackendUsed, str] = {
120 OcrBackendUsed.NONE: "Extracted",
121 OcrBackendUsed.TESSERACT: "Tesseract OCR",
122 OcrBackendUsed.VISION: "Vision OCR",
123}
126class ExtractEvent(BaseModel):
127 """Emitted with page-level extraction progress.
129 OCR on a PDF or image fires one event per page as xberg finishes it, with
130 ``page`` the distinct pages xberg has OCR'd so far and ``total_pages`` the
131 total xberg reports for the file. An archive's members fire none. Extraction
132 then fires once per file with ``page == total_pages`` so subscribers see
133 "extracted N pages" before the embed phase ticks. ``ocr_backend`` is the OCR
134 backend that produced the pages, or ``none`` when no page was OCR'd.
135 """
137 file: str
138 page: int
139 total_pages: int
140 ocr_backend: OcrBackendUsed
142 @property
143 def step(self) -> str:
144 """The step that produced these pages, as a progress line names it."""
145 return _EXTRACT_STEP_NAMES[self.ocr_backend]
148class OcrStartEvent(BaseModel):
149 """Emitted once when Tesseract starts OCR on a file, before its first page completes.
151 ``total_pages`` is the file's page count, read before extraction.
152 """
154 file: str
155 total_pages: int
157 @property
158 def status_text(self) -> str:
159 """The progress line for this event, naming the file's page count."""
160 return f"Tesseract OCR on {self.file} ({self.total_pages} pages in the file)"
163class EmbedEvent(BaseModel):
164 """Emitted per batch during embedding."""
166 file: str
167 chunk: int
168 total_chunks: int
171class CrawlStartEvent(BaseModel):
172 """Emitted when a crawl operation begins."""
174 url: str
175 depth: int
178# Sentinel used in CrawlPageEvent.total when the crawl's final page count is
179# not yet known (BFS streaming, page N emitted before N+1 is discovered).
180# Consumers (plugin, TUI, CLI) treat total <= 0 as indeterminate progress.
181CRAWL_TOTAL_UNKNOWN = -1
184class CrawlPageEvent(BaseModel):
185 """Emitted per page during crawling."""
187 url: str
188 current: int
189 total: int
192class CrawlPageFailedEvent(BaseModel):
193 """Emitted when a crawled page yields nothing to save, with the reason why."""
195 url: str
196 reason: str
199class CrawlDoneEvent(BaseModel):
200 """Emitted when a crawl operation completes."""
202 pages_crawled: int
203 files_written: int
206class SyncDoneEvent(BaseModel):
207 """Emitted when the sync operation completes."""
209 added: int
210 updated: int
211 removed: int
212 failed: int
213 skipped: int = 0
214 relocated: int = 0
217class SyncWaitEvent(BaseModel):
218 """Emitted once when a sync starts to wait for another sync on the same library.
220 *holder* and *pid* name the running sync; both are None when it cannot be read.
221 """
223 message: str
224 holder: str | None = None
225 pid: int | None = None
228class SetupStartEvent(BaseModel):
229 """Emitted when a setup/bootstrap operation begins."""
231 component: str
232 size_estimate_bytes: int | None = None
235class SetupProgressEvent(BaseModel):
236 """Emitted periodically during a setup/bootstrap operation."""
238 component: str
239 downloaded_bytes: int
240 total_bytes: int | None = None
241 detail: str = ""
244class SetupDoneEvent(BaseModel):
245 """Emitted when a setup/bootstrap operation completes."""
247 component: str
248 success: bool
249 error: str | None = None
252class WikiPhase(StrEnum):
253 """Stage of a wiki build or synthesis run."""
255 EXTRACT = "extract"
256 GENERATE = "generate"
257 INDEX = "index"
260class WikiPhaseEvent(BaseModel):
261 """Emitted when a wiki run enters a new phase.
263 ``total`` is the number of units the phase will process (sources for a
264 build, clusters for synthesis), 0 where the phase has no unit count.
265 """
267 phase: WikiPhase
268 total: int = 0
271class WikiPageEvent(BaseModel):
272 """Emitted after each source (build) or cluster (synthesis) is written."""
274 label: str
275 pages: int
276 current: int
277 total: int
280class AnalyzeEvent(BaseModel):
281 """Emitted after each batch of files corpus analysis reads."""
283 done: int
284 total: int
285 file: str
288class UpgradePhase(StrEnum):
289 """Stage of a first open that upgrades the library's format."""
291 WAITING = "waiting"
292 TABLE = "table"
293 DONE = "done"
294 FAILED = "failed"
297class TableWork(StrEnum):
298 """What an upgrade does to the table in progress, as the line words it."""
300 UPDATING = "updating"
301 CHECKING = "checking"
302 ADDING_IDS = "adding ids to"
305UPGRADE_HEADLINE = (
306 "Upgrading your library (one time). "
307 "Large libraries take about a minute; nothing is re-embedded."
308)
310_UPGRADE_TEXT: dict[UpgradePhase, str] = {
311 UpgradePhase.WAITING: (
312 "Waiting for another lilbee process to finish with the library · {seconds} s"
313 ),
314 UpgradePhase.TABLE: (
315 "Tables done: {tables_done} of {tables_total} · {work} {table} · {seconds} s"
316 ),
317 UpgradePhase.DONE: "Library upgrade finished in {elapsed_s:.1f} s.",
318 UpgradePhase.FAILED: "Library upgrade stopped after {elapsed_s:.1f} s.",
319}
320_RUNNING_PHASES = frozenset({UpgradePhase.WAITING, UpgradePhase.TABLE})
323class LibraryUpgradeEvent(BaseModel):
324 """Emitted while an open upgrades the library: the wait, each table, then the end.
326 ``tables_done`` of ``tables_total`` counts finished tables, so it only rises;
327 ``work``, ``table`` and ``rows`` say what is done to the one in progress;
328 ``elapsed_s`` is the time since the upgrade began, wait included.
329 """
331 phase: UpgradePhase
332 work: TableWork | None = None
333 table: str = ""
334 rows: int = 0
335 tables_done: int = 0
336 tables_total: int = 0
337 elapsed_s: float = 0.0
339 @property
340 def running(self) -> bool:
341 """Whether the upgrade is still going, so its time keeps rising."""
342 return self.phase in _RUNNING_PHASES
344 def line_at(self, elapsed_s: float) -> str:
345 """The line every surface shows for this event, *elapsed_s* into the upgrade."""
346 fields = self.model_dump() | {"elapsed_s": elapsed_s, "seconds": int(elapsed_s)}
347 return _UPGRADE_TEXT[self.phase].format(**fields)
349 @property
350 def text(self) -> str:
351 """The line every surface shows for this event, at the time it was emitted."""
352 return self.line_at(self.elapsed_s)