Coverage for src/lilbee/data/store/core.py: 100%
1025 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""The ``Store`` class: high-level LanceDB read/write API used across lilbee."""
3from __future__ import annotations
5import logging
6import math
7import os
8from collections.abc import Callable, Iterable, Sequence
9from contextlib import AbstractContextManager
10from datetime import UTC, datetime
11from pathlib import Path
12from typing import TYPE_CHECKING
14import pyarrow as pa
16from lilbee.core.config import (
17 CHUNK_CONCEPTS_TABLE,
18 CHUNKS_TABLE,
19 CITATIONS_TABLE,
20 ENTITIES_TABLE,
21 ENTITY_SCHEMA_TABLE,
22 INGEST_SOURCE_COLUMNS,
23 MEMORIES_TABLE,
24 META_TABLE,
25 PAGE_TEXTS_TABLE,
26 SOURCES_TABLE,
27 WIKI_MENTIONS_TABLE,
28 Config,
29)
30from lilbee.core.vectors import Vector
31from lilbee.retrieval.embedding_profiles import resolve_embedding_profile
32from lilbee.runtime.lock import LOCK_TIMEOUT, LockTimeoutError, write_lock
34from .fusion import adaptive_weight_scale, fuse_arms, normalized_bm25, vector_similarity
35from .lance_helpers import (
36 _CHUNK_COLUMN,
37 _chunk_type_predicate,
38 _has_fts_index,
39 _has_scalar_index,
40 _has_vector_index,
41 _safe_delete_unlocked,
42 _sources_search_filter,
43 ensure_table,
44 escape_sql_string,
45 refs_compatible,
46 table_names,
47)
48from .ranking import mmr_rerank
49from .schema import (
50 _citations_schema,
51 _entity_schema_state_schema,
52 _meta_schema,
53 _page_texts_schema,
54 _sources_schema,
55 _wiki_mentions_schema,
56)
57from .types import (
58 ENTITY_SCHEMA_DELETE_ALL_PREDICATE,
59 META_DELETE_ALL_PREDICATE,
60 META_SCHEMA_VERSION,
61 READ_CONSISTENCY_INTERVAL,
62 SOURCE_STAT_UNKNOWN,
63 ChunkType,
64 ChunkWrite,
65 CitationRecord,
66 EmbeddingModelMismatchError,
67 EntitySchemaState,
68 MemoryKind,
69 MemoryRow,
70 PageTextRecord,
71 RemoveResult,
72 SearchChunk,
73 SourceMeta,
74 SourceRecord,
75 SourceStat,
76 SourceStatBackfill,
77 SourceType,
78 StoreMeta,
79)
81if TYPE_CHECKING:
82 import lance
83 import lancedb
84 import lancedb.table
85 from lancedb.index import FTS
87log = logging.getLogger(__name__)
89# Batched ingest flushes contend with long store operations (a search-triggered
90# FTS optimize can hold the lock past the interactive 30s), and failing the
91# flush replans and re-embeds the whole batch. Give the batch path more
92# patience than interactive writes before it gives up.
93BATCH_LOCK_TIMEOUT = 120.0
94# Lock budget for index builds reached from the read path: a query must not
95# stall behind a long ingest, so it skips the build and retries next search.
96_READ_LOCK_TIMEOUT = 2.0
99def _drop_unsupported_far_rows(
100 results: list[SearchChunk], max_distance: float
101) -> list[SearchChunk]:
102 """Apply ``max_distance`` to rows whose only signal is the vector arm.
104 A row the BM25 arm also matched keeps lexical support regardless of its
105 vector distance; dropping it on distance alone would re-bury exactly the
106 identifier hits rank fusion exists to preserve.
107 """
108 if max_distance <= 0:
109 return results
110 return [
111 r
112 for r in results
113 if r.bm25_score is not None or r.distance is None or r.distance <= max_distance
114 ]
117_MAX_THRESHOLD = 1.0
118_MAX_FILTER_ITERATIONS = 20 # safety cap to prevent runaway loops
121def _is_fts_position_overflow(exc: Exception) -> bool:
122 """True when *exc* is LanceDB's positional-FTS list-encoding overflow.
124 A positional index (built by an intermediate dev commit) raises e.g.
125 "Max offset N exceeds length of values M" on optimize(); a positionless
126 rebuild is the remediation. Matched on message because LanceDB raises it
127 as a generic error type.
128 """
129 msg = str(exc).lower()
130 return "offset" in msg and "exceeds" in msg
133def _lexical_rows(
134 table: lancedb.table.Table,
135 query_text: str,
136 limit: int,
137 chunk_type: ChunkType | None,
138 column: str = _CHUNK_COLUMN,
139) -> list[SearchChunk]:
140 """BM25 rows for *query_text* over a single FTS *column*.
142 ``MatchQuery`` pins the column and matches plain terms, so an unpinned search
143 cannot widen to the title index and a quoted span cannot reach LanceDB as a
144 phrase (which the positionless index rejects). This is the one place FTS
145 queries are built; every arm goes through it.
146 """
147 from lancedb.query import MatchQuery
149 query = table.search(MatchQuery(query_text, column), query_type="fts").limit(limit)
150 if chunk_type:
151 query = query.where(_chunk_type_predicate(chunk_type))
152 return [SearchChunk(**r) for r in query.to_list()]
155# Vector ANN index. IVF_PQ compresses vectors so search scales to millions;
156# refine_factor re-ranks the PQ candidates against full vectors to recover recall.
157# The index type is carried by the lancedb IvfPq config at build time.
158_VECTOR_METRIC = "cosine"
159_ANN_NPROBES_FLOOR = 20
160# Fraction of IVF partitions probed per query. 0.05 was the "fast" end of the
161# recall/latency curve and measurably cost recall at scale: on the 8.8M-passage
162# MS MARCO index it gave recall@100 of 67%, where probing more partitions
163# reached 87% (docs/benchmarks/retrieval-msmarco.md). 0.15 is the "balanced /
164# high-recall" range for IVF and recovers most of that gap; refine_factor below
165# still re-ranks the survivors against full vectors.
166_ANN_NPROBES_PARTITION_FRACTION = 0.15
167_ANN_REFINE_FACTOR = 10
169# Stat columns of ``_sources``; mirrors the field names in ``schema._sources_schema``
170# and ``types.SourceRecord``. Legacy tables that predate these columns are migrated
171# in place with the SOURCE_STAT_UNKNOWN sentinel.
172_SOURCE_STAT_COLUMNS = ("size_bytes", "mtime_ns", "stat_captured_ns")
174# Extraction-metadata columns of ``_sources``; nullable strings, so legacy
175# tables migrate in place with NULL (meaning "extractor reported nothing").
176_SOURCE_META_COLUMNS = ("title", "authors", "created_at")
178# Document-title column of the chunks table; nullable so pre-title rows and
179# writers that carry no title (wiki pages) read as NULL.
180_TITLE_COLUMN = "title"
182# (table, source column) pairs deleted when a source's rows are replaced. The
183# concept nodes/edges tables carry no source column (corpus-level aggregates),
184# so only the per-chunk concept mapping is source-scoped.
185_PER_SOURCE_TABLES = (
186 *(
187 (name, INGEST_SOURCE_COLUMNS[name])
188 for name in (CHUNKS_TABLE, PAGE_TEXTS_TABLE, CHUNK_CONCEPTS_TABLE, ENTITIES_TABLE)
189 ),
190 # Per-(subject, source), so removing a source drops its mention evidence
191 # here with its chunks; the wiki refresh only ever re-adds a source, never
192 # has to remember to subtract a deleted one. Not in INGEST_SOURCE_COLUMNS
193 # (a wiki table, not an ingest one), so its column is named directly.
194 (WIKI_MENTIONS_TABLE, "source"),
195)
197# (table, source column) pairs re-keyed when a source is relocated (moved on
198# disk, same content). Extends the per-source set with the wiki citation's raw
199# source_filename so citations keep pointing at the source after a move.
200_RELOCATABLE_TABLES = (
201 *_PER_SOURCE_TABLES,
202 (CITATIONS_TABLE, INGEST_SOURCE_COLUMNS[CITATIONS_TABLE]),
203)
205# Sentinel: relocation must leave the stored title untouched (extraction-derived).
206_KEEP_TITLE = "\x00keep"
208# Stat backfills replace this many source rows per locked write: the first
209# sync after a stat-column upgrade backfills every source, and an unchunked
210# replace would join millions of filenames into one delete predicate.
211_SOURCE_STAT_BATCH_ROWS = 2000
213# Rows per Arrow batch when the aggregate scan walks the whole chunks table;
214# bounds the decoded-text working set while the scan stays columnar.
215_TERM_SCAN_BATCH_ROWS = 20_000
217# The title arm collapses each matched document to one row, so it over-fetches
218# to gather enough distinct documents before deduping. Bounded so a title that
219# hits a huge document can't scan the whole corpus.
220_TITLE_FETCH_FACTOR = 20
221_TITLE_MIN_FETCH = 200
222_TITLE_FETCH_CEILING = 4096
225def _ann_nprobes(row_count: int) -> int:
226 """Partitions to probe: a fixed fraction of the IVF partition count (~sqrt(N)), floored."""
227 partitions = math.isqrt(max(row_count, 0))
228 return max(_ANN_NPROBES_FLOOR, math.ceil(partitions * _ANN_NPROBES_PARTITION_FRACTION))
231def _check_vector_dims(records: list[dict], embedding_dim: int) -> None:
232 """Raise ``ValueError`` when any record's vector is not *embedding_dim* wide."""
233 for rec in records:
234 vec = rec.get("vector", [])
235 if len(vec) != embedding_dim:
236 raise ValueError(
237 f"Vector dimension mismatch: expected {embedding_dim}, "
238 f"got {len(vec)} (source={rec.get('source', '?')})"
239 )
242def _citations_for_wiki_predicate(wiki_source: str) -> str:
243 """SQL predicate selecting every citation row belonging to *wiki_source*."""
244 return f"wiki_source = '{escape_sql_string(wiki_source)}'"
247def _get_distance(chunk: SearchChunk) -> float:
248 """Extract distance as a sortable float (inf for None)."""
249 return chunk.distance if chunk.distance is not None else float("inf")
252def _count_within_threshold(sorted_results: list[SearchChunk], threshold: float) -> int:
253 """Count results whose distance is within the given threshold."""
254 for i, r in enumerate(sorted_results):
255 if _get_distance(r) > threshold:
256 return i
257 return len(sorted_results)
260class Store:
261 """LanceDB vector store: wraps all DB operations with config-driven defaults."""
263 def __init__(self, config: Config) -> None:
264 self._config = config
265 self._fts_ready: bool = False
266 self._title_fts_ready: bool = False
267 self._doc_prefix_warned: bool = False
268 # Scalar indexes (source/chunk_type) are built at ingest; a serve-only
269 # store builds them lazily from the search path.
270 self._scalar_ready: bool = False
271 self._db: lancedb.DBConnection | None = None
272 # Cache of {filename: ingested_at} rebuilt only when sources
273 # mutate; callers (temporal filter) hit it per-query.
274 self._source_ingested_cache: dict[str, str] | None = None
276 def _fts_config(self) -> FTS:
277 """Shared FTS index config: positionless (with_position=True overflows
278 LanceDB's list encoding on optimize()) and stemmed for the configured
279 corpus language."""
280 from lancedb.index import FTS
282 return FTS(with_position=False, language=self._config.fts_language)
284 def _index_build_lock(self, blocking: bool) -> AbstractContextManager[None]:
285 """Write lock for index builds: full budget from ingest, short from the read path."""
286 return self._write_lock() if blocking else self._write_lock(_READ_LOCK_TIMEOUT)
288 def _write_lock(self, timeout: float = LOCK_TIMEOUT) -> AbstractContextManager[None]:
289 """Acquire the write lock keyed on *this* store's data directory.
291 A per-instance ``Lilbee`` writes to its own ``lancedb_dir``; locking the
292 global ``cfg`` dir instead would leave those writes uncoordinated across
293 processes.
294 """
295 return write_lock(self._config.lancedb_dir, timeout)
297 def _invalidate_source_cache(self) -> None:
298 """Drop the cached {filename: ingested_at} map."""
299 self._source_ingested_cache = None
301 def source_ingested_at_map(self) -> dict[str, str]:
302 """Return {filename: ingested_at} for every source, cached until mutation.
304 Best-effort: a reader racing a concurrent invalidation can store a
305 pre-mutation snapshot. The only consumer (temporal query filter)
306 treats a missing/stale entry as "do not filter," so staleness
307 degrades ranking precision, never correctness.
308 """
309 if self._source_ingested_cache is not None:
310 return self._source_ingested_cache
311 mapping = {s["filename"]: s.get("ingested_at", "") for s in self.get_sources()}
312 self._source_ingested_cache = mapping
313 return mapping
315 def _chunks_schema(self) -> pa.Schema:
316 return pa.schema(
317 [
318 pa.field("source", pa.utf8()),
319 pa.field("content_type", pa.utf8()),
320 pa.field("chunk_type", pa.utf8()),
321 pa.field("page_start", pa.int32()),
322 pa.field("page_end", pa.int32()),
323 pa.field("line_start", pa.int32()),
324 pa.field("line_end", pa.int32()),
325 pa.field("chunk", pa.utf8()),
326 pa.field("chunk_index", pa.int32()),
327 pa.field(_TITLE_COLUMN, pa.utf8()),
328 pa.field("vector", pa.list_(pa.float32(), self._config.embedding_dim)),
329 ]
330 )
332 def _chunks_table(self) -> lancedb.table.Table:
333 """Open/create the chunks table, adding the title column to pre-title tables."""
334 table = ensure_table(self.get_db(), CHUNKS_TABLE, self._chunks_schema())
335 if _TITLE_COLUMN not in table.schema.names:
336 table.add_columns({_TITLE_COLUMN: "CAST(NULL AS STRING)"})
337 self._backfill_stem_titles_unlocked(table)
338 return table
340 def _backfill_stem_titles_unlocked(self, table: lancedb.table.Table) -> None:
341 """Backfill filename-stem titles for pre-upgrade rows. Caller holds ``write_lock()``.
343 Without this the title arm only matches documents ingested after the
344 upgrade. Extracted (H1/EXIF) titles still need ``lilbee rebuild``;
345 failure leaves NULLs, the pre-backfill behavior.
346 """
347 from lilbee.data.title import derive_title # circular at module scope
349 try:
350 rows = table.search().select(["source"]).limit(None).to_list()
351 sources = sorted({r["source"] for r in rows})
352 filled = 0
353 for source in sources:
354 title = derive_title(source)
355 if not title:
356 continue
357 escaped = source.replace("'", "''")
358 table.update(where=f"source = '{escaped}'", values={_TITLE_COLUMN: title})
359 filled += 1
360 log.info(
361 "Backfilled filename titles for %d of %d existing sources; run "
362 "`lilbee rebuild` to derive titles from document content",
363 filled,
364 len(sources),
365 )
366 except Exception:
367 log.warning(
368 "Title backfill failed; pre-upgrade rows keep NULL titles until `lilbee rebuild`",
369 exc_info=True,
370 )
372 def get_meta(self) -> StoreMeta | None:
373 """Return the persisted store metadata row, or ``None`` if unset."""
374 table = self.open_table(META_TABLE)
375 if table is None:
376 return None
377 rows = table.search().limit(None).to_list()
378 if not rows:
379 return None
380 # _meta is meant to hold one row, but a swallowed delete on rewrite could
381 # leave a stale one behind; take the newest so identity reads stay
382 # deterministic rather than returning an arbitrary row.
383 row = max(rows, key=lambda r: r["updated_at"])
384 return StoreMeta(
385 embedding_model=row["embedding_model"],
386 embedding_dim=int(row["embedding_dim"]),
387 schema_version=int(row["schema_version"]),
388 updated_at=row["updated_at"],
389 )
391 def _write_meta_unlocked(self, *, embedding_model: str, embedding_dim: int) -> None:
392 """Overwrite the single ``_meta`` row with the supplied identity.
394 Caller must hold ``write_lock()``. Args are passed explicitly rather than
395 re-read from ``self._config`` so the caller can snapshot cfg at a coherent
396 instant and not race with a concurrent ``set_embedding_model``.
397 """
398 db = self.get_db()
399 table = ensure_table(db, META_TABLE, _meta_schema())
400 _safe_delete_unlocked(table, META_DELETE_ALL_PREDICATE)
401 table.add(
402 [
403 {
404 "embedding_model": embedding_model,
405 "embedding_dim": embedding_dim,
406 "schema_version": META_SCHEMA_VERSION,
407 "updated_at": datetime.now(UTC).isoformat(),
408 }
409 ]
410 )
412 def _has_chunks(self) -> bool:
413 """Return True when the chunks table exists and has at least one row."""
414 chunks = self.open_table(CHUNKS_TABLE)
415 return chunks is not None and chunks.count_rows() > 0
417 def has_chunks(self) -> bool:
418 """Public predicate: True iff the store currently holds at least one chunk."""
419 return self._has_chunks()
421 def initialize_meta_if_legacy(self) -> bool:
422 """Pin a legacy store's identity to the current cfg if not already set.
424 Returns ``True`` when a meta row was just written. No-op when meta already
425 exists or no chunks are present. This is the path that converts a
426 pre-upgrade store (chunks present, no ``_meta``) into a gated store. It
427 snapshots cfg under the write lock to keep the recorded identity coherent
428 with what the gate is comparing against.
429 """
430 if self.get_meta() is not None:
431 return False
432 if not self._has_chunks():
433 return False
434 embedding_model = self._config.embedding_model
435 embedding_dim = self._config.embedding_dim
436 with self._write_lock():
437 # Re-check under the lock so two callers do not both warn-and-write.
438 if self.get_meta() is not None:
439 return False
440 log.warning(
441 "Legacy store has chunks but no _meta row. Initializing _meta from "
442 "the current configuration (embedding_model=%s, embedding_dim=%d). "
443 "If you changed embedding_model before upgrading, run `lilbee rebuild` "
444 "to ensure the store is consistent.",
445 embedding_model,
446 embedding_dim,
447 )
448 self._write_meta_unlocked(embedding_model=embedding_model, embedding_dim=embedding_dim)
449 return True
451 def _ensure_embedding_compat(self) -> None:
452 """Raise when the persisted embedding identity drifts from cfg.
454 Pure check, no side effects. Migration of legacy stores (chunks present,
455 no ``_meta``) is the caller's responsibility via ``initialize_meta_if_legacy``;
456 rewriting a legacy bare-repo ``_meta`` row to the canonical full ref is
457 the caller's responsibility via ``canonicalize_meta_if_legacy``. This
458 method stays safe to call from inside an existing ``write_lock()`` (no
459 recursive lock attempt). cfg fields are snapshotted at entry so the
460 comparison is coherent even if another thread mutates them mid-call.
461 """
462 current_model = self._config.embedding_model
463 current_dim = self._config.embedding_dim
464 meta = self.get_meta()
465 if meta is None:
466 return
467 if refs_compatible(
468 meta["embedding_model"], current_model, meta["embedding_dim"], current_dim
469 ):
470 return
471 raise EmbeddingModelMismatchError(
472 persisted_model=meta["embedding_model"],
473 persisted_dim=meta["embedding_dim"],
474 current_model=current_model,
475 current_dim=current_dim,
476 )
478 def _warn_stale_doc_prefix(self) -> None:
479 """Warn once when the embedding family's document prefix postdates this store.
481 Queries would carry the family prefix while stored documents do not;
482 the mismatch degrades quality silently until a rebuild re-embeds.
483 """
484 if self._doc_prefix_warned:
485 return
486 self._doc_prefix_warned = True
487 meta = self.get_meta()
488 if meta is None:
489 return
490 profile = resolve_embedding_profile(self._config.embedding_model)
491 if profile.doc_prefix and meta["schema_version"] < profile.doc_prefix_since:
492 log.warning(
493 "This index predates '%s' document prefixes: queries are prefixed "
494 "but stored documents are not. Run `lilbee rebuild` to re-embed.",
495 self._config.embedding_model,
496 )
498 def assert_embedding_compatible(self) -> None:
499 """Run the full embedding-identity gate (legacy init, canonicalize, check).
501 Mirrors the gate ``search`` applies. Callers that write under a fresh
502 embedder (import) use this to fail before any destructive work when the
503 store was built by a different model.
504 """
505 self.initialize_meta_if_legacy()
506 self.canonicalize_meta_if_legacy()
507 self._ensure_embedding_compat()
509 def _needs_canonical_meta_rewrite(
510 self, meta: StoreMeta | None, current_model: str, current_dim: int
511 ) -> bool:
512 """True iff *meta* is the legacy form and refs-compatible with current cfg."""
513 if meta is None or meta["embedding_model"] == current_model:
514 return False
515 return refs_compatible(
516 meta["embedding_model"], current_model, meta["embedding_dim"], current_dim
517 )
519 def canonicalize_meta_if_legacy(self) -> bool:
520 """Rewrite a legacy bare-repo ``_meta`` row to the canonical full ref.
522 Pre-canonical lilbee persisted only ``<org>/<repo>`` in
523 ``_meta.embedding_model``. The current code persists the full
524 ``<org>/<repo>/<filename>.gguf``. When the two refer to the same
525 model under :func:`refs_compatible` but differ as raw strings, the
526 meta row is rewritten so the legacy name never surfaces. Returns
527 ``True`` on write; ``False`` when missing, already canonical, or
528 incompatible (the gate handles incompatibility).
529 """
530 current_model = self._config.embedding_model
531 current_dim = self._config.embedding_dim
532 if not self._needs_canonical_meta_rewrite(self.get_meta(), current_model, current_dim):
533 return False
534 with self._write_lock():
535 meta = self.get_meta() # re-read under the lock for racing callers
536 if not self._needs_canonical_meta_rewrite(meta, current_model, current_dim):
537 return False
538 assert meta is not None # filtered above # noqa: S101
539 log.info(
540 "Migrating legacy embedding ref in store meta: %r -> %r",
541 meta["embedding_model"],
542 current_model,
543 )
544 self._write_meta_unlocked(embedding_model=current_model, embedding_dim=current_dim)
545 return True
547 def get_db(self) -> lancedb.DBConnection:
548 if self._db is None:
549 import lancedb as _lancedb
551 self._config.lancedb_dir.mkdir(parents=True, exist_ok=True)
552 self._db = _lancedb.connect(
553 str(self._config.lancedb_dir),
554 read_consistency_interval=READ_CONSISTENCY_INTERVAL,
555 )
556 return self._db
558 def open_table(self, name: str) -> lancedb.table.Table | None:
559 """Open a table if it exists, otherwise return None."""
560 db = self.get_db()
561 if name not in table_names(db):
562 return None
563 return db.open_table(name)
565 def ensure_fts_index(self, *, blocking: bool = True) -> None:
566 """Create the chunks FTS index, or run ``optimize()`` once it exists.
568 ``optimize()`` folds newly added rows into the FTS index and also
569 runs LanceDB's default compaction + version pruning (default prune
570 window: 7 days). Work scales with recent deltas rather than total
571 chunk count, so large corpora no longer pay the full
572 ``create_index(config=FTS(), replace=True)`` rebuild cost on every sync.
574 ``blocking=False`` (the search path) marks an existing index ready
575 without the lock and skips maintenance when another process holds it,
576 so a long concurrent ingest cannot stall or fail a query.
577 """
578 probe = self.open_table(CHUNKS_TABLE)
579 if probe is None:
580 return
581 if _has_fts_index(probe):
582 self._fts_ready = True
583 try:
584 with self._index_build_lock(blocking):
585 self._ensure_fts_index_unlocked()
586 except LockTimeoutError:
587 if blocking:
588 raise
589 log.debug("Skipped FTS index maintenance; another process holds the write lock")
591 def _ensure_fts_index_unlocked(self) -> None:
592 """Body of ``ensure_fts_index``. Caller holds ``write_lock()``."""
593 table = self.open_table(CHUNKS_TABLE)
594 if table is None:
595 return
596 try:
597 if _has_fts_index(table):
598 self._fts_ready = True
599 try:
600 # One optimize folds new rows into every index on the table.
601 table.optimize()
602 log.debug("FTS index optimized on '%s'", CHUNKS_TABLE)
603 except Exception as exc:
604 if _is_fts_position_overflow(exc):
605 # Positional indexes overflow on optimize(); rebuild
606 # positionless once.
607 self._rebuild_fts_positionless(table)
608 else:
609 log.warning(
610 "FTS optimize() failed; the existing index still serves hybrid search",
611 exc_info=True,
612 )
613 else:
614 # Positionless: with_position=True overflows LanceDB's list
615 # encoding on optimize(), and nothing issues phrase queries.
616 table.create_index(_CHUNK_COLUMN, config=self._fts_config(), replace=False)
617 self._fts_ready = True
618 log.debug("FTS index created on '%s'", CHUNKS_TABLE)
619 # Only the opt-in title arm needs the title index.
620 if self._config.title_search:
621 self._ensure_title_fts_unlocked(table)
622 except Exception:
623 log.debug("FTS index ensure failed (empty table?)", exc_info=True)
625 def _ensure_title_fts_unlocked(self, table: lancedb.table.Table) -> None:
626 """Create the title FTS index when the column exists. Caller holds ``write_lock()``.
628 Failure never blocks the chunk index: the title arm feature-detects the
629 index per query, so a store without it simply searches without titles.
630 """
631 if _TITLE_COLUMN not in table.schema.names or _has_fts_index(table, _TITLE_COLUMN):
632 self._title_fts_ready = _has_fts_index(table, _TITLE_COLUMN)
633 return
634 try:
635 # Positionless for the same reason as the chunk index.
636 table.create_index(_TITLE_COLUMN, config=self._fts_config(), replace=False)
637 self._title_fts_ready = True
638 log.debug("Title FTS index created on '%s'", CHUNKS_TABLE)
639 except Exception:
640 # Only reached with title_search enabled, so a silent failure means
641 # the user's opted-in title arm quietly does nothing. Warn, don't hide.
642 log.warning(
643 "Title FTS index creation failed; the title-search arm will "
644 "contribute nothing until it can be built",
645 exc_info=True,
646 )
648 def ensure_title_fts_index(self, *, blocking: bool = True) -> None:
649 """Build the title FTS index for a title_search toggle after startup.
651 Without this, a process that latched ``_fts_ready`` before the toggle
652 never builds the index and the title arm stays silently dead until the
653 next ingest or restart.
654 """
655 table = self.open_table(CHUNKS_TABLE)
656 if table is None:
657 return
658 if _has_fts_index(table, _TITLE_COLUMN):
659 self._title_fts_ready = True
660 return
661 try:
662 with self._index_build_lock(blocking):
663 self._ensure_title_fts_unlocked(table)
664 except LockTimeoutError:
665 if blocking:
666 raise
667 log.debug("Skipped title FTS build; another process holds the write lock")
669 def _rebuild_fts_positionless(self, table: lancedb.table.Table) -> None:
670 """Replace positional FTS indexes with positionless ones. Caller holds the lock.
672 The one-shot remediation for a store whose index was built
673 ``with_position=True`` and now overflows on every ``optimize()``. The
674 title index is rebuilt too when the title arm is enabled.
675 """
676 try:
677 table.create_index(_CHUNK_COLUMN, config=self._fts_config(), replace=True)
678 if self._config.title_search and _TITLE_COLUMN in table.schema.names:
679 table.create_index(_TITLE_COLUMN, config=self._fts_config(), replace=True)
680 log.warning("Rebuilt the FTS index positionless after a positional-index overflow")
681 except Exception:
682 log.warning(
683 "Positionless FTS rebuild failed; the existing index still serves",
684 exc_info=True,
685 )
687 # Tables and (column, index_type) pairs the query path filters by.
688 # chunk_concepts serves the concept boost (ConceptGraph._chunk_concepts_from);
689 # without its index every boosted query full-scans the table.
690 _SCALAR_TARGETS: tuple[tuple[str, tuple[tuple[str, str], ...]], ...] = (
691 (CHUNKS_TABLE, (("source", "BTREE"), ("chunk_type", "BITMAP"))),
692 (CHUNK_CONCEPTS_TABLE, (("chunk_source", "BTREE"),)),
693 )
695 def ensure_scalar_indexes(self, *, blocking: bool = True) -> None:
696 """Build scalar indexes on the columns lilbee filters by.
698 ``source`` and ``chunk_type`` predicates run as prefilters (LanceDB's
699 default), but without an index each is a full-table scan. Readiness
700 latches only when every target table exists and is covered, so a table
701 created later (chunk_concepts under serve ordering) still gets its
702 index on a following call. The lock is taken only when there is
703 something to build; ``blocking=False`` (the search path) skips the
704 build when another process holds it instead of stalling the query.
705 """
706 pending = []
707 complete = True
708 for name, columns in self._SCALAR_TARGETS:
709 table = self.open_table(name)
710 if table is None:
711 complete = False
712 continue
713 names = table.schema.names
714 if any(c in names and not _has_scalar_index(table, c) for c, _ in columns):
715 pending.append((name, columns))
716 if not pending:
717 self._scalar_ready = complete
718 return
719 try:
720 with self._index_build_lock(blocking):
721 for name, columns in pending:
722 self._ensure_scalar_index_on(name, columns)
723 self._scalar_ready = complete
724 except LockTimeoutError:
725 if blocking:
726 raise
727 log.debug("Skipped scalar index build; another process holds the write lock")
729 def _ensure_scalar_index_on(
730 self, table_name: str, columns: tuple[tuple[str, str], ...]
731 ) -> None:
732 """Build the given (column, index_type) scalar indexes on *table_name*.
734 Caller holds ``write_lock()``. Each column gets its own try so one
735 failure does not skip the rest; a failure on a populated table warns
736 (the prefilter speedup is silently lost) while an empty table's is debug.
737 """
738 table = self.open_table(table_name)
739 if table is None:
740 return
741 names = table.schema.names
742 fail_level = logging.WARNING if table.count_rows() > 0 else logging.DEBUG
743 for column, index_type in columns:
744 if column not in names or _has_scalar_index(table, column):
745 continue
746 try:
747 table.create_scalar_index(column, index_type=index_type, replace=False)
748 log.debug("Scalar (%s) index created on '%s.%s'", index_type, table_name, column)
749 except Exception:
750 log.log(
751 fail_level,
752 "Scalar index create failed on '%s.%s'",
753 table_name,
754 column,
755 exc_info=True,
756 )
758 def ensure_vector_index(self, *, force: bool = False) -> bool:
759 """Build or refresh the ANN vector index when the corpus is large enough.
761 Below ``cfg.ann_index_threshold`` (or when it is 0) the store keeps exact
762 flat search, which is faster and exact for small vaults and is all a
763 laptop needs. Once an index exists, ``optimize()`` folds new rows in.
764 Pass ``force=True`` to build regardless of the threshold (publish flow).
765 Returns True when an index was created or refreshed.
766 """
767 threshold = self._config.ann_index_threshold
768 with self._write_lock():
769 table = self.open_table(CHUNKS_TABLE)
770 if table is None:
771 return False
772 if _has_vector_index(table):
773 table.optimize()
774 log.debug("Vector index optimized on '%s'", CHUNKS_TABLE)
775 return True
776 if not force and (threshold <= 0 or table.count_rows() < threshold):
777 return False
778 from lancedb.index import IvfPq
780 try:
781 table.create_index("vector", config=IvfPq(distance_type=_VECTOR_METRIC))
782 log.info("Vector ANN index created on '%s'", CHUNKS_TABLE)
783 return True
784 except Exception:
785 log.warning(
786 "Vector ANN index build failed on '%s' at %d rows; search falls back "
787 "to exact flat scan, which is slow at this scale. Free up memory/disk "
788 "and re-run to rebuild the index.",
789 CHUNKS_TABLE,
790 table.count_rows(),
791 exc_info=True,
792 )
793 return False
795 def _add_chunks_unlocked(self, records: list[dict]) -> int:
796 """Add chunk records and return the count. Caller must hold ``write_lock()``."""
797 embedding_model = self._config.embedding_model
798 embedding_dim = self._config.embedding_dim
799 self._ensure_embedding_compat()
800 self._fts_ready = False
801 self._scalar_ready = False
802 if not records:
803 return 0
804 _check_vector_dims(records, embedding_dim)
805 table = self._chunks_table()
806 table.add(records)
807 if self.get_meta() is None:
808 self._write_meta_unlocked(embedding_model=embedding_model, embedding_dim=embedding_dim)
809 return len(records)
811 def add_chunks(self, records: list[dict]) -> int:
812 """Add chunk records to the store. Returns count added.
814 Raises ``EmbeddingModelMismatchError`` if the persisted ``_meta`` row was
815 written under a different embedding model than the current ``cfg``. On the
816 first write to a fresh store, ``_meta`` is initialized from the current cfg.
818 The gate runs inside the write lock and uses a single cfg snapshot so a
819 concurrent ``set_embedding_model`` cannot slip a write in past a stale
820 compatibility check.
821 """
822 with self._write_lock():
823 return self._add_chunks_unlocked(records)
825 def replace_chunks(self, records: list[dict], predicate: str) -> int:
826 """Replace the chunk rows matching *predicate* with *records* under one write lock.
828 Same compatibility and dimension gates as :meth:`add_chunks`, both run
829 before the delete so a rejected write leaves the existing rows in place.
830 The lock serializes writers; delete and add are separate commits, so a
831 concurrent reader can briefly see the rows absent, and callers retry on
832 a crash between the two. A delete failure propagates, following the same
833 rule as :meth:`_delete_by_sources_unlocked`: swallowed, the caller would
834 read a success-shaped result over rows still describing the old body.
835 Returns the count added.
836 """
837 with self._write_lock():
838 self._ensure_embedding_compat()
839 _check_vector_dims(records, self._config.embedding_dim)
840 table = self._chunks_table()
841 table.delete(predicate)
842 return self._add_chunks_unlocked(records)
844 def _stamp_meta_unlocked(self, embedding_model: str, embedding_dim: int) -> None:
845 """Write the embedder identity on the first write to a fresh store."""
846 if self.get_meta() is None:
847 self._write_meta_unlocked(embedding_model=embedding_model, embedding_dim=embedding_dim)
849 def absorb_rows(self, name: str, rows: pa.Table) -> int:
850 """Append an ingest shard's *rows* to table *name*, creating it if absent.
852 The rows already carry their embeddings, so this is the merge path for a
853 multi-GPU sync: the per-worker stores are folded in whole and the indexes
854 are rebuilt corpus-wide afterwards.
855 """
856 with self._write_lock():
857 self._ensure_embedding_compat()
858 self._fts_ready = False
859 self._scalar_ready = False
860 ensure_table(self.get_db(), name, rows.schema).add(rows)
861 self._stamp_meta_unlocked(self._config.embedding_model, self._config.embedding_dim)
862 return int(rows.num_rows)
864 def _assert_schemas_match(
865 self,
866 name: str,
867 target: Path,
868 shard_tables: list[Path],
869 sources: Sequence[lance.LanceDataset],
870 ) -> None:
871 """Refuse shards whose schema differs from the table's.
873 A mismatched vector width is the realistic case: two workers built with
874 different embedding models. Adoption commits fragment metadata and never
875 passes the rows through a writer, so nothing else would notice.
876 """
877 import lance
879 expected = lance.dataset(str(target)).schema
880 for shard_table, source in zip(shard_tables, sources, strict=True):
881 if not source.schema.equals(expected):
882 raise ValueError(
883 f"Cannot adopt {shard_table} into {name}: its schema does not match "
884 f"the index. A shard built with a different embedding model cannot be "
885 f"folded in; re-ingest it with the model this index was built on."
886 )
888 def adopt_fragments(self, name: str, shard_tables: list[Path]) -> int:
889 """Take over *shard_tables*' data files for table *name* without copying rows.
891 Every fragment's data file is hard-linked into this table's directory and
892 the whole set committed as one metadata-only append, so the rows are never
893 read or rewritten and the bytes exist once on disk with two names. That is
894 the difference between a merge that costs the corpus and one that costs its
895 fragment count: the copy path rewrites every vector, and because the shard
896 stores are kept as resume state the corpus would then be on disk twice.
898 Whole-fragment only, so this serves the first full merge. A re-sync merges
899 named sources, where a fragment holds both touched and untouched rows and
900 the scoped row copy is both correct and already cheap.
902 Returns the rows adopted. Raises ValueError when a shard's schema differs
903 from the table's, and OSError when a shard is on another filesystem (hard
904 links cannot cross one) or a data file name collides. The parent is left
905 as it was in every failure: schemas are checked before anything is linked,
906 links made before a failure are removed, and the commit happens once at
907 the end.
908 """
909 import lance
911 with self._write_lock():
912 self._ensure_embedding_compat()
913 target = self._config.lancedb_dir / f"{name}.lance"
914 sources = [lance.dataset(str(shard_table)) for shard_table in shard_tables]
915 if not sources:
916 return 0
917 if name not in table_names(self.get_db()):
918 ensure_table(self.get_db(), name, sources[0].schema)
919 # Before anything is linked. Committing a fragment whose schema does
920 # not match writes an index that reads back as a panic inside Arrow
921 # rather than an error: the row copy is rejected by the writer, and
922 # adoption has no writer to reject it.
923 self._assert_schemas_match(name, target, shard_tables, sources)
924 # A table created empty has no data directory yet: nothing has been
925 # written into it, and the links below need somewhere to go.
926 (target / "data").mkdir(parents=True, exist_ok=True)
927 adopted: list[lance.FragmentMetadata] = []
928 linked: list[Path] = []
929 rows = 0
930 try:
931 for shard_table, source in zip(shard_tables, sources, strict=True):
932 for fragment in source.get_fragments():
933 meta = fragment.metadata
934 for data_file in meta.files:
935 filename = Path(data_file.path).name
936 link = target / "data" / filename
937 os.link(shard_table / "data" / filename, link)
938 linked.append(link)
939 adopted.append(meta)
940 rows += source.count_rows()
941 except OSError:
942 # A link made before the failure is a file the manifest never
943 # names, so nothing would ever remove it. The caller falls back
944 # to copying rows.
945 for link in linked:
946 link.unlink(missing_ok=True)
947 raise
948 if not adopted:
949 return 0
950 self._fts_ready = False
951 self._scalar_ready = False
952 existing = lance.dataset(str(target))
953 lance.LanceDataset.commit(
954 str(target),
955 lance.LanceOperation.Append(adopted),
956 read_version=existing.version,
957 )
958 self._stamp_meta_unlocked(self._config.embedding_model, self._config.embedding_dim)
959 return rows
961 def bm25_probe(
962 self, query_text: str, top_k: int = 5, chunk_type: ChunkType | None = None
963 ) -> list[SearchChunk]:
964 """Quick BM25-only search for confidence checking. Returns up to top_k results.
966 When *chunk_type* is set, only chunks of that type ("raw" or "wiki") are returned.
967 """
968 table = self.open_table(CHUNKS_TABLE)
969 if table is None:
970 return []
971 if not self._fts_ready:
972 self.ensure_fts_index(blocking=False)
973 if not self._fts_ready:
974 return []
975 try:
976 results = _lexical_rows(table, query_text, top_k, chunk_type)
977 norms = normalized_bm25([r.bm25_score or 0.0 for r in results])
978 return [
979 r.model_copy(update={"score": norm}) for r, norm in zip(results, norms, strict=True)
980 ]
981 except Exception:
982 log.debug("BM25 probe failed", exc_info=True)
983 return []
985 def search(
986 self,
987 query_vector: Vector,
988 top_k: int | None = None,
989 max_distance: float | None = None,
990 query_text: str | None = None,
991 chunk_type: ChunkType | None = None,
992 ) -> list[SearchChunk]:
993 """Search for similar chunks. Hybrid when FTS available, else vector-only.
995 Results with distance > max_distance are filtered out (vector-only path).
996 Pass max_distance=0 to disable filtering.
997 When *chunk_type* is set, only chunks of that type ("raw" or "wiki") are returned.
999 Raises ``EmbeddingModelMismatchError`` if the persisted ``_meta`` row was
1000 written under a different embedding model than the current ``cfg``.
1001 """
1002 if top_k is None:
1003 top_k = self._config.top_k
1004 if max_distance is None:
1005 max_distance = self._config.max_distance
1006 table = self.open_table(CHUNKS_TABLE)
1007 if table is None:
1008 return []
1009 self.initialize_meta_if_legacy()
1010 self.canonicalize_meta_if_legacy()
1011 self._ensure_embedding_compat()
1012 self._warn_stale_doc_prefix()
1014 if not self._scalar_ready:
1015 # A serve-only store never ran ingest, where scalar indexes are
1016 # built; without them the source/chunk_type prefilters full-scan.
1017 self.ensure_scalar_indexes(blocking=False)
1019 if query_text and not self._fts_ready:
1020 self.ensure_fts_index(blocking=False)
1021 if query_text and self._config.title_search and not self._title_fts_ready:
1022 self.ensure_title_fts_index(blocking=False)
1024 if query_text and self._fts_ready:
1025 try:
1026 return self._hybrid_search(
1027 table, query_text, query_vector, top_k, max_distance, chunk_type
1028 )
1029 except Exception:
1030 # Falling back changes recall characteristics for the query;
1031 # a corpus-wide FTS breakage must not present as silence.
1032 log.warning("Hybrid search failed, falling back to vector-only", exc_info=True)
1034 rows = self._vector_arm(
1035 table, query_vector, top_k * self._config.candidate_multiplier, chunk_type
1036 )
1037 log.debug(
1038 "Vector search: query=%r, candidates=%d, max_distance=%.2f",
1039 query_text or "vector-only",
1040 len(rows),
1041 max_distance,
1042 )
1043 if rows:
1044 log.debug("Top 5 distances: %s", [r.distance for r in rows[:5]])
1045 results = self._filter_and_rerank(rows, query_vector, top_k, max_distance)
1046 return [
1047 r.model_copy(
1048 update={"score": vector_similarity(r.distance) if r.distance is not None else 0.0}
1049 )
1050 for r in results
1051 ]
1053 def _vector_arm(
1054 self,
1055 table: lancedb.table.Table,
1056 query_vector: Vector,
1057 limit: int,
1058 chunk_type: ChunkType | None,
1059 ) -> list[SearchChunk]:
1060 """Vector-arm candidates with the ANN recall recovery applied.
1062 When ``chunk_type`` is set, the predicate is pushed into the query so
1063 the limit applies *after* the type filter; post-filtering would
1064 silently starve wiki-only queries whose matches live past the window.
1065 """
1066 query = table.search(query_vector).metric(_VECTOR_METRIC).limit(limit)
1067 if _has_vector_index(table):
1068 # IVF_PQ is lossy; probe more partitions and refine against full
1069 # vectors so recall stays close to the exact flat scan.
1070 query = query.nprobes(_ann_nprobes(table.count_rows()))
1071 query = query.refine_factor(_ANN_REFINE_FACTOR)
1072 if chunk_type:
1073 query = query.where(_chunk_type_predicate(chunk_type))
1074 return [SearchChunk(**r) for r in query.to_list()]
1076 def _fts_arm(
1077 self,
1078 table: lancedb.table.Table,
1079 query_text: str,
1080 limit: int,
1081 chunk_type: ChunkType | None,
1082 ) -> list[SearchChunk]:
1083 """BM25-arm candidates over the chunk text."""
1084 return _lexical_rows(table, query_text, limit, chunk_type)
1086 def _title_arm(
1087 self,
1088 table: lancedb.table.Table,
1089 query_text: str,
1090 limit: int,
1091 chunk_type: ChunkType | None,
1092 ) -> list[SearchChunk]:
1093 """One BM25 row per document whose title matches, in title-relevance order.
1095 Every chunk of a document carries the same title, so all of its chunks
1096 tie on BM25 and a plain ``limit`` would return an arbitrary tie-ordered
1097 subset of a single document. Instead this over-fetches, collapses each
1098 source to one deterministic representative (its first chunk), and returns
1099 the top *limit* documents ordered by title score -- so "a query naming a
1100 document by title surfaces its chunks" holds as one stable row per doc.
1102 Empty when the store predates the title column or its FTS index (old
1103 indexes keep working) and empty on any query-time failure: the optional
1104 title arm must never take down the healthy chunk arm, so its failure
1105 degrades to no-titles, mirroring ``bm25_probe``.
1106 """
1107 if not _has_fts_index(table, _TITLE_COLUMN):
1108 return []
1109 # Every chunk of one document ties on title BM25, so a fixed window can
1110 # fill up with a single long document's chunks and starve every other
1111 # title-matching document. Widen the fetch until enough distinct
1112 # documents surface, the matches run out, or the ceiling is hit.
1113 fetch = max(limit * _TITLE_FETCH_FACTOR, _TITLE_MIN_FETCH)
1114 while True:
1115 try:
1116 rows = _lexical_rows(table, query_text, fetch, chunk_type, column=_TITLE_COLUMN)
1117 except Exception:
1118 log.debug("Title arm search failed; contributing no title rows", exc_info=True)
1119 return []
1120 best: dict[str, SearchChunk] = {}
1121 for row in rows:
1122 seen = best.get(row.source)
1123 if seen is None or row.chunk_index < seen.chunk_index:
1124 best[row.source] = row
1125 if len(best) >= limit or len(rows) < fetch or fetch >= _TITLE_FETCH_CEILING:
1126 break
1127 fetch = min(fetch * 4, _TITLE_FETCH_CEILING)
1128 ordered = sorted(best.values(), key=lambda r: (-(r.bm25_score or 0.0), r.source))
1129 return ordered[:limit]
1131 def _hybrid_search(
1132 self,
1133 table: lancedb.table.Table,
1134 query_text: str,
1135 query_vector: Vector,
1136 top_k: int,
1137 max_distance: float,
1138 chunk_type: ChunkType | None = None,
1139 ) -> list[SearchChunk]:
1140 """Multi-arm retrieval fused by weighted reciprocal rank; the fused ordering is final.
1142 A vector arm and a chunk-BM25 arm always run; a title-BM25 arm joins
1143 when ``cfg.title_search`` is on. Each row's fused score is the
1144 weight-normalized sum of its arm contributions: the vector arm has
1145 weight 1.0, the lexical arm ``cfg.lexical_fusion_weight`` (scaled per
1146 query when ``cfg.adaptive_fusion`` is on), the title arm
1147 ``cfg.title_search_weight``. So a row a single peaked arm is certain
1148 about scores that arm's share of the total weight, not a fixed 0.5.
1150 Each arm fetches exactly ``top_k`` rows. Deeper pools measurably hurt
1151 rank fusion by flooding the fused top-k with both-arm mediocrity and
1152 burying single-arm certainty (lexical identifier hits above all). No
1153 MMR runs here: lexical passages are often mutually similar, which MMR
1154 penalizes, trading relevant hits for off-topic neighbors.
1156 Title rows carry ``bm25_score``, so a title match counts as lexical
1157 support for the distance exemption like any other lexical hit.
1158 """
1159 title_rows: list[SearchChunk] = []
1160 if self._config.title_search:
1161 title_rows = self._title_arm(table, query_text, top_k, chunk_type)
1162 vector_rows = self._vector_arm(table, query_vector, top_k, chunk_type)
1163 base_lexical_weight = self._config.lexical_fusion_weight
1164 base_title_weight = self._config.title_search_weight
1165 lexical_weight = base_lexical_weight
1166 title_weight = base_title_weight
1167 if self._config.adaptive_fusion:
1168 # Quiet the lexical arms per query by vector confidence. The title
1169 # arm is lexical too, so the same factor scales it.
1170 scale = adaptive_weight_scale(vector_rows, self._config.adaptive_fusion_margin)
1171 lexical_weight = base_lexical_weight * scale
1172 title_weight = base_title_weight * scale
1173 fused = fuse_arms(
1174 vector_rows,
1175 self._fts_arm(table, query_text, top_k, chunk_type),
1176 title_rows,
1177 lexical_weight=lexical_weight,
1178 title_weight=title_weight,
1179 )
1180 fused = _drop_unsupported_far_rows(fused, max_distance)
1181 return fused[:top_k]
1183 def _filter_and_rerank(
1184 self,
1185 results: list[SearchChunk],
1186 query_vector: Vector,
1187 top_k: int,
1188 max_distance: float,
1189 ) -> list[SearchChunk]:
1190 """Apply the configured distance filter, then MMR-rerank down to top_k."""
1191 if max_distance > 0:
1192 before = len(results)
1193 if self._config.adaptive_threshold:
1194 results = self._adaptive_filter(results, top_k, max_distance)
1195 filter_name = "adaptive"
1196 else:
1197 results = self._fixed_filter(results, max_distance)
1198 filter_name = "fixed"
1199 log.debug(
1200 "After %s filter: %d/%d results, threshold=%.2f",
1201 filter_name,
1202 len(results),
1203 before,
1204 max_distance,
1205 )
1206 if len(results) > top_k:
1207 results = mmr_rerank(query_vector, results, top_k, self._config.mmr_lambda)
1208 return results
1210 def _adaptive_filter(
1211 self, results: list[SearchChunk], top_k: int, initial_threshold: float
1212 ) -> list[SearchChunk]:
1213 """Widen cosine distance threshold when too few results.
1214 Inspired by grantflow's (grantflow-ai/grantflow) adaptive retrieval
1215 pattern which widens thresholds on recursive retry. Step size and
1216 cap are configurable via ``self._config.adaptive_threshold_step``.
1218 Pre-sorts results by distance for a single-pass cutoff search.
1219 Step size is ``self._config.adaptive_threshold_step`` (default 0.2).
1220 """
1221 cap = max(initial_threshold, _MAX_THRESHOLD)
1222 step = self._config.adaptive_threshold_step
1224 sorted_results = sorted(results, key=_get_distance)
1226 threshold = initial_threshold
1227 for _ in range(_MAX_FILTER_ITERATIONS):
1228 if threshold > cap:
1229 break
1230 cutoff = _count_within_threshold(sorted_results, threshold)
1231 if cutoff >= top_k:
1232 return sorted_results[:cutoff]
1233 threshold += step
1234 # Final pass at cap
1235 cutoff = _count_within_threshold(sorted_results, cap)
1236 return sorted_results[:cutoff]
1238 def _fixed_filter(self, results: list[SearchChunk], threshold: float) -> list[SearchChunk]:
1239 """Simple fixed threshold filter - keep only results within distance threshold."""
1240 return [r for r in results if _get_distance(r) <= threshold]
1242 def add_entities(self, records: list[dict]) -> int:
1243 """Append typed entity rows; creates the table on first write.
1245 Additive to the store: existing tables and schemas are untouched, and
1246 stores without this table behave as if nothing was ever extracted.
1247 """
1248 if not records:
1249 return 0
1250 from lilbee.retrieval.entities.schema import _entities_schema
1252 with self._write_lock():
1253 db = self.get_db()
1254 table = ensure_table(db, ENTITIES_TABLE, _entities_schema())
1255 table.add(records)
1256 return len(records)
1258 def entity_schema_state(self) -> EntitySchemaState | None:
1259 """The persisted entity schema row, or ``None`` when never induced.
1261 The schema is machine state induced from the corpus and lives inside
1262 the index, so it travels with the data.
1263 """
1264 table = self.open_table(ENTITY_SCHEMA_TABLE)
1265 if table is None:
1266 return None
1267 rows = table.search().limit(None).to_list()
1268 if not rows:
1269 return None
1270 # One row by contract; take the newest if a rewrite ever left a stale one.
1271 row = max(rows, key=lambda r: r["updated_at"])
1272 return EntitySchemaState(
1273 schema_json=str(row["schema_json"]),
1274 applied=bool(row["applied"]),
1275 source_count=int(row["source_count"]),
1276 updated_at=str(row["updated_at"]),
1277 )
1279 def save_entity_schema(self, schema_json: str, *, applied: bool, source_count: int) -> None:
1280 """Overwrite the single persisted entity schema row."""
1281 with self._write_lock():
1282 db = self.get_db()
1283 table = ensure_table(db, ENTITY_SCHEMA_TABLE, _entity_schema_state_schema())
1284 _safe_delete_unlocked(table, ENTITY_SCHEMA_DELETE_ALL_PREDICATE)
1285 table.add(
1286 [
1287 {
1288 "schema_json": schema_json,
1289 "applied": applied,
1290 "source_count": source_count,
1291 "updated_at": datetime.now(UTC).isoformat(),
1292 }
1293 ]
1294 )
1296 def mark_entity_schema_applied(self) -> None:
1297 """Record that a full extraction pass completed under the stored schema."""
1298 state = self.entity_schema_state()
1299 if state is None:
1300 return
1301 self.save_entity_schema(
1302 state["schema_json"], applied=True, source_count=state["source_count"]
1303 )
1305 def entity_value_counts(self, entity_type: str) -> tuple[int, int]:
1306 """(mentions, distinct normalized values) for one entity type.
1308 Full scan by design: a count is a corpus property. Streaming batches
1309 keep memory flat at any corpus size.
1310 """
1311 table = self.open_table(ENTITIES_TABLE)
1312 if table is None:
1313 return 0, 0
1314 mentions = 0
1315 values: set[str] = set()
1316 arrow = table.to_arrow().select(["type", "normalized_value"])
1317 for batch in arrow.to_batches(max_chunksize=_TERM_SCAN_BATCH_ROWS):
1318 types = batch.column("type").to_pylist()
1319 vals = batch.column("normalized_value").to_pylist()
1320 for t_, v in zip(types, vals, strict=True):
1321 if t_ == entity_type:
1322 mentions += 1
1323 values.add(v)
1324 return mentions, len(values)
1326 def entity_association_counts(self, counted: str, grouped_by: str) -> dict[str, int]:
1327 """Distinct *counted*-type values co-occurring with each *grouped_by* value.
1329 Co-occurrence is per chunk: two entities extracted from the same
1330 ``(source, chunk_index)`` are associated. This is the GROUP BY that
1331 answers "how many X is each Y associated with".
1332 """
1333 table = self.open_table(ENTITIES_TABLE)
1334 if table is None:
1335 return {}
1336 per_chunk: dict[tuple[str, int], tuple[set[str], set[str]]] = {}
1337 arrow = table.to_arrow().select(["type", "normalized_value", "source", "chunk_index"])
1338 for batch in arrow.to_batches(max_chunksize=_TERM_SCAN_BATCH_ROWS):
1339 rows = zip(
1340 batch.column("type").to_pylist(),
1341 batch.column("normalized_value").to_pylist(),
1342 batch.column("source").to_pylist(),
1343 batch.column("chunk_index").to_pylist(),
1344 strict=True,
1345 )
1346 for t_, v, src, idx in rows:
1347 if t_ not in (counted, grouped_by):
1348 continue
1349 counted_vals, group_vals = per_chunk.setdefault((src, idx), (set(), set()))
1350 (counted_vals if t_ == counted else group_vals).add(v)
1351 associations: dict[str, set[str]] = {}
1352 for counted_vals, group_vals in per_chunk.values():
1353 for group_value in group_vals:
1354 associations.setdefault(group_value, set()).update(counted_vals)
1355 return {k: len(v) for k, v in sorted(associations.items())}
1357 def count_term_mentions(self, term: str) -> tuple[int, int]:
1358 """(matching chunks, distinct matching sources) for a case-insensitive
1359 substring scan of the WHOLE chunks table.
1361 This is deliberately a full scan, not a top-k search: a count is a
1362 corpus property, and any retrieval shortcut undercounts it. Streaming
1363 Arrow batches keeps the working set to one batch of text at a time,
1364 so cost is linear in corpus size and memory stays flat.
1365 """
1366 table = self.open_table(CHUNKS_TABLE)
1367 if table is None:
1368 return 0, 0
1369 needle = term.lower()
1370 chunk_hits = 0
1371 sources: set[str] = set()
1372 arrow = table.to_arrow().select(["source", "chunk"])
1373 for batch in arrow.to_batches(max_chunksize=_TERM_SCAN_BATCH_ROWS):
1374 texts = batch.column("chunk").to_pylist()
1375 names = batch.column("source").to_pylist()
1376 for name, text in zip(names, texts, strict=True):
1377 if text and needle in text.lower():
1378 chunk_hits += 1
1379 sources.add(name)
1380 return chunk_hits, len(sources)
1382 def count_chunks(self) -> int:
1383 """Total chunks in the store."""
1384 table = self.open_table(CHUNKS_TABLE)
1385 return table.count_rows() if table is not None else 0
1387 def get_chunks_by_source(self, source: str) -> list[SearchChunk]:
1388 """Return every chunk whose ``source`` equals *source*.
1390 The database does the filtering, so only the matching rows are read.
1391 A query failure raises rather than falling back to a whole-table scan:
1392 a document's chunks are a bounded read, and the scan that would rescue
1393 it costs the entire index, vectors included, in memory.
1394 """
1395 table = self.open_table(CHUNKS_TABLE)
1396 if table is None:
1397 return []
1398 escaped = escape_sql_string(source)
1399 rows = table.search().where(f"source = '{escaped}'").limit(None).to_list()
1400 return [SearchChunk(**r) for r in rows]
1402 def get_chunks_by_indices(self, source: str, indices: Sequence[int]) -> list[SearchChunk]:
1403 """Return *source*'s chunks whose ``chunk_index`` is in *indices*.
1405 Rows come back in ``chunk_index`` order; indices past either end of
1406 the document are simply absent from the result. Filtering happens in
1407 the database for the same reason as :meth:`get_chunks_by_source`:
1408 neighbor expansion runs once per hit source per query, so a
1409 whole-table rescue would spike memory on the hottest path there is.
1410 """
1411 if not indices:
1412 return []
1413 table = self.open_table(CHUNKS_TABLE)
1414 if table is None:
1415 return []
1416 escaped = escape_sql_string(source)
1417 wanted = ", ".join(str(int(i)) for i in indices)
1418 predicate = f"source = '{escaped}' AND chunk_index IN ({wanted})"
1419 rows = table.search().where(predicate).limit(None).to_list()
1420 return sorted((SearchChunk(**r) for r in rows), key=lambda c: c.chunk_index)
1422 def _delete_by_sources_unlocked(self, sources: list[str]) -> None:
1423 """Delete the sources' chunks, page texts, and chunk-concept rows.
1425 Caller must hold ``write_lock()``. One ``IN`` delete per table covers
1426 every source, so a batched flush pays a constant number of predicate
1427 deletes instead of one set per document. A delete failure propagates:
1428 swallowed, it would leave every flushed file silently stale; raised,
1429 the flush fails and the files replan on the next sync.
1430 """
1431 quoted = ", ".join(f"'{escape_sql_string(source)}'" for source in sources)
1432 for name, column in _PER_SOURCE_TABLES:
1433 table = self.open_table(name)
1434 if table is not None:
1435 table.delete(f"{column} IN ({quoted})")
1437 def _delete_by_source_unlocked(self, source: str) -> None:
1438 """Delete a single source's chunks, page texts, and chunk-concept rows."""
1439 self._delete_by_sources_unlocked([source])
1441 def delete_by_source(self, source: str) -> None:
1442 """Delete a source's chunks and page texts."""
1443 with self._write_lock():
1444 self._delete_by_source_unlocked(source)
1445 self._invalidate_source_cache()
1447 def add_page_texts(self, records: list[dict]) -> int:
1448 """Add per-page text rows (no vectors). Returns count added."""
1449 if not records:
1450 return 0
1451 with self._write_lock():
1452 db = self.get_db()
1453 table = ensure_table(db, PAGE_TEXTS_TABLE, _page_texts_schema())
1454 table.add(records)
1455 return len(records)
1457 def get_page_texts(self, source: str | None = None) -> list[PageTextRecord]:
1458 """Return per-page text rows, all or for a single *source*."""
1459 table = self.open_table(PAGE_TEXTS_TABLE)
1460 if table is None:
1461 return []
1462 query = table.search()
1463 if source is not None:
1464 query = query.where(f"source = '{escape_sql_string(source)}'")
1465 rows: list[PageTextRecord] = query.limit(None).to_list()
1466 return rows
1468 def page_texts_arrow(self, source: str | None = None) -> pa.Table:
1469 """Return per-page text rows as an Arrow table in a single scan.
1471 The columnar sibling of :meth:`get_page_texts`: the export path keeps the
1472 whole set in Arrow (no per-row Python objects) from read through file
1473 write. Empty with the canonical schema when the table or *source* is empty.
1474 """
1475 table = self.open_table(PAGE_TEXTS_TABLE)
1476 if table is None:
1477 return _page_texts_schema().empty_table()
1478 query = table.search().select(["source", "page", "text", "content_type"])
1479 if source is not None:
1480 query = query.where(f"source = '{escape_sql_string(source)}'")
1481 return query.limit(None).to_arrow()
1483 def sources_arrow(self) -> pa.Table:
1484 """Return each tracked source's extraction metadata as an Arrow table.
1486 The columnar sibling of :meth:`get_sources`, keyed by ``source`` so it
1487 joins straight onto the page-text table. ``get_sources`` builds a dict per
1488 source, which on a corpus of millions of single-page documents costs more
1489 than the text being exported. An index written before the metadata columns
1490 existed gets them as nulls rather than missing, so the join has one shape.
1491 """
1492 import pyarrow as pa
1493 import pyarrow.compute as pc
1495 table = self.open_table(SOURCES_TABLE)
1496 columns = ["source", *SourceMeta._fields]
1497 if table is None:
1498 return pa.schema([pa.field(name, pa.utf8()) for name in columns]).empty_table()
1499 present = [name for name in SourceMeta._fields if name in table.schema.names]
1500 arrow = table.search().select(["filename", *present]).limit(None).to_arrow()
1501 arrow = arrow.rename_columns(["source", *present])
1502 for name in SourceMeta._fields:
1503 if name not in present:
1504 arrow = arrow.append_column(name, pa.nulls(arrow.num_rows, pa.utf8()))
1505 arrow = arrow.select(columns)
1506 if pc.count_distinct(arrow.column("source")).as_py() == arrow.num_rows:
1507 return arrow
1508 # One row per source, so a caller joining on it cannot fan out. A doubled
1509 # row (a source re-merged from a shard) would otherwise multiply every page
1510 # it owns. Last wins, matching the dict this replaced; single-threaded
1511 # because that is the only execution mode with an ordered aggregate.
1512 grouped = arrow.group_by("source", use_threads=False).aggregate(
1513 [(name, "last") for name in SourceMeta._fields]
1514 )
1515 renamed = grouped.rename_columns(
1516 [name.removesuffix("_last") for name in grouped.schema.names]
1517 )
1518 return renamed.select(columns)
1520 def wiki_chunk_sources(self) -> set[str]:
1521 """Return the distinct sources of the chunk rows written by the wiki layer."""
1522 table = self.open_table(CHUNKS_TABLE)
1523 if table is None:
1524 return set()
1525 rows = (
1526 table.search()
1527 .where(f"chunk_type = '{ChunkType.WIKI}'")
1528 .select(["source"])
1529 .limit(None)
1530 .to_list()
1531 )
1532 return {row["source"] for row in rows}
1534 def wiki_citation_sources(self) -> set[str]:
1535 """Return the distinct wiki_source values present in the citations table."""
1536 table = self.open_table(CITATIONS_TABLE)
1537 if table is None:
1538 return set()
1539 rows = table.search().select(["wiki_source"]).limit(None).to_list()
1540 return {row["wiki_source"] for row in rows}
1542 def get_sources(
1543 self,
1544 *,
1545 search: str | None = None,
1546 limit: int | None = None,
1547 offset: int = 0,
1548 ) -> list[SourceRecord]:
1549 """Return source records, filtered by *search* and sliced by offset/limit."""
1550 table = self.open_table(SOURCES_TABLE)
1551 if table is None:
1552 return []
1553 query = table.search()
1554 where = _sources_search_filter(search, include_title="title" in table.schema.names)
1555 if where is not None:
1556 query = query.where(where)
1557 if offset:
1558 query = query.offset(offset)
1559 query = query.limit(limit)
1560 result: list[SourceRecord] = query.to_list() # type: ignore[assignment]
1561 return result
1563 def count_sources(self, *, search: str | None = None) -> int:
1564 """Count tracked sources matching *search* without materializing rows."""
1565 table = self.open_table(SOURCES_TABLE)
1566 if table is None:
1567 return 0
1568 where = _sources_search_filter(search, include_title="title" in table.schema.names)
1569 count: int = table.count_rows() if where is None else table.count_rows(filter=where)
1570 return count
1572 def _source_row(
1573 self,
1574 filename: str,
1575 file_hash: str,
1576 chunk_count: int,
1577 source_type: str,
1578 stat: SourceStat | None,
1579 meta: SourceMeta | None = None,
1580 ) -> dict:
1581 """Build one ``_sources`` row, defaulting absent stat to the unknown sentinel.
1583 Absent extraction metadata persists as NULL, matching rows written
1584 before the metadata columns existed.
1585 """
1586 meta = meta or SourceMeta()
1587 return {
1588 "filename": filename,
1589 "file_hash": file_hash,
1590 "ingested_at": datetime.now(UTC).isoformat(),
1591 "chunk_count": chunk_count,
1592 "source_type": source_type,
1593 "size_bytes": stat.size_bytes if stat else SOURCE_STAT_UNKNOWN,
1594 "mtime_ns": stat.mtime_ns if stat else SOURCE_STAT_UNKNOWN,
1595 "stat_captured_ns": stat.captured_ns if stat else SOURCE_STAT_UNKNOWN,
1596 "title": meta.title or None,
1597 "authors": meta.authors or None,
1598 "created_at": meta.created_at or None,
1599 }
1601 def _sources_table(self) -> lancedb.table.Table:
1602 """Open/create ``_sources``, adding the stat and metadata columns to older tables."""
1603 table = ensure_table(self.get_db(), SOURCES_TABLE, _sources_schema())
1604 defaults = {name: f"CAST({SOURCE_STAT_UNKNOWN} AS BIGINT)" for name in _SOURCE_STAT_COLUMNS}
1605 defaults |= {name: "CAST(NULL AS STRING)" for name in _SOURCE_META_COLUMNS}
1606 missing = {name: sql for name, sql in defaults.items() if name not in table.schema.names}
1607 if missing:
1608 table.add_columns(missing)
1609 return table
1611 def _replace_source_rows_unlocked(self, rows: list[dict]) -> None:
1612 """Replace source rows: one batched delete plus one batched add.
1614 Caller must hold ``write_lock()``. A per-file delete+add pair costs two
1615 LanceDB version commits, so bulk ingest folds every file in a flush into
1616 a single pair.
1617 """
1618 table = self._sources_table()
1619 filenames = ", ".join(f"'{escape_sql_string(r['filename'])}'" for r in rows)
1620 # Skip the add when the delete failed: adding over a stale row would leave
1621 # two _sources rows for one filename. The file replans on the next sync.
1622 if not _safe_delete_unlocked(table, f"filename IN ({filenames})"):
1623 return
1624 table.add(rows)
1626 def upsert_source(
1627 self,
1628 filename: str,
1629 file_hash: str,
1630 chunk_count: int,
1631 source_type: SourceType = SourceType.DOCUMENT,
1632 stat: SourceStat | None = None,
1633 meta: SourceMeta | None = None,
1634 ) -> None:
1635 """Add or update a source tracking record."""
1636 row = self._source_row(filename, file_hash, chunk_count, source_type, stat, meta)
1637 with self._write_lock():
1638 self._replace_source_rows_unlocked([row])
1639 self._invalidate_source_cache()
1641 def update_source_stats(self, backfills: list[SourceStatBackfill]) -> None:
1642 """Record size/mtime for already-tracked sources in batched locked writes."""
1643 if not backfills:
1644 return
1645 for start in range(0, len(backfills), _SOURCE_STAT_BATCH_ROWS):
1646 rows = [
1647 {
1648 **bf.record,
1649 "size_bytes": bf.stat.size_bytes,
1650 "mtime_ns": bf.stat.mtime_ns,
1651 "stat_captured_ns": bf.stat.captured_ns,
1652 }
1653 for bf in backfills[start : start + _SOURCE_STAT_BATCH_ROWS]
1654 ]
1655 with self._write_lock():
1656 self._replace_source_rows_unlocked(rows)
1657 self._invalidate_source_cache()
1659 def optimize_sources(self) -> None:
1660 """Compact the sources table; per-flush upserts otherwise accrete tiny versions."""
1661 with self._write_lock():
1662 table = self.open_table(SOURCES_TABLE)
1663 if table is None:
1664 return
1665 try:
1666 table.optimize()
1667 except Exception:
1668 log.debug("Sources table optimize failed", exc_info=True)
1670 def write_chunks_batch(self, items: list[ChunkWrite]) -> int:
1671 """Write several documents' chunks in one locked transaction. Returns chunks added.
1673 One ``write_lock`` acquisition covers the batch's cleanup deletes, page
1674 texts, chunk add, and source upserts, so a reader never observes a
1675 half-applied batch. Page texts land after the cleanup and before the
1676 source rows, so a page-text failure leaves the rows stale and the files
1677 replan next sync; a document with no chunks still persists its page
1678 texts and source row. The embedding-identity gate and per-vector
1679 dimension check mirror ``add_chunks``; a dimension mismatch raises and
1680 the whole batch is rejected.
1681 """
1682 if not items:
1683 return 0
1684 with self._write_lock(timeout=BATCH_LOCK_TIMEOUT):
1685 embedding_model = self._config.embedding_model
1686 embedding_dim = self._config.embedding_dim
1687 self._ensure_embedding_compat()
1688 self._fts_ready = False
1689 self._scalar_ready = False
1690 all_records = [rec for it in items for rec in it.records]
1691 _check_vector_dims(all_records, embedding_dim)
1692 db = self.get_db()
1693 self._cleanup_batch_unlocked(items)
1694 self._add_page_texts_unlocked(db, items)
1695 self._add_chunk_records_unlocked(all_records, embedding_model, embedding_dim)
1696 self._replace_source_rows_unlocked(self._batch_source_rows(items))
1697 self._invalidate_source_cache()
1698 return len(all_records)
1700 def _cleanup_batch_unlocked(self, items: list[ChunkWrite]) -> None:
1701 """One ``IN`` delete per table for the flagged documents. Caller holds ``write_lock()``."""
1702 cleanup_sources = [it.source for it in items if it.needs_cleanup]
1703 if cleanup_sources:
1704 self._delete_by_sources_unlocked(cleanup_sources)
1706 def _add_page_texts_unlocked(self, db: lancedb.DBConnection, items: list[ChunkWrite]) -> None:
1707 """Add the batch's page-text rows. Caller holds ``write_lock()``."""
1708 page_rows = [row for it in items for row in (it.page_texts or [])]
1709 if page_rows:
1710 ensure_table(db, PAGE_TEXTS_TABLE, _page_texts_schema()).add(page_rows)
1712 def _add_chunk_records_unlocked(
1713 self,
1714 all_records: list[dict],
1715 embedding_model: str,
1716 embedding_dim: int,
1717 ) -> None:
1718 """Add the batch's chunk rows, writing meta on first use. Caller holds ``write_lock()``."""
1719 if not all_records:
1720 return
1721 self._chunks_table().add(all_records)
1722 if self.get_meta() is None:
1723 self._write_meta_unlocked(embedding_model=embedding_model, embedding_dim=embedding_dim)
1725 def _batch_source_rows(self, items: list[ChunkWrite]) -> list[dict]:
1726 """One ``_sources`` row per batched document."""
1727 return [
1728 self._source_row(
1729 it.source, it.file_hash, len(it.records), it.source_type, it.stat, it.meta
1730 )
1731 for it in items
1732 ]
1734 def _delete_source_unlocked(self, filename: str) -> None:
1735 """Remove the *filename* source record. Caller must hold ``write_lock()``."""
1736 table = self.open_table(SOURCES_TABLE)
1737 if table is not None:
1738 _safe_delete_unlocked(table, f"filename = '{escape_sql_string(filename)}'")
1740 def delete_source(self, filename: str) -> None:
1741 """Remove a source file tracking record."""
1742 with self._write_lock():
1743 self._delete_source_unlocked(filename)
1744 self._invalidate_source_cache()
1746 def _remove_many_unlocked(self, names: list[str]) -> None:
1747 """Delete the documents' chunks and source records together.
1749 All deletes run under the caller's single ``write_lock()`` so no
1750 reader can observe chunks whose source record is already gone; one
1751 ``IN`` delete per table covers the whole set.
1752 """
1753 self._delete_by_sources_unlocked(names)
1754 quoted = ", ".join(f"'{escape_sql_string(name)}'" for name in names)
1755 table = self.open_table(SOURCES_TABLE)
1756 if table is not None:
1757 _safe_delete_unlocked(table, f"filename IN ({quoted})")
1759 def relocate_sources(self, moves: list[tuple[str, str, SourceStat | None]]) -> None:
1760 """Re-key moved sources from old filename to new, preserving their chunks.
1762 A source whose file moved (same content hash, new path) keeps its chunks
1763 and embeddings; only its filename key and disk stat change. Each per-source
1764 table's source column, the citation source_filename, and the sources row are
1765 updated in place under one write lock, so a move costs no re-extraction or
1766 re-embedding. ``moves`` is ``(old_name, new_name, new_stat)`` tuples.
1768 Each table is opened once; the re-key is then a targeted per-move update.
1769 A single-statement batch would need a ``CASE`` expression, which LanceDB's
1770 update SQL does not support, and a delete+re-add across the vector tables is
1771 not worth its risk for what is a rare mass relabel.
1772 """
1773 if not moves:
1774 return
1775 from lilbee.data.title import derive_title # circular at module scope
1777 with self._write_lock():
1778 tables = [(self.open_table(name), column) for name, column in _RELOCATABLE_TABLES]
1779 sources = self.open_table(SOURCES_TABLE)
1780 for old, new, stat in moves:
1781 where_old = f"= '{escape_sql_string(old)}'"
1782 new_title = self._relocated_title(sources, old, new, derive_title)
1783 for table, column in tables:
1784 if table is None:
1785 continue
1786 values: dict[str, object] = {column: new}
1787 # Stem titles track the filename; re-derive them on the same
1788 # handle and statement as the re-key.
1789 if new_title is not _KEEP_TITLE and _TITLE_COLUMN in table.schema.names:
1790 values[_TITLE_COLUMN] = new_title
1791 table.update(where=f"{column} {where_old}", values=values)
1792 if sources is not None:
1793 row_values: dict[str, object] = {"filename": new}
1794 if new_title is not _KEEP_TITLE:
1795 row_values["title"] = new_title
1796 if stat is not None:
1797 row_values["size_bytes"] = stat.size_bytes
1798 row_values["mtime_ns"] = stat.mtime_ns
1799 row_values["stat_captured_ns"] = stat.captured_ns
1800 sources.update(where=f"filename {where_old}", values=row_values)
1801 self._invalidate_source_cache()
1803 def _relocated_title(
1804 self,
1805 sources: lancedb.table.Table | None,
1806 old: str,
1807 new: str,
1808 derive: Callable[[str], str],
1809 ) -> str | None:
1810 """New title for a moved source, or ``_KEEP_TITLE`` when it must not change.
1812 Extraction-derived titles survive a move (the content is unchanged);
1813 a stem-derived title tracks the filename it was derived from, so it is
1814 re-derived from the new name instead of matching the old one forever.
1815 """
1816 if sources is None:
1817 return _KEEP_TITLE
1818 try:
1819 rows = (
1820 sources.search()
1821 .where(f"filename = '{escape_sql_string(old)}'")
1822 .select(["title"])
1823 .limit(1)
1824 .to_list()
1825 )
1826 except Exception:
1827 return _KEEP_TITLE
1828 if not rows or "title" not in rows[0]:
1829 return _KEEP_TITLE
1830 stored = rows[0]["title"] or ""
1831 if stored != (derive(old) or ""):
1832 return _KEEP_TITLE
1833 return derive(new) or None
1835 def remove_documents(self, names: list[str]) -> RemoveResult:
1836 """Remove documents from the knowledge base by source name.
1838 Looks up known sources and deletes their chunks and source records. Never
1839 touches files on disk: source bytes are the user's, and a linked-in corpus
1840 must never be deleted. Durable, file-aware removal (skip-markers, unlinking
1841 a top-level link) lives in :func:`lilbee.app.ingest.remove_documents_durably`.
1843 Returns a RemoveResult with removed and not_found lists.
1844 """
1845 known = {s["filename"] for s in self.get_sources()}
1846 removed = [name for name in names if name in known]
1847 not_found = [name for name in names if name not in known]
1849 if removed:
1850 # One lock acquisition and one IN-delete per table for the whole
1851 # set, mirroring the batched flush path, instead of a LanceDB
1852 # version commit per document.
1853 with self._write_lock():
1854 self._remove_many_unlocked(removed)
1855 self._invalidate_source_cache()
1857 return RemoveResult(removed=removed, not_found=not_found)
1859 def clear_table(self, name: str, predicate: str) -> bool:
1860 """Delete rows matching *predicate* from *name*. Acquires write lock.
1862 Returns whether the delete succeeded, so a caller recording the
1863 outcome (prune, the legacy migration) does not report success over a
1864 swallowed failure.
1865 """
1866 with self._write_lock():
1867 table = self.open_table(name)
1868 if table is None:
1869 return True
1870 return _safe_delete_unlocked(table, predicate)
1872 def clear_and_add(self, name: str, schema: pa.Schema, rows: list[dict], predicate: str) -> None:
1873 """Replace the rows matching *predicate* with *rows* in one locked write.
1875 Delete and add run under a single write lock, so a reader never observes
1876 the table emptied mid-rebuild. A delete failure propagates, following the
1877 same rule as :meth:`_delete_by_sources_unlocked`: adding over rows whose
1878 predecessors are still there duplicates them, and reporting success would
1879 leave the caller acting on state it thinks it replaced.
1880 """
1881 with self._write_lock():
1882 db = self.get_db()
1883 table = ensure_table(db, name, schema)
1884 table.delete(predicate)
1885 if rows:
1886 table.add(rows)
1888 def add_citations(self, records: list[CitationRecord]) -> int:
1889 """Add citation records to the store. Returns count added."""
1890 if not records:
1891 return 0
1892 with self._write_lock():
1893 db = self.get_db()
1894 table = ensure_table(db, CITATIONS_TABLE, _citations_schema())
1895 table.add(records)
1896 return len(records)
1898 def get_citations_for_wiki(self, wiki_source: str) -> list[CitationRecord]:
1899 """Get all citations for a wiki page."""
1900 table = self.open_table(CITATIONS_TABLE)
1901 if table is None:
1902 return []
1903 escaped = escape_sql_string(wiki_source)
1904 rows: list[CitationRecord] = table.search().where(f"wiki_source = '{escaped}'").to_list()
1905 return rows
1907 def get_citations_for_source(self, source_filename: str) -> list[CitationRecord]:
1908 """Get all citations that reference a source document (reverse lookup)."""
1909 table = self.open_table(CITATIONS_TABLE)
1910 if table is None:
1911 return []
1912 escaped = escape_sql_string(source_filename)
1913 rows: list[CitationRecord] = (
1914 table.search().where(f"source_filename = '{escaped}'").to_list()
1915 )
1916 return rows
1918 def delete_citations_for_wiki(self, wiki_source: str) -> bool:
1919 """Delete all citations for a wiki page. Returns whether the delete succeeded."""
1920 return self.clear_table(CITATIONS_TABLE, _citations_for_wiki_predicate(wiki_source))
1922 def delete_all_wiki_rows(self) -> bool:
1923 """Delete every wiki chunk row, every citation, and every mention.
1924 Returns whether all deletes succeeded, so a caller cannot report a wipe
1925 over a swallowed failure. Only the wiki layer writes citations and
1926 mentions, so wiping it empties those tables outright. All deletes run
1927 before the results combine.
1928 """
1929 chunks_cleared = self.clear_table(CHUNKS_TABLE, f"chunk_type = '{ChunkType.WIKI}'")
1930 citations_cleared = self.clear_table(CITATIONS_TABLE, "1 = 1")
1931 mentions_cleared = self.clear_wiki_mentions()
1932 return chunks_cleared and citations_cleared and mentions_cleared
1934 def replace_citations_for_wiki(self, wiki_source: str, records: list[CitationRecord]) -> None:
1935 """Swap a wiki page's citations for *records* under one write lock.
1937 The lock keeps another writer from landing between the delete and the
1938 add; the two remain separate commits, so a crash in between leaves the
1939 rows absent until the page is regenerated or accepted again.
1940 """
1941 self.clear_and_add(
1942 CITATIONS_TABLE,
1943 _citations_schema(),
1944 [dict(rec) for rec in records],
1945 _citations_for_wiki_predicate(wiki_source),
1946 )
1948 def replace_wiki_mentions_for_source(self, source: str, rows: list[dict]) -> None:
1949 """Swap one source's wiki mention rows for *rows* under one write lock.
1951 A source contributes its whole mention set at once, so an incremental
1952 wiki refresh replaces exactly that source's rows and leaves every other
1953 source's evidence in place for the corpus-wide aggregate.
1954 """
1955 escaped = escape_sql_string(source)
1956 self.clear_and_add(
1957 WIKI_MENTIONS_TABLE,
1958 _wiki_mentions_schema(),
1959 rows,
1960 f"source = '{escaped}'",
1961 )
1963 def wiki_mention_rows(self, slugs: Iterable[str] | None = None) -> list[dict]:
1964 """Every wiki mention row, or only those for *slugs*.
1966 The wiki aggregates these across sources to rebuild the stub index. A
1967 full refresh reads them all; an incremental one reads only the slugs it
1968 touched, to recompute their corpus-wide totals without a full scan.
1969 """
1970 table = self.open_table(WIKI_MENTIONS_TABLE)
1971 if table is None:
1972 return []
1973 query = table.search()
1974 if slugs is not None:
1975 wanted = list(slugs)
1976 if not wanted:
1977 return []
1978 joined = ", ".join(f"'{escape_sql_string(s)}'" for s in wanted)
1979 query = query.where(f"slug IN ({joined})")
1980 rows: list[dict] = query.limit(None).to_list()
1981 return rows
1983 def clear_wiki_mentions(self) -> bool:
1984 """Drop every wiki mention row (a full rebuild starts from nothing)."""
1985 return self.clear_table(WIKI_MENTIONS_TABLE, "1 = 1")
1987 def has_wiki_mentions(self) -> bool:
1988 """Whether any wiki mention row exists.
1990 A cold store or one migrated from the file-only index has none, which
1991 forces a refresh to rebuild in full and seed the table before an
1992 incremental pass can aggregate over it.
1993 """
1994 table = self.open_table(WIKI_MENTIONS_TABLE)
1995 return table is not None and table.count_rows() > 0
1997 def _memories_schema(self) -> pa.Schema:
1998 return pa.schema(
1999 [
2000 pa.field("id", pa.utf8()),
2001 pa.field("owner", pa.utf8()),
2002 pa.field("shared", pa.bool_()),
2003 pa.field("kind", pa.utf8()),
2004 pa.field("source", pa.utf8()),
2005 pa.field("text", pa.utf8()),
2006 pa.field("vector", pa.list_(pa.float32(), self._config.embedding_dim)),
2007 pa.field("created_at", pa.utf8()),
2008 pa.field("updated_at", pa.utf8()),
2009 ]
2010 )
2012 def _duplicate_memory_id_unlocked(
2013 self, table: lancedb.table.Table, record: MemoryRow
2014 ) -> str | None:
2015 """Return the id of a near-duplicate same-owner, same-kind memory, if any."""
2016 if table.count_rows() == 0:
2017 return None
2018 predicate = (
2019 f"owner = '{escape_sql_string(record.owner)}' "
2020 f"AND kind = '{escape_sql_string(record.kind)}'"
2021 )
2022 rows = table.search(record.vector).metric("cosine").where(predicate).limit(1).to_list()
2023 if rows and rows[0].get("_distance", 1.0) <= self._config.memory_dedup_distance:
2024 return str(rows[0]["id"])
2025 return None
2027 def _evict_overflow_unlocked(self, table: lancedb.table.Table, owner: str) -> None:
2028 """Delete oldest memories for *owner* so an incoming insert stays within the cap."""
2029 cap = self._config.memory_max_per_owner
2030 predicate = f"owner = '{escape_sql_string(owner)}'"
2031 rows = table.search().where(predicate).limit(None).to_list()
2032 if len(rows) < cap:
2033 return
2034 rows.sort(key=lambda r: r.get("created_at", ""))
2035 for row in rows[: len(rows) - (cap - 1)]:
2036 _safe_delete_unlocked(table, f"id = '{escape_sql_string(str(row['id']))}'")
2038 def add_memory(self, record: MemoryRow) -> str:
2039 """Insert *record*, or update the nearest same-owner duplicate in place.
2041 Returns the stored id. Raises ``EmbeddingModelMismatchError`` when the store
2042 was built under a different embedding model, and ``ValueError`` on a vector
2043 dimension mismatch.
2044 """
2045 if len(record.vector) != self._config.embedding_dim:
2046 raise ValueError(
2047 f"Memory vector dimension mismatch: expected "
2048 f"{self._config.embedding_dim}, got {len(record.vector)}"
2049 )
2050 with self._write_lock():
2051 embedding_model = self._config.embedding_model
2052 embedding_dim = self._config.embedding_dim
2053 self._ensure_embedding_compat()
2054 db = self.get_db()
2055 table = ensure_table(db, MEMORIES_TABLE, self._memories_schema())
2056 duplicate_id = self._duplicate_memory_id_unlocked(table, record)
2057 if duplicate_id is not None and _safe_delete_unlocked(
2058 table, f"id = '{escape_sql_string(duplicate_id)}'"
2059 ):
2060 # Only reuse the id once the old row is actually gone; a swallowed
2061 # delete failure would otherwise leave two rows with the same id.
2062 record.id = duplicate_id
2063 self._evict_overflow_unlocked(table, record.owner)
2064 table.add([record.model_dump(mode="json")])
2065 if self.get_meta() is None:
2066 self._write_meta_unlocked(
2067 embedding_model=embedding_model, embedding_dim=embedding_dim
2068 )
2069 return record.id
2071 def get_memories(
2072 self,
2073 *,
2074 owner_predicate: str,
2075 kind: MemoryKind | None = None,
2076 ) -> list[MemoryRow]:
2077 """Return memories matching *owner_predicate* and optional *kind*, newest first."""
2078 table = self.open_table(MEMORIES_TABLE)
2079 if table is None:
2080 return []
2081 clauses = [f"({owner_predicate})"]
2082 if kind is not None:
2083 clauses.append(f"kind = '{escape_sql_string(kind)}'")
2084 rows = table.search().where(" AND ".join(clauses)).limit(None).to_list()
2085 memories = [MemoryRow(**r) for r in rows]
2086 memories.sort(key=lambda m: m.created_at, reverse=True)
2087 return memories
2089 def search_memories(
2090 self,
2091 query_vector: Vector,
2092 *,
2093 owner_predicate: str,
2094 top_k: int,
2095 max_distance: float,
2096 ) -> list[MemoryRow]:
2097 """Vector-recall FACT memories within *max_distance*, best first."""
2098 table = self.open_table(MEMORIES_TABLE)
2099 if table is None or top_k <= 0:
2100 return []
2101 self._ensure_embedding_compat()
2102 predicate = f"({owner_predicate}) AND kind = '{MemoryKind.FACT}'"
2103 rows = table.search(query_vector).metric("cosine").where(predicate).limit(top_k).to_list()
2104 return [MemoryRow(**r) for r in rows if r.get("_distance", 1.0) <= max_distance]
2106 def update_memory(self, memory_id: str, *, shared: bool, owner: str) -> bool:
2107 """Set the *shared* flag on *owner*'s memory. Returns True when found and owned.
2109 The ``owner`` predicate scopes the mutation to the caller's namespace so an
2110 agent cannot flip another owner's (or the human's) memory.
2111 """
2112 with self._write_lock():
2113 table = self.open_table(MEMORIES_TABLE)
2114 if table is None:
2115 return False
2116 predicate = self._owned_memory_predicate(memory_id, owner)
2117 rows = table.search().where(predicate).limit(1).to_list()
2118 if not rows:
2119 return False
2120 record = MemoryRow(**rows[0])
2121 record.shared = shared
2122 record.updated_at = datetime.now(UTC).isoformat()
2123 # If the delete fails, do not add the modified copy: that would leave
2124 # two rows for one id. Report not-updated instead.
2125 if not _safe_delete_unlocked(table, predicate):
2126 return False
2127 table.add([record.model_dump(mode="json")])
2128 return True
2130 def delete_memory(self, memory_id: str, *, owner: str) -> bool:
2131 """Delete *owner*'s memory by id. Returns True when a matching row was deleted.
2133 The ``owner`` predicate scopes the delete to the caller's namespace so an
2134 agent cannot destroy another owner's (or the human's) memory.
2135 """
2136 with self._write_lock():
2137 table = self.open_table(MEMORIES_TABLE)
2138 if table is None:
2139 return False
2140 predicate = self._owned_memory_predicate(memory_id, owner)
2141 if not table.search().where(predicate).limit(1).to_list():
2142 return False
2143 # Report the real outcome: a swallowed delete failure must not be
2144 # reported as a successful forget.
2145 return _safe_delete_unlocked(table, predicate)
2147 @staticmethod
2148 def _owned_memory_predicate(memory_id: str, owner: str) -> str:
2149 """SQL predicate matching a single memory id within *owner*'s namespace."""
2150 return f"id = '{escape_sql_string(memory_id)}' AND owner = '{escape_sql_string(owner)}'"
2152 def rebuild_memory_embeddings(self, embed: Callable[[list[str]], list[Vector]]) -> int:
2153 """Re-embed every memory under the current model, recreating the table.
2155 The vector column dimension is immutable, so a different-dim model needs a
2156 fresh table; recreating unconditionally also covers the same-dim case. Memory
2157 text is human-authored and re-embeddable, so no data is lost. Returns the count.
2159 The snapshot, embed, and table rebuild all run under the write lock so a
2160 concurrent ``add_memory`` cannot commit into the read-then-drop window and
2161 be erased; it either lands before the snapshot or blocks until the rebuild
2162 finishes.
2163 """
2164 with self._write_lock():
2165 table = self.open_table(MEMORIES_TABLE)
2166 if table is None:
2167 return 0
2168 rows = table.search().limit(None).to_list()
2169 if not rows:
2170 return 0
2171 memories = [MemoryRow(**r) for r in rows]
2172 vectors = embed([m.text for m in memories])
2173 for memory, vector in zip(memories, vectors, strict=True):
2174 # MemoryRow serializes to JSON on write, which has no ndarray encoding.
2175 memory.vector = vector.tolist()
2176 db = self.get_db()
2177 db.drop_table(MEMORIES_TABLE)
2178 new_table = ensure_table(db, MEMORIES_TABLE, self._memories_schema())
2179 new_table.add([m.model_dump(mode="json") for m in memories])
2180 return len(memories)
2182 def close(self) -> None:
2183 """Release the database connection and reset state."""
2184 self._db = None
2185 self._fts_ready = False
2186 self._title_fts_ready = False
2187 self._scalar_ready = False
2189 def drop_all(self) -> None:
2190 """Drop every table except ``_memories`` -- used by rebuild.
2192 Memory is user-authored data with no on-disk source, not derived from
2193 documents, so a rebuild preserves it. Only a factory reset (which deletes
2194 the data directory) clears it.
2195 """
2196 with self._write_lock():
2197 self._fts_ready = False
2198 self._title_fts_ready = False
2199 self._scalar_ready = False
2200 db = self.get_db()
2201 for name in table_names(db):
2202 if name == MEMORIES_TABLE:
2203 continue
2204 db.drop_table(name)
2205 self._invalidate_source_cache()