Coverage for src/lilbee/data/types.py: 100%
134 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Shared ingest types and constants."""
3from __future__ import annotations
5import hashlib
6from dataclasses import dataclass
7from enum import StrEnum
8from pathlib import Path
9from typing import NamedTuple, NotRequired, TypedDict
11from pydantic import BaseModel
12from rich.highlighter import ReprHighlighter
13from rich.text import Text
15from lilbee.core.vectors import Vector
16from lilbee.data.store import (
17 ChunkType,
18 ConceptRecords,
19 IndexMismatch,
20 PageTextRecord,
21 SourceMeta,
22 SourceStat,
23 SourceStatBackfill,
24)
25from lilbee.runtime.progress import OcrBackendUsed
27# PDF and image content types route to paginated extraction; every other format
28# routes to markdown extraction. content_type is derived per-file in
29# discovery.classify_file (PDFs and images grouped; others keyed by extension).
30PDF_CONTENT_TYPE = "pdf"
31IMAGE_CONTENT_TYPE = "image"
32CODE_CONTENT_TYPE = "code"
33MARKDOWN_OUTPUT = "markdown"
34MARKDOWN_MIME = "text/markdown"
35# Sync summary note for a skipped document whose extraction ran with OCR off.
36SKIPPED_OCR_OFF_NOTE = ": ocr is off"
39def is_under(key: str, root: str) -> bool:
40 """Whether source *key* is *root* or lies below it, on whole path segments."""
41 return key == root or key.startswith(f"{root}/")
44@dataclass(frozen=True)
45class ShardId:
46 """Which slice of the corpus one ingest worker owns.
48 *corpus_root* is the data root holding the corpus's ``.lilbeeignore``,
49 which every worker uses in place of its own private data root.
50 """
52 index: int
53 count: int
54 corpus_root: Path
56 def owns(self, key: str) -> bool:
57 """Whether source *key* belongs to this slice.
59 Hashed with blake2b, not ``hash()``, which is salted per process: two
60 runs would deal the same corpus differently and every resume would
61 re-embed what a sibling already holds.
62 """
63 digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest()
64 return int.from_bytes(digest, "big") % self.count == self.index
67class FileToProcess(NamedTuple):
68 """A file queued for ingestion with its metadata."""
70 name: str
71 path: Path
72 content_type: str
73 file_hash: str
74 needs_cleanup: bool
75 stat: SourceStat | None = None
78class MemberRecords(NamedTuple):
79 """One archive member's records, written as its own source."""
81 name: str
82 content_type: str
83 records: list[ChunkRecord]
84 page_texts: list[PageTextRecord]
85 meta: SourceMeta
88class DocumentRecords(NamedTuple):
89 """One file's records, its source metadata, and the OCR its extraction ran with."""
91 records: list[ChunkRecord]
92 meta: SourceMeta
93 ocr: OcrReport | None = None
96class SkippedSource(BaseModel):
97 """One file a held-out record holds out of the index, and why."""
99 filename: str
100 reason: str
103class FileChangePlan(NamedTuple):
104 """Outcome of diffing disk files against the tracked sources."""
106 files_to_process: list[FileToProcess]
107 added: dict[str, None]
108 updated: dict[str, None]
109 unchanged: int
110 stat_backfills: list[SourceStatBackfill]
111 # Files a held-out record holds out; not in the index, so never counted as unchanged.
112 held_out: list[str]
115class OcrBackendName(StrEnum):
116 """OCR backends lilbee selects in OcrConfig: xberg's tesseract or lilbee's vision plugin."""
118 TESSERACT = "tesseract"
119 LILBEE_VISION = "lilbee-vision"
122class OcrReport(BaseModel, frozen=True):
123 """Which OCR backend one extraction ran and how many pages it OCR'd."""
125 backend: OcrBackendUsed
126 pages: int = 0
129class EmbeddingBackendName(StrEnum):
130 """Embedding backends registered with xberg. lilbee registers its own embedder
131 as a plugin so the semantic chunker detects boundaries with the same model that
132 vectorizes chunks."""
134 LILBEE = "lilbee"
137class TokenizerBackendName(StrEnum):
138 """Tokenizer backends registered with xberg. lilbee registers its embedder's
139 tokenizer so ChunkSizing counts chunk budgets in the same tokens the embedder
140 consumes, instead of a chars-per-token heuristic. Separate registry from the
141 embedding backend, so sharing the ``lilbee`` name is fine."""
143 LILBEE = "lilbee"
146class ExtractMode(StrEnum):
147 """Extraction topology: paginated (PDFs/images) vs markdown output (text formats)."""
149 MARKDOWN = "markdown"
150 PAGINATED = "paginated"
153class ChunkRecord(TypedDict):
154 """A single store-ready chunk record matching store.CHUNKS_SCHEMA."""
156 source: str
157 content_type: str
158 chunk_type: ChunkType
159 page_start: int
160 page_end: int
161 line_start: int
162 line_end: int
163 chunk: str
164 chunk_index: int
165 vector: Vector
166 # Stamped once per document by the pipeline (see produce_records); None
167 # when the title is empty, so chunk rows persist NULL like the _sources table.
168 title: NotRequired[str | None]
171class SyncResult(BaseModel):
172 """Summary of a sync operation."""
174 added: list[str] = []
175 updated: list[str] = []
176 removed: list[str] = []
177 unchanged: int = 0
178 # Sources recognized as moved (same content hash, new location): re-keyed to
179 # the new name in place, so their chunks and embeddings were reused, not rebuilt.
180 relocated: list[str] = []
181 failed: list[str] = []
182 skipped: list[str] = []
183 # The OCR each skipped document ran with; files that never reach OCR are absent.
184 skipped_ocr: dict[str, OcrReport] = {}
185 # Files an earlier sync held out, so this run did not attempt them.
186 held_out: list[SkippedSource] = []
187 # Chunks whose text exceeded the embedder's char budget and were truncated
188 # before embedding. Non-zero means some tail content did not reach the index.
189 truncated: int = 0
190 # Folders the walk could not open, so the files in them were not looked at.
191 unreadable_folders: list[str] = []
192 # Set when the index was built with another embedder than the one configured:
193 # the sync left it as it is, and search refuses it until a rebuild or a switch back.
194 index_mismatch: IndexMismatch | None = None
196 def _lines(self) -> list[list[tuple[str, str]]]:
197 """The summary as lines of ``(text, style)`` segments; ``""`` means unstyled."""
198 lines: list[list[tuple[str, str]]] = [
199 [(f"Added: {len(self.added)}", "")],
200 [(f"Updated: {len(self.updated)}", "")],
201 [(f"Removed: {len(self.removed)}", "")],
202 [(f"Unchanged: {self.unchanged}", "")],
203 ]
204 if self.index_mismatch is not None:
205 lines.append([("Index mismatch:", "red"), (f" {self.index_mismatch.message}", "")])
206 if self.relocated:
207 lines.append([(f"Relocated: {len(self.relocated)}", "")])
208 lines += [
209 [(f"Held out: {len(self.held_out)}", "")],
210 [(f"Skipped: {len(self.skipped)}", "")],
211 [(f"Failed: {len(self.failed)}", "")],
212 [(f"Truncated: {self.truncated}", "")],
213 ]
214 lines += [
215 [(" ", ""), (h.filename, "yellow"), (f": {h.reason}", "")] for h in self.held_out
216 ]
217 lines += [[(" ", ""), (name, "yellow"), self._skip_note(name)] for name in self.skipped]
218 lines += [[(" ", ""), (name, "red")] for name in self.failed]
219 lines += [
220 [("Could not open folder: ", "yellow"), (folder, "")]
221 for folder in self.unreadable_folders
222 ]
223 return lines
225 def _skip_note(self, name: str) -> tuple[str, str]:
226 """The OCR-off note for a skipped file, or an empty segment."""
227 report = self.skipped_ocr.get(name)
228 if report is not None and report.backend is OcrBackendUsed.NONE:
229 return (SKIPPED_OCR_OFF_NOTE, "")
230 return ("", "")
232 def __str__(self) -> str:
233 return "\n".join(
234 "".join(f"[{style}]{text}[/{style}]" if style else text for text, style in line)
235 for line in self._lines()
236 )
238 def __repr__(self) -> str:
239 return (
240 f"SyncResult(added={len(self.added)}, updated={len(self.updated)}, "
241 f"removed={len(self.removed)}, unchanged={self.unchanged}, "
242 f"held_out={len(self.held_out)}, skipped={len(self.skipped)}, "
243 f"failed={len(self.failed)}, truncated={self.truncated})"
244 )
246 def __rich__(self) -> Text:
247 """Render the summary with every filename and reason as literal text."""
248 rendered = Text("\n").join(Text.assemble(*line) for line in self._lines())
249 return ReprHighlighter()(rendered)
252@dataclass
253class _IngestResult:
254 """Outcome of a single file ingestion attempt.
256 ``records`` carries the produced (extracted + embedded) chunks until the
257 batched flush writes them; ``None`` on a failed file. ``needs_cleanup``
258 travels with the records so the flush can delete the source's old chunks in
259 the same transaction. ``page_texts`` carries the per-page text dataset rows
260 and ``concept_records`` the file's concept-table rows, and ``entity_rows``
261 the file's typed-entity rows, all written by the same flush. ``meta``
262 carries the document's extraction-time metadata for the source row.
263 ``skip_reason`` is set when the file was refused rather than attempted, and
264 it decides the outcome ahead of the chunk count. ``ocr`` is the OCR the
265 document extraction ran with; ``None`` for files that never reach OCR.
266 """
268 name: str
269 path: Path
270 chunk_count: int
271 error: Exception | None
272 file_hash: str = ""
273 skip_reason: str | None = None
274 records: list[ChunkRecord] | None = None
275 needs_cleanup: bool = True
276 page_texts: list[PageTextRecord] | None = None
277 stat: SourceStat | None = None
278 concept_records: ConceptRecords | None = None
279 entity_rows: list[dict] | None = None
280 meta: SourceMeta | None = None
281 members: list[MemberRecords] | None = None
282 ocr: OcrReport | None = None