Coverage for src/lilbee/data/types.py: 100%
92 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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 FileChangePlan(NamedTuple):
62 """Outcome of diffing disk files against the tracked sources."""
64 files_to_process: list[FileToProcess]
65 added: dict[str, None]
66 updated: dict[str, None]
67 unchanged: int
68 stat_backfills: list[SourceStatBackfill]
71class OcrBackendName(StrEnum):
72 """OCR backends lilbee selects in OcrConfig: xberg's tesseract or lilbee's vision plugin."""
74 TESSERACT = "tesseract"
75 LILBEE_VISION = "lilbee-vision"
78class EmbeddingBackendName(StrEnum):
79 """Embedding backends registered with xberg. lilbee registers its own embedder
80 as a plugin so the semantic chunker detects boundaries with the same model that
81 vectorizes chunks."""
83 LILBEE = "lilbee"
86class TokenizerBackendName(StrEnum):
87 """Tokenizer backends registered with xberg. lilbee registers its embedder's
88 tokenizer so ChunkSizing counts chunk budgets in the same tokens the embedder
89 consumes, instead of a chars-per-token heuristic. Separate registry from the
90 embedding backend, so sharing the ``lilbee`` name is fine."""
92 LILBEE = "lilbee"
95class ExtractMode(StrEnum):
96 """Extraction topology: paginated (PDFs/images) vs markdown output (text formats)."""
98 MARKDOWN = "markdown"
99 PAGINATED = "paginated"
102class ChunkRecord(TypedDict):
103 """A single store-ready chunk record matching store.CHUNKS_SCHEMA."""
105 source: str
106 content_type: str
107 chunk_type: ChunkType
108 page_start: int
109 page_end: int
110 line_start: int
111 line_end: int
112 chunk: str
113 chunk_index: int
114 vector: Vector
115 # Stamped once per document by the pipeline (see produce_records); None
116 # when the title is empty, so chunk rows persist NULL like the _sources table.
117 title: NotRequired[str | None]
120class SyncResult(BaseModel):
121 """Summary of a sync operation."""
123 added: list[str] = []
124 updated: list[str] = []
125 removed: list[str] = []
126 unchanged: int = 0
127 # Sources recognized as moved (same content hash, new location): re-keyed to
128 # the new name in place, so their chunks and embeddings were reused, not rebuilt.
129 relocated: list[str] = []
130 failed: list[str] = []
131 skipped: list[str] = []
132 # Chunks whose text exceeded the embedder's char budget and were truncated
133 # before embedding. Non-zero means some tail content did not reach the index.
134 truncated: int = 0
136 def __str__(self) -> str:
137 lines = [
138 f"Added: {len(self.added)}",
139 f"Updated: {len(self.updated)}",
140 f"Removed: {len(self.removed)}",
141 f"Unchanged: {self.unchanged}",
142 ]
143 if self.relocated:
144 lines.append(f"Relocated: {len(self.relocated)}")
145 lines += [
146 f"Skipped: {len(self.skipped)}",
147 f"Failed: {len(self.failed)}",
148 f"Truncated: {self.truncated}",
149 ]
150 for f in self.skipped:
151 lines.append(f" [yellow]{f}[/yellow]")
152 for f in self.failed:
153 lines.append(f" [red]{f}[/red]")
154 return "\n".join(lines)
156 def __repr__(self) -> str:
157 return (
158 f"SyncResult(added={len(self.added)}, updated={len(self.updated)}, "
159 f"removed={len(self.removed)}, unchanged={self.unchanged}, "
160 f"skipped={len(self.skipped)}, failed={len(self.failed)}, "
161 f"truncated={self.truncated})"
162 )
164 def __rich__(self) -> str:
165 return self.__str__()
168@dataclass
169class _IngestResult:
170 """Outcome of a single file ingestion attempt.
172 ``records`` carries the produced (extracted + embedded) chunks until the
173 batched flush writes them; ``None`` on a failed file. ``needs_cleanup``
174 travels with the records so the flush can delete the source's old chunks in
175 the same transaction. ``page_texts`` carries the per-page text dataset rows
176 and ``concept_records`` the file's concept-table rows, and ``entity_rows``
177 the file's typed-entity rows, all written by the same flush. ``meta``
178 carries the document's extraction-time metadata for the source row.
179 """
181 name: str
182 path: Path
183 chunk_count: int
184 error: Exception | None
185 file_hash: str = ""
186 records: list[ChunkRecord] | None = None
187 needs_cleanup: bool = True
188 page_texts: list[PageTextRecord] | None = None
189 stat: SourceStat | None = None
190 concept_records: ConceptRecords | None = None
191 entity_rows: list[dict] | None = None
192 meta: SourceMeta | None = None