Coverage for src/lilbee/data/store/types.py: 100%
222 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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 ``WIKI`` is an indexed wiki page; no listing shows it.
43 """
45 DOCUMENT = "document"
46 IMPORTED = "imported"
47 WIKI = "wiki"
50class HeldOut(StrEnum):
51 """Why a ``_sources`` row holds its path out of later syncs.
53 ``REMOVED``: the user removed it; the row also keeps the path a citation names.
54 ``FAILED``: ingestion produced nothing at that content; a row that holds a
55 document keeps it listed.
56 """
58 FAILED = "failed"
59 REMOVED = "removed"
62class HeldRecord(NamedTuple):
63 """Why a sync holds a path out, and the content hash it holds it out at."""
65 kind: HeldOut
66 file_hash: str
67 reason: str
70# The reason a removal records, shown with the held-out file.
71REMOVED_REASON = "removed via remove (add the source again to restore)"
74class SourceMeta(NamedTuple):
75 """Document-level metadata captured at extraction time.
77 ``title`` is always derivable (extraction metadata or the cleaned filename
78 stem); ``authors`` and ``created_at`` are only present when the extractor
79 reports them. Empty strings persist as NULL so old and new rows read alike.
80 """
82 title: str = ""
83 authors: str = ""
84 created_at: str = ""
87class ChunkWrite(NamedTuple):
88 """One document's chunks plus its source-table update, for a batched write.
90 ``Store.write_chunks_batch`` folds many of these into a single locked
91 transaction so bulk ingest doesn't pay a write-lock acquisition per document.
92 ``page_texts`` rows land in the same transaction, after the cleanup delete
93 and before the source row. ``source_type`` lets the detached import path
94 reuse the same atomic write while still tagging its rows ``IMPORTED``.
95 """
97 source: str
98 file_hash: str
99 records: list[dict]
100 needs_cleanup: bool
101 stat: SourceStat | None = None
102 page_texts: list[dict] | None = None
103 source_type: SourceType = SourceType.DOCUMENT
104 meta: SourceMeta | None = None
107class ChunkType(StrEnum):
108 """Values for the ``chunk_type`` column.
110 Documents ingest as ``RAW``, extracted tables as ``TABLE`` (when table
111 extraction is on), and wiki pages written by the wiki producer as
112 ``WIKI``. Callers filter with ``Store.search(chunk_type=...)``; a ``RAW``
113 filter also covers table chunks, since both are document content.
114 """
116 RAW = "raw"
117 TABLE = "table"
118 WIKI = "wiki"
121# ``schema_version`` is an integer for forward-compat. Bump only if we ever need to
122# add or rename a meta column without forcing every store to drop_all.
123# 2: nomic-embed document prefixes are stamped at ingest (see
124# lilbee.retrieval.embedding_profiles doc_prefix_since).
125# 3: every table keys a source by ``source_id``; ``_sources`` alone holds its path.
126META_SCHEMA_VERSION = 3
128# Always-true predicate used to clear the single-row ``_meta`` table before re-insert.
129# Lance's ``Table.delete`` requires a SQL where clause; this matches every row without
130# coupling the deletion to any specific column's value domain.
131META_DELETE_ALL_PREDICATE = "schema_version IS NOT NULL"
133# Same, for the single-row ``_entity_schema`` table.
134ENTITY_SCHEMA_DELETE_ALL_PREDICATE = "updated_at IS NOT NULL"
137class EntitySchemaState(TypedDict):
138 """Single-row state of the induced entity schema.
140 ``applied`` records whether a full extraction pass completed under this
141 schema; an interrupted pass leaves it False so the next sync redoes the
142 (idempotent) pass. ``source_count`` is how many documents the index held
143 when the schema was induced, which is what the next sync compares against
144 to decide the corpus has drifted far enough to re-induce.
145 """
147 schema_json: str
148 applied: bool
149 source_count: int
150 updated_at: str
153class SearchScope(StrEnum):
154 """What the user wants to search over.
156 Values are used as-is on CLI flags, MCP params, and HTTP query strings.
157 ``BOTH`` resolves to a ``None`` ``chunk_type`` (no filter); the two
158 others map 1:1 to the chunks-table values.
159 """
161 RAW = ChunkType.RAW
162 WIKI = ChunkType.WIKI
163 BOTH = "both"
166def scope_to_chunk_type(scope: SearchScope | str | None) -> ChunkType | None:
167 """Translate a user-facing scope into a ``Store.search`` ``chunk_type`` arg.
169 ``None``/``"both"`` → no filter. ``"raw"`` / ``"wiki"`` → the matching
170 ``ChunkType``. Raises ``ValueError`` on any other string.
171 """
172 if scope is None:
173 return None
174 normalized = SearchScope(scope)
175 if normalized is SearchScope.BOTH:
176 return None
177 return ChunkType(normalized.value)
180class SearchChunk(BaseModel):
181 """A search result from LanceDB.
182 Every store search path sets ``score``: canonical [0, 1] relevance,
183 higher = better. Ranking, filtering, and selection compare only this
184 field; the arm-specific fields below it are provenance.
185 Vector-arm rows carry ``distance``; FTS-arm rows carry ``bm25_score``;
186 reranked rows additionally carry ``rerank_score`` (higher = better).
187 Memory pseudo-rows carry ``memory_id`` and a ``memory:<id>`` source.
188 """
190 model_config = ConfigDict(populate_by_name=True)
192 source: str
193 content_type: str
194 chunk_type: ChunkType = ChunkType.RAW
196 @field_validator("chunk_type", mode="before")
197 @classmethod
198 def _coerce_none_chunk_type(cls, v: str | None) -> str:
199 """LanceDB rows from before the chunk_type column was added return None."""
200 return v if v is not None else ChunkType.RAW
202 # Document title at ingest time; None on rows written before the column
203 # existed (or by writers that carry no title, e.g. wiki pages).
204 title: str | None = None
206 page_start: int
207 page_end: int
208 line_start: int
209 line_end: int
210 chunk: str
211 chunk_index: int
212 vector: list[float] = Field(repr=False)
213 distance: float | None = Field(None, alias="_distance")
214 # Legacy ``_relevance_score`` passthrough. No store path populates it and
215 # no ranking code reads it; it survives only as a display-compatible field
216 # for rows produced by external LanceDB rerankers.
217 relevance_score: float | None = Field(None, validation_alias="_relevance_score")
218 # FTS/BM25-only rows carry a raw, unbounded ``_score``. It lives in its own
219 # field so it never contaminates the canonical ``score``; the
220 # confidence-based expansion-skip reads it (squashed to [0, 1]), and the
221 # relevance filter treats its presence as lexical support.
222 bm25_score: float | None = Field(None, validation_alias="_score")
223 rerank_score: float | None = None
224 # Canonical relevance in [0, 1], set by the store on every search path:
225 # normalized reciprocal-rank fusion on the hybrid path, clamped cosine
226 # similarity on vector-only, list-normalized BM25 on FTS-only probes.
227 score: float | None = None
228 # Set only on memory pseudo-rows built from a recalled MemoryRow (never by
229 # the store): the memory id, so consumers can tell a memory source apart
230 # from a retrieved passage and link it back to /api/memories.
231 memory_id: str | None = None
234class SourceRecord(TypedDict):
235 """A tracked source document record.
237 The stat columns are absent on rows read from stores created before they
238 existed; ``source_stat`` is the accessor that folds absence and the
239 ``SOURCE_STAT_UNKNOWN`` sentinel into ``None``.
240 """
242 filename: str
243 file_hash: str
244 ingested_at: str
245 chunk_count: int
246 source_type: str
247 size_bytes: NotRequired[int]
248 mtime_ns: NotRequired[int]
249 stat_captured_ns: NotRequired[int]
250 # Extraction-time document metadata; absent or None on rows written
251 # before the columns existed, and None when the extractor reported none.
252 title: NotRequired[str | None]
253 authors: NotRequired[str | None]
254 created_at: NotRequired[str | None]
257class SourceHead(TypedDict):
258 """The columns of a tracked source that a listing shows."""
260 filename: str
261 file_hash: str
262 chunk_count: int
263 ingested_at: str
266@dataclass(frozen=True)
267class SourceSummary:
268 """Totals over the tracked sources and the most recently ingested few."""
270 document_count: int
271 total_chunks: int
272 kinds: dict[str, int]
273 newest: list[SourceHead]
276DOCUMENTS_PAGE_DEFAULT = 50
277"""Rows in one page of the document listing when the caller names no limit."""
279DOCUMENTS_PAGE_MAX = 1000
280"""The most rows one page of the document listing may hold."""
283# Sentinel for the stat columns on rows written before they existed (or for
284# detached imports with no backing file). Planning treats it as "unknown: re-hash".
285SOURCE_STAT_UNKNOWN = -1
288class SourceStat(NamedTuple):
289 """File size and mtime captured when a source was hashed, plus the capture time.
291 ``captured_ns`` is the wall-clock time the stat was taken; the sync planner
292 hashes a file whose mtime is not strictly older than it (racily clean).
293 """
295 size_bytes: int
296 mtime_ns: int
297 captured_ns: int = SOURCE_STAT_UNKNOWN
300def source_stat(record: SourceRecord) -> SourceStat | None:
301 """Stored stat for a source row, or None when unknown.
303 The stat columns are nullable ``int64``, so a row can carry an explicit
304 ``None`` (an import, or a write before the columns existed) as well as a
305 missing key or the ``SOURCE_STAT_UNKNOWN`` sentinel. All three mean "no
306 usable stat": return None so the caller re-hashes instead of crashing on
307 ``int(None)``.
308 """
309 size = record.get("size_bytes")
310 mtime = record.get("mtime_ns")
311 captured = record.get("stat_captured_ns")
312 if size is None or mtime is None or size == SOURCE_STAT_UNKNOWN or mtime == SOURCE_STAT_UNKNOWN:
313 return None
314 captured_ns = SOURCE_STAT_UNKNOWN if captured is None else int(captured)
315 return SourceStat(int(size), int(mtime), captured_ns)
318class SourceStatBackfill(NamedTuple):
319 """An already-tracked source row paired with its freshly verified stat."""
321 record: SourceRecord
322 stat: SourceStat
325class SourceMove(NamedTuple):
326 """A file under a new name, and the absent sources of its content it can take the rows of."""
328 candidates: tuple[str, ...]
329 new: str
330 stat: SourceStat | None
333class SourceRelocation(NamedTuple):
334 """One hold of a re-key: the old name each new name took, and how many moves it settled."""
336 taken: dict[str, str]
337 settled: int
340class SourcesCheckpoint(NamedTuple):
341 """A point in the life of the source table.
343 Every source row stamped before ``time`` is committed, and every later
344 commit moves the table off ``version``, which is None with no table.
345 """
347 time: str
348 version: int | None
351class PageTextRecord(TypedDict):
352 """One row of the per-page text dataset, matching ``_page_texts``.
354 The export dataset additionally carries the source's extraction metadata
355 (denormalized onto every page row) so an export/import cycle preserves it;
356 these are absent on rows read from the ``_page_texts`` table and on datasets
357 exported before the columns existed.
358 """
360 source: str
361 page: int
362 text: str
363 content_type: str
364 title: NotRequired[str | None]
365 authors: NotRequired[str | None]
366 created_at: NotRequired[str | None]
369class CitationRecord(TypedDict):
370 """A citation linking a wiki chunk to a specific source location."""
372 wiki_source: str
373 wiki_chunk_index: int
374 citation_key: str
375 claim_type: str
376 source_filename: str
377 source_hash: str
378 page_start: int
379 page_end: int
380 line_start: int
381 line_end: int
382 excerpt: str
383 created_at: str
386class MemoryKind(StrEnum):
387 """Whether a memory is an always-injected preference or a similarity-recalled fact."""
389 PREFERENCE = "preference"
390 FACT = "fact"
393class MemorySource(StrEnum):
394 """Provenance of a memory: user-typed, LLM-extracted, or agent-written."""
396 MANUAL = "manual"
397 EXTRACTED = "extracted"
398 AGENT = "agent"
401# Memory owner namespaces. ``"local"`` is the single human (TUI/CLI/REST); agents own
402# ``"agent:<id>"`` namespaces. The prefix lives only here so it is never hand-spliced.
403LOCAL_OWNER = "local"
404AGENT_OWNER_PREFIX = "agent:"
406# Memory pseudo-source values on SearchChunk: ``"memory:<id>"`` with a
407# ``"memory"`` content type. Built by retrieval, never stored in the chunks
408# table, so recalled facts show in sources without a backing document.
409MEMORY_SOURCE_PREFIX = "memory:"
410MEMORY_CONTENT_TYPE = "memory"
413def agent_owner(agent_id: str) -> str:
414 """Owner string for an agent identity (``"opencode"`` -> ``"agent:opencode"``)."""
415 return f"{AGENT_OWNER_PREFIX}{agent_id}"
418def is_agent_owner(owner: str) -> bool:
419 """True when *owner* is an agent namespace rather than the local human."""
420 return owner.startswith(AGENT_OWNER_PREFIX)
423def memory_source(memory_id: str) -> str:
424 """Source value for a recalled memory (``"abc"`` -> ``"memory:abc"``)."""
425 return f"{MEMORY_SOURCE_PREFIX}{memory_id}"
428def is_memory_source(source: str) -> bool:
429 """True when *source* names a recalled memory rather than a document."""
430 return source.startswith(MEMORY_SOURCE_PREFIX)
433class MemoryRow(BaseModel):
434 """A long-term memory entry in the per-library ``_memories`` table.
436 Built from a LanceDB row via ``MemoryRow(**row)`` (which coerces the ``kind``
437 and ``source`` strings to enums) and written back via ``model_dump(mode="json")``.
438 Extra keys like a search ``_distance`` are ignored on construction.
439 """
441 model_config = ConfigDict(extra="ignore")
443 id: str
444 owner: str
445 shared: bool
446 kind: MemoryKind
447 source: MemorySource
448 text: str
449 vector: list[float] = Field(repr=False)
450 created_at: str
451 updated_at: str
454class StoreMeta(TypedDict):
455 """Single-row store metadata recording the embedding model used to build the store.
457 Compatibility is checked before every read and write. When ``cfg.embedding_model``
458 or ``cfg.embedding_dim`` drifts from the persisted row, the store refuses to serve
459 until ``lilbee rebuild`` (CLI) or ``POST /api/sync {"force_rebuild": true}`` (HTTP)
460 rewrites the chunks under the new model.
462 ``updated_at`` is an ISO 8601 UTC timestamp produced by ``datetime.isoformat()``;
463 kept as ``str`` to match the LanceDB ``utf8`` schema column.
464 """
466 embedding_model: str
467 embedding_dim: int
468 schema_version: int
469 updated_at: str
472class IndexMismatch(BaseModel):
473 """The embedder drift between a persisted index and the configured model.
475 ``adoptable`` is true when the dimensions agree, so switching the embedding
476 model back to ``persisted_model`` makes the index searchable without a rebuild.
477 """
479 persisted_model: str
480 persisted_dim: int
481 current_model: str
482 current_dim: int
483 adoptable: bool
484 message: str
487class EmbeddingModelMismatchError(RuntimeError):
488 """Raised when stored vectors were built with a different embedder than ``cfg``.
490 Carries the persisted and configured refs and dims so each surface renders its
491 own recovery affordance (TUI prompt, CLI command, REST body) from the facts.
492 """
494 def __init__(
495 self,
496 *,
497 persisted_model: str,
498 persisted_dim: int,
499 current_model: str,
500 current_dim: int,
501 ) -> None:
502 self.persisted_model = persisted_model
503 self.persisted_dim = persisted_dim
504 self.current_model = current_model
505 self.current_dim = current_dim
506 super().__init__(self._build_message())
508 @property
509 def dims_match(self) -> bool:
510 """True when the index is adoptable by switching embedder alone (same dim)."""
511 return self.persisted_dim == self.current_dim
513 def describe(self) -> IndexMismatch:
514 """The drift as a serializable record for sync results and status payloads."""
515 return IndexMismatch(
516 persisted_model=self.persisted_model,
517 persisted_dim=self.persisted_dim,
518 current_model=self.current_model,
519 current_dim=self.current_dim,
520 adoptable=self.dims_match,
521 message=str(self),
522 )
524 def _build_message(self) -> str:
525 if self.dims_match:
526 return (
527 f"This index was built with embedding model '{self.persisted_model}', "
528 f"but lilbee is configured to use '{self.current_model}'. Configure lilbee "
529 f"to use '{self.persisted_model}' to search this index, or rebuild it under "
530 f"'{self.current_model}'."
531 )
532 return (
533 f"This index was built with embedding model '{self.persisted_model}' "
534 f"(dim {self.persisted_dim}), which differs from the current "
535 f"'{self.current_model}' (dim {self.current_dim}). The dimensions differ, "
536 f"so rebuild the index under '{self.current_model}' to use it."
537 )
540@dataclass
541class RemoveResult:
542 """Result of a remove: *removed* can also name held-out files and root labels."""
544 removed: list[str]
545 not_found: list[str]