Coverage for src/lilbee/data/types.py: 100%
107 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +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
13from lilbee.core.vectors import Vector
14from lilbee.data.store import (
15 ChunkType,
16 ConceptRecords,
17 PageTextRecord,
18 SourceMeta,
19 SourceStat,
20 SourceStatBackfill,
21)
23# PDF and image content types route to paginated extraction; every other format
24# routes to markdown extraction. content_type is derived per-file in
25# discovery.classify_file (PDFs and images grouped; others keyed by extension).
26PDF_CONTENT_TYPE = "pdf"
27IMAGE_CONTENT_TYPE = "image"
28MARKDOWN_OUTPUT = "markdown"
29MARKDOWN_MIME = "text/markdown"
32@dataclass(frozen=True)
33class ShardId:
34 """Which slice of the corpus one ingest worker owns."""
36 index: int
37 count: int
39 def owns(self, key: str) -> bool:
40 """Whether source *key* belongs to this slice.
42 Hashed with blake2b, not ``hash()``, which is salted per process: two
43 runs would deal the same corpus differently and every resume would
44 re-embed what a sibling already holds.
45 """
46 digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest()
47 return int.from_bytes(digest, "big") % self.count == self.index
50class FileToProcess(NamedTuple):
51 """A file queued for ingestion with its metadata."""
53 name: str
54 path: Path
55 content_type: str
56 file_hash: str
57 needs_cleanup: bool
58 stat: SourceStat | None = None
61class MemberRecords(NamedTuple):
62 """One archive member's records, written as its own source."""
64 name: str
65 content_type: str
66 records: list[ChunkRecord]
67 page_texts: list[PageTextRecord]
68 meta: SourceMeta
71class SkippedSource(BaseModel):
72 """One file a skip marker holds out of the index, and why."""
74 filename: str
75 reason: str
78class FileChangePlan(NamedTuple):
79 """Outcome of diffing disk files against the tracked sources."""
81 files_to_process: list[FileToProcess]
82 added: dict[str, None]
83 updated: dict[str, None]
84 unchanged: int
85 stat_backfills: list[SourceStatBackfill]
86 # Files a skip marker holds out; not in the index, so never counted as unchanged.
87 held_out: list[str]
90class OcrBackendName(StrEnum):
91 """OCR backends lilbee selects in OcrConfig: xberg's tesseract or lilbee's vision plugin."""
93 TESSERACT = "tesseract"
94 LILBEE_VISION = "lilbee-vision"
97class EmbeddingBackendName(StrEnum):
98 """Embedding backends registered with xberg. lilbee registers its own embedder
99 as a plugin so the semantic chunker detects boundaries with the same model that
100 vectorizes chunks."""
102 LILBEE = "lilbee"
105class TokenizerBackendName(StrEnum):
106 """Tokenizer backends registered with xberg. lilbee registers its embedder's
107 tokenizer so ChunkSizing counts chunk budgets in the same tokens the embedder
108 consumes, instead of a chars-per-token heuristic. Separate registry from the
109 embedding backend, so sharing the ``lilbee`` name is fine."""
111 LILBEE = "lilbee"
114class ExtractMode(StrEnum):
115 """Extraction topology: paginated (PDFs/images) vs markdown output (text formats)."""
117 MARKDOWN = "markdown"
118 PAGINATED = "paginated"
121class ChunkRecord(TypedDict):
122 """A single store-ready chunk record matching store.CHUNKS_SCHEMA."""
124 source: str
125 content_type: str
126 chunk_type: ChunkType
127 page_start: int
128 page_end: int
129 line_start: int
130 line_end: int
131 chunk: str
132 chunk_index: int
133 vector: Vector
134 # Stamped once per document by the pipeline (see produce_records); None
135 # when the title is empty, so chunk rows persist NULL like the _sources table.
136 title: NotRequired[str | None]
139class SyncResult(BaseModel):
140 """Summary of a sync operation."""
142 added: list[str] = []
143 updated: list[str] = []
144 removed: list[str] = []
145 unchanged: int = 0
146 # Sources recognized as moved (same content hash, new location): re-keyed to
147 # the new name in place, so their chunks and embeddings were reused, not rebuilt.
148 relocated: list[str] = []
149 failed: list[str] = []
150 skipped: list[str] = []
151 # Files an earlier sync skip-marked, so this run did not attempt them.
152 held_out: list[SkippedSource] = []
153 # Chunks whose text exceeded the embedder's char budget and were truncated
154 # before embedding. Non-zero means some tail content did not reach the index.
155 truncated: int = 0
157 def __str__(self) -> str:
158 lines = [
159 f"Added: {len(self.added)}",
160 f"Updated: {len(self.updated)}",
161 f"Removed: {len(self.removed)}",
162 f"Unchanged: {self.unchanged}",
163 ]
164 if self.relocated:
165 lines.append(f"Relocated: {len(self.relocated)}")
166 lines += [
167 f"Held out: {len(self.held_out)}",
168 f"Skipped: {len(self.skipped)}",
169 f"Failed: {len(self.failed)}",
170 f"Truncated: {self.truncated}",
171 ]
172 for held in self.held_out:
173 lines.append(f" [yellow]{held.filename}[/yellow]: {held.reason}")
174 for f in self.skipped:
175 lines.append(f" [yellow]{f}[/yellow]")
176 for f in self.failed:
177 lines.append(f" [red]{f}[/red]")
178 return "\n".join(lines)
180 def __repr__(self) -> str:
181 return (
182 f"SyncResult(added={len(self.added)}, updated={len(self.updated)}, "
183 f"removed={len(self.removed)}, unchanged={self.unchanged}, "
184 f"held_out={len(self.held_out)}, skipped={len(self.skipped)}, "
185 f"failed={len(self.failed)}, truncated={self.truncated})"
186 )
188 def __rich__(self) -> str:
189 return self.__str__()
192@dataclass
193class _IngestResult:
194 """Outcome of a single file ingestion attempt.
196 ``records`` carries the produced (extracted + embedded) chunks until the
197 batched flush writes them; ``None`` on a failed file. ``needs_cleanup``
198 travels with the records so the flush can delete the source's old chunks in
199 the same transaction. ``page_texts`` carries the per-page text dataset rows
200 and ``concept_records`` the file's concept-table rows, and ``entity_rows``
201 the file's typed-entity rows, all written by the same flush. ``meta``
202 carries the document's extraction-time metadata for the source row.
203 ``skip_reason`` is set when the file was refused rather than attempted, and
204 it decides the outcome ahead of the chunk count.
205 """
207 name: str
208 path: Path
209 chunk_count: int
210 error: Exception | None
211 file_hash: str = ""
212 skip_reason: str | None = None
213 records: list[ChunkRecord] | None = None
214 needs_cleanup: bool = True
215 page_texts: list[PageTextRecord] | None = None
216 stat: SourceStat | None = None
217 concept_records: ConceptRecords | None = None
218 entity_rows: list[dict] | None = None
219 meta: SourceMeta | None = None
220 members: list[MemberRecords] | None = None