Coverage for src/lilbee/data/store/types.py: 100%
172 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Public dataclasses, TypedDicts, enums, and constants for the store package."""
3from __future__ import annotations
5from dataclasses import dataclass
6from datetime import timedelta
7from enum import StrEnum
8from typing import NamedTuple, NotRequired, TypedDict
10from pydantic import BaseModel, ConfigDict, Field, field_validator
12# How often readers re-check the manifest for new versions from other processes.
13# Zero means strong consistency (every read checks); higher values reduce disk I/O
14# on slow media (HDD) at the cost of serving slightly stale data.
15READ_CONSISTENCY_INTERVAL = timedelta(seconds=5)
18@dataclass
19class ConceptRecords:
20 """Rows for the three concept tables, built from one or more files' chunks."""
22 nodes: list[dict]
23 edges: list[dict]
24 chunk_concepts: list[dict]
26 @classmethod
27 def merged(cls, batches: list[ConceptRecords]) -> ConceptRecords:
28 """Concatenate several record sets into one batched write unit."""
29 return cls(
30 nodes=[row for batch in batches for row in batch.nodes],
31 edges=[row for batch in batches for row in batch.edges],
32 chunk_concepts=[row for batch in batches for row in batch.chunk_concepts],
33 )
36class SourceType(StrEnum):
37 """Values for the ``_sources.source_type`` column.
39 ``DOCUMENT`` mirrors a file under ``documents/`` and is managed by the
40 file-driven sync. ``IMPORTED`` is detached: it came from ``lilbee import``
41 and has no backing file, so sync must not treat it as a missing document.
42 """
44 DOCUMENT = "document"
45 IMPORTED = "imported"
48class SourceMeta(NamedTuple):
49 """Document-level metadata captured at extraction time.
51 ``title`` is always derivable (extraction metadata or the cleaned filename
52 stem); ``authors`` and ``created_at`` are only present when the extractor
53 reports them. Empty strings persist as NULL so old and new rows read alike.
54 """
56 title: str = ""
57 authors: str = ""
58 created_at: str = ""
61class ChunkWrite(NamedTuple):
62 """One document's chunks plus its source-table update, for a batched write.
64 ``Store.write_chunks_batch`` folds many of these into a single locked
65 transaction so bulk ingest doesn't pay a write-lock acquisition per document.
66 ``page_texts`` rows land in the same transaction, after the cleanup delete
67 and before the source row. ``source_type`` lets the detached import path
68 reuse the same atomic write while still tagging its rows ``IMPORTED``.
69 """
71 source: str
72 file_hash: str
73 records: list[dict]
74 needs_cleanup: bool
75 stat: SourceStat | None = None
76 page_texts: list[dict] | None = None
77 source_type: SourceType = SourceType.DOCUMENT
78 meta: SourceMeta | None = None
81class ChunkType(StrEnum):
82 """Values for the ``chunk_type`` column.
84 Documents ingest as ``RAW``, extracted tables as ``TABLE`` (when table
85 extraction is on), and wiki pages written by the wiki producer as
86 ``WIKI``. Callers filter with ``Store.search(chunk_type=...)``; a ``RAW``
87 filter also covers table chunks, since both are document content.
88 """
90 RAW = "raw"
91 TABLE = "table"
92 WIKI = "wiki"
95# ``schema_version`` is an integer for forward-compat. Bump only if we ever need to
96# add or rename a meta column without forcing every store to drop_all.
97# 2: nomic-embed document prefixes are stamped at ingest (see
98# lilbee.retrieval.embedding_profiles doc_prefix_since).
99META_SCHEMA_VERSION = 2
101# Always-true predicate used to clear the single-row ``_meta`` table before re-insert.
102# Lance's ``Table.delete`` requires a SQL where clause; this matches every row without
103# coupling the deletion to any specific column's value domain.
104META_DELETE_ALL_PREDICATE = "schema_version IS NOT NULL"
106# Same, for the single-row ``_entity_schema`` table.
107ENTITY_SCHEMA_DELETE_ALL_PREDICATE = "updated_at IS NOT NULL"
110class EntitySchemaState(TypedDict):
111 """Single-row state of the induced entity schema.
113 ``applied`` records whether a full extraction pass completed under this
114 schema; an interrupted pass leaves it False so the next sync redoes the
115 (idempotent) pass. ``source_count`` is how many documents the index held
116 when the schema was induced, which is what the next sync compares against
117 to decide the corpus has drifted far enough to re-induce.
118 """
120 schema_json: str
121 applied: bool
122 source_count: int
123 updated_at: str
126class SearchScope(StrEnum):
127 """What the user wants to search over.
129 Values are used as-is on CLI flags, MCP params, and HTTP query strings.
130 ``BOTH`` resolves to a ``None`` ``chunk_type`` (no filter); the two
131 others map 1:1 to the chunks-table values.
132 """
134 RAW = ChunkType.RAW
135 WIKI = ChunkType.WIKI
136 BOTH = "both"
139def scope_to_chunk_type(scope: SearchScope | str | None) -> ChunkType | None:
140 """Translate a user-facing scope into a ``Store.search`` ``chunk_type`` arg.
142 ``None``/``"both"`` → no filter. ``"raw"`` / ``"wiki"`` → the matching
143 ``ChunkType``. Raises ``ValueError`` on any other string.
144 """
145 if scope is None:
146 return None
147 normalized = SearchScope(scope)
148 if normalized is SearchScope.BOTH:
149 return None
150 return ChunkType(normalized.value)
153class SearchChunk(BaseModel):
154 """A search result from LanceDB.
155 Every store search path sets ``score``: canonical [0, 1] relevance,
156 higher = better. Ranking, filtering, and selection compare only this
157 field; the arm-specific fields below it are provenance.
158 Vector-arm rows carry ``distance``; FTS-arm rows carry ``bm25_score``;
159 reranked rows additionally carry ``rerank_score`` (higher = better).
160 """
162 model_config = ConfigDict(populate_by_name=True)
164 source: str
165 content_type: str
166 chunk_type: ChunkType = ChunkType.RAW
168 @field_validator("chunk_type", mode="before")
169 @classmethod
170 def _coerce_none_chunk_type(cls, v: str | None) -> str:
171 """LanceDB rows from before the chunk_type column was added return None."""
172 return v if v is not None else ChunkType.RAW
174 # Document title at ingest time; None on rows written before the column
175 # existed (or by writers that carry no title, e.g. wiki pages).
176 title: str | None = None
178 page_start: int
179 page_end: int
180 line_start: int
181 line_end: int
182 chunk: str
183 chunk_index: int
184 vector: list[float] = Field(repr=False)
185 distance: float | None = Field(None, alias="_distance")
186 # Legacy ``_relevance_score`` passthrough. No store path populates it and
187 # no ranking code reads it; it survives only as a display-compatible field
188 # for rows produced by external LanceDB rerankers.
189 relevance_score: float | None = Field(None, validation_alias="_relevance_score")
190 # FTS/BM25-only rows carry a raw, unbounded ``_score``. It lives in its own
191 # field so it never contaminates the canonical ``score``; the
192 # confidence-based expansion-skip reads it (squashed to [0, 1]), and the
193 # relevance filter treats its presence as lexical support.
194 bm25_score: float | None = Field(None, validation_alias="_score")
195 rerank_score: float | None = None
196 # Canonical relevance in [0, 1], set by the store on every search path:
197 # normalized reciprocal-rank fusion on the hybrid path, clamped cosine
198 # similarity on vector-only, list-normalized BM25 on FTS-only probes.
199 score: float | None = None
202class SourceRecord(TypedDict):
203 """A tracked source document record.
205 The stat columns are absent on rows read from stores created before they
206 existed; ``source_stat`` is the accessor that folds absence and the
207 ``SOURCE_STAT_UNKNOWN`` sentinel into ``None``.
208 """
210 filename: str
211 file_hash: str
212 ingested_at: str
213 chunk_count: int
214 source_type: str
215 size_bytes: NotRequired[int]
216 mtime_ns: NotRequired[int]
217 stat_captured_ns: NotRequired[int]
218 # Extraction-time document metadata; absent or None on rows written
219 # before the columns existed, and None when the extractor reported none.
220 title: NotRequired[str | None]
221 authors: NotRequired[str | None]
222 created_at: NotRequired[str | None]
225# Sentinel for the stat columns on rows written before they existed (or for
226# detached imports with no backing file). Planning treats it as "unknown: re-hash".
227SOURCE_STAT_UNKNOWN = -1
230class SourceStat(NamedTuple):
231 """File size and mtime captured when a source was hashed, plus the capture time.
233 ``captured_ns`` is the wall-clock time the stat was taken; the sync planner
234 hashes a file whose mtime is not strictly older than it (racily clean).
235 """
237 size_bytes: int
238 mtime_ns: int
239 captured_ns: int = SOURCE_STAT_UNKNOWN
242def source_stat(record: SourceRecord) -> SourceStat | None:
243 """Stored stat for a source row, or None when unknown.
245 The stat columns are nullable ``int64``, so a row can carry an explicit
246 ``None`` (an import, or a write before the columns existed) as well as a
247 missing key or the ``SOURCE_STAT_UNKNOWN`` sentinel. All three mean "no
248 usable stat": return None so the caller re-hashes instead of crashing on
249 ``int(None)``.
250 """
251 size = record.get("size_bytes")
252 mtime = record.get("mtime_ns")
253 captured = record.get("stat_captured_ns")
254 if size is None or mtime is None or size == SOURCE_STAT_UNKNOWN or mtime == SOURCE_STAT_UNKNOWN:
255 return None
256 captured_ns = SOURCE_STAT_UNKNOWN if captured is None else int(captured)
257 return SourceStat(int(size), int(mtime), captured_ns)
260class SourceStatBackfill(NamedTuple):
261 """An already-tracked source row paired with its freshly verified stat."""
263 record: SourceRecord
264 stat: SourceStat
267class PageTextRecord(TypedDict):
268 """One row of the per-page text dataset, matching ``_page_texts``.
270 The export dataset additionally carries the source's extraction metadata
271 (denormalized onto every page row) so an export/import cycle preserves it;
272 these are absent on rows read from the ``_page_texts`` table and on datasets
273 exported before the columns existed.
274 """
276 source: str
277 page: int
278 text: str
279 content_type: str
280 title: NotRequired[str | None]
281 authors: NotRequired[str | None]
282 created_at: NotRequired[str | None]
285class CitationRecord(TypedDict):
286 """A citation linking a wiki chunk to a specific source location."""
288 wiki_source: str
289 wiki_chunk_index: int
290 citation_key: str
291 claim_type: str
292 source_filename: str
293 source_hash: str
294 page_start: int
295 page_end: int
296 line_start: int
297 line_end: int
298 excerpt: str
299 created_at: str
302class MemoryKind(StrEnum):
303 """Whether a memory is an always-injected preference or a similarity-recalled fact."""
305 PREFERENCE = "preference"
306 FACT = "fact"
309class MemorySource(StrEnum):
310 """Provenance of a memory: user-typed, LLM-extracted, or agent-written."""
312 MANUAL = "manual"
313 EXTRACTED = "extracted"
314 AGENT = "agent"
317# Memory owner namespaces. ``"local"`` is the single human (TUI/CLI/REST); agents own
318# ``"agent:<id>"`` namespaces. The prefix lives only here so it is never hand-spliced.
319LOCAL_OWNER = "local"
320AGENT_OWNER_PREFIX = "agent:"
323def agent_owner(agent_id: str) -> str:
324 """Owner string for an agent identity (``"opencode"`` -> ``"agent:opencode"``)."""
325 return f"{AGENT_OWNER_PREFIX}{agent_id}"
328def is_agent_owner(owner: str) -> bool:
329 """True when *owner* is an agent namespace rather than the local human."""
330 return owner.startswith(AGENT_OWNER_PREFIX)
333class MemoryRow(BaseModel):
334 """A long-term memory entry in the per-library ``_memories`` table.
336 Built from a LanceDB row via ``MemoryRow(**row)`` (which coerces the ``kind``
337 and ``source`` strings to enums) and written back via ``model_dump(mode="json")``.
338 Extra keys like a search ``_distance`` are ignored on construction.
339 """
341 model_config = ConfigDict(extra="ignore")
343 id: str
344 owner: str
345 shared: bool
346 kind: MemoryKind
347 source: MemorySource
348 text: str
349 vector: list[float] = Field(repr=False)
350 created_at: str
351 updated_at: str
354class StoreMeta(TypedDict):
355 """Single-row store metadata recording the embedding model used to build the store.
357 Compatibility is checked before every read and write. When ``cfg.embedding_model``
358 or ``cfg.embedding_dim`` drifts from the persisted row, the store refuses to serve
359 until ``lilbee rebuild`` (CLI) or ``POST /api/sync {"force_rebuild": true}`` (HTTP)
360 rewrites the chunks under the new model.
362 ``updated_at`` is an ISO 8601 UTC timestamp produced by ``datetime.isoformat()``;
363 kept as ``str`` to match the LanceDB ``utf8`` schema column.
364 """
366 embedding_model: str
367 embedding_dim: int
368 schema_version: int
369 updated_at: str
372class EmbeddingModelMismatchError(RuntimeError):
373 """Raised when stored vectors were built with a different embedder than ``cfg``.
375 Carries the persisted and configured refs and dims so each surface renders its
376 own recovery affordance (TUI prompt, CLI command, REST body) from the facts.
377 """
379 def __init__(
380 self,
381 *,
382 persisted_model: str,
383 persisted_dim: int,
384 current_model: str,
385 current_dim: int,
386 ) -> None:
387 self.persisted_model = persisted_model
388 self.persisted_dim = persisted_dim
389 self.current_model = current_model
390 self.current_dim = current_dim
391 super().__init__(self._build_message())
393 @property
394 def dims_match(self) -> bool:
395 """True when the index is adoptable by switching embedder alone (same dim)."""
396 return self.persisted_dim == self.current_dim
398 def _build_message(self) -> str:
399 if self.dims_match:
400 return (
401 f"This index was built with embedding model '{self.persisted_model}', "
402 f"but lilbee is configured to use '{self.current_model}'. Configure lilbee "
403 f"to use '{self.persisted_model}' to search this index, or rebuild it under "
404 f"'{self.current_model}'."
405 )
406 return (
407 f"This index was built with embedding model '{self.persisted_model}' "
408 f"(dim {self.persisted_dim}), which differs from the current "
409 f"'{self.current_model}' (dim {self.current_dim}). The dimensions differ, "
410 f"so rebuild the index under '{self.current_model}' to use it."
411 )
414@dataclass
415class RemoveResult:
416 """Result of a remove_documents operation."""
418 removed: list[str]
419 not_found: list[str]