Coverage for src/lilbee/data/store/core.py: 100%
1947 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""The ``Store`` class: high-level LanceDB read/write API used across lilbee."""
3from __future__ import annotations
5import hashlib
6import json
7import logging
8import math
9import threading
10import time
11import uuid
12from collections import Counter
13from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
14from concurrent.futures import ThreadPoolExecutor
15from contextlib import AbstractContextManager, ExitStack, contextmanager, suppress
16from datetime import UTC, datetime
17from pathlib import Path
18from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
20import pyarrow as pa
21import pyarrow.compute as pc
23from lilbee.core import settings
24from lilbee.core.config import (
25 CHUNK_CONCEPTS_TABLE,
26 CHUNKS_TABLE,
27 CITATIONS_TABLE,
28 CRAWL_PAGES_TABLE,
29 ENTITIES_TABLE,
30 ENTITY_SCHEMA_TABLE,
31 MEMORIES_TABLE,
32 META_TABLE,
33 PAGE_TEXTS_TABLE,
34 SOURCES_TABLE,
35 WIKI_MENTIONS_TABLE,
36 Config,
37)
38from lilbee.core.health_warnings import HealthWarning, WarningCode
39from lilbee.core.security import write_private_text
40from lilbee.core.vectors import Vector
41from lilbee.core.web_names import is_web_source
42from lilbee.retrieval.embedding_profiles import resolve_embedding_profile
43from lilbee.runtime.lock import (
44 LOCK_TIMEOUT,
45 LockingUnsupportedError,
46 LockTimeoutError,
47 write_lock,
48)
49from lilbee.runtime.progress import (
50 UPGRADE_HEADLINE,
51 DetailedProgressCallback,
52 EventType,
53 LibraryUpgradeEvent,
54 TableWork,
55 UpgradePhase,
56 noop_callback,
57)
59from .crawl_pages import CrawlPages, import_legacy_crawl_files, legacy_crawl_files
60from .fusion import adaptive_weight_scale, fuse_arms, normalized_bm25, vector_similarity
61from .held_import import legacy_files, read_legacy_records, remove_legacy_files
62from .lance_helpers import (
63 _CHUNK_COLUMN,
64 _chunk_type_predicate,
65 _has_fts_index,
66 _has_scalar_index,
67 _has_vector_index,
68 _index_registry,
69 _listed_indices,
70 _safe_delete_unlocked,
71 _scalar_index_dangling,
72 _vector_index_dangling,
73 ensure_table,
74 escape_sql_string,
75 open_committed,
76 refs_compatible,
77 table_names,
78)
79from .locations import SourceLocations, root_location, under_sql
80from .ranking import mmr_rerank
81from .schema import (
82 _citations_schema,
83 _entity_schema_state_schema,
84 _meta_schema,
85 _page_texts_schema,
86 _sources_schema,
87 _wiki_mentions_schema,
88)
89from .types import (
90 ENTITY_SCHEMA_DELETE_ALL_PREDICATE,
91 META_DELETE_ALL_PREDICATE,
92 META_SCHEMA_VERSION,
93 READ_CONSISTENCY_INTERVAL,
94 SOURCE_STAT_UNKNOWN,
95 ChunkType,
96 ChunkWrite,
97 CitationRecord,
98 EmbeddingModelMismatchError,
99 EntitySchemaState,
100 HeldOut,
101 HeldRecord,
102 MemoryKind,
103 MemoryRow,
104 PageTextRecord,
105 RemoveResult,
106 SearchChunk,
107 SourceHead,
108 SourceMeta,
109 SourceMove,
110 SourceRecord,
111 SourceRelocation,
112 SourcesCheckpoint,
113 SourceStat,
114 SourceStatBackfill,
115 SourceSummary,
116 SourceType,
117 StoreMeta,
118)
120if TYPE_CHECKING:
121 from lancedb.db import LanceDBConnection
122 from lancedb.expr import Expr
123 from lancedb.index import FTS, IndexConfig
124 from lancedb.query import LanceFtsQueryBuilder, LanceQueryBuilder
125 from lancedb.table import LanceTable
127log = logging.getLogger(__name__)
129# Batched ingest flushes contend with long store operations (a search-triggered
130# FTS optimize can hold the lock past the interactive 30s), and failing the
131# flush replans and re-embeds the whole batch. Give the batch path more
132# patience than interactive writes before it gives up.
133BATCH_LOCK_TIMEOUT = 120.0
134# Lock budget for index builds reached from the read path: a query must not
135# stall behind a long ingest, so it skips the build and retries next search.
136_READ_LOCK_TIMEOUT = 2.0
137# What makes a read-path index build skip rather than fail the query.
138_UPKEEP_SKIPPED = (LockTimeoutError, LockingUnsupportedError)
141def _drop_unsupported_far_rows(
142 results: list[SearchChunk], max_distance: float
143) -> list[SearchChunk]:
144 """Apply ``max_distance`` to rows whose only signal is the vector arm.
146 A row the BM25 arm also matched keeps lexical support regardless of its
147 vector distance; dropping it on distance alone would re-bury exactly the
148 identifier hits rank fusion exists to preserve.
149 """
150 if max_distance <= 0:
151 return results
152 return [
153 r
154 for r in results
155 if r.bm25_score is not None or r.distance is None or r.distance <= max_distance
156 ]
159_MAX_THRESHOLD = 1.0
160_MAX_FILTER_ITERATIONS = 20 # safety cap to prevent runaway loops
163def _is_fts_position_overflow(exc: Exception) -> bool:
164 """True when *exc* is LanceDB's positional-FTS list-encoding overflow.
166 A positional index (built by an intermediate dev commit) raises e.g.
167 "Max offset N exceeds length of values M" on optimize(); a positionless
168 rebuild is the remediation. Matched on message because LanceDB raises it
169 as a generic error type.
170 """
171 msg = str(exc).lower()
172 return "offset" in msg and "exceeds" in msg
175def _stable_top(
176 table: LanceTable,
177 query: LanceQueryBuilder,
178 limit: int,
179 score: str,
180 *,
181 ascending: bool = False,
182) -> list[dict[str, Any]]:
183 """The best *limit* rows of *query*, one set in one order on every call.
185 LanceDB returns equal scores in an arbitrary order that changes between calls,
186 and its limit cuts them at an arbitrary row. Rank fusion turns that order into
187 different scores and so into different results. A tie that crosses the cut is
188 widened until it ends, and ties then rank by source and chunk.
190 The vector and text columns cost a wide fetch 50 times its time. Widening so
191 reads only the score and the ranking key, and whole rows only for the winners.
193 A tie past ``_TIE_FETCH_CEILING`` rows keeps the stable order, but the rows
194 inside the tie are the ones LanceDB returned.
195 """
197 def tie_crosses_cut(found: list[dict[str, Any]], width: int) -> bool:
198 return (
199 len(found) == width
200 and width < _TIE_FETCH_CEILING
201 and found[limit - 1][score] == found[-1][score]
202 )
204 def best(found: list[dict[str, Any]]) -> list[dict[str, Any]]:
205 sign = 1 if ascending else -1
206 ranked = sorted(found, key=lambda r: (sign * r[score], r[_SOURCE_ID], r[_CHUNK_INDEX]))
207 return ranked[:limit]
209 width = limit + 1
210 rows = query.limit(width).to_list()
211 if not tie_crosses_cut(rows, width):
212 return best(rows)
213 query = query.select([_SOURCE_ID, _CHUNK_INDEX, score])
214 while tie_crosses_cut(rows, width):
215 width = min(width * _TIE_WIDEN_FACTOR, _TIE_FETCH_CEILING)
216 rows = query.limit(width).to_list()
217 return _whole_rows(table, best(rows), score)
220def _whole_rows(table: LanceTable, keys: list[dict[str, Any]], score: str) -> list[dict[str, Any]]:
221 """The whole row of each of *keys*, in that order, with its *score*.
223 The read is by source and chunk rather than by row id: a row id changes when a
224 compaction lands between the fetch of the keys and this read. A row deleted in
225 that time drops out.
226 """
227 if not keys:
228 return []
229 indexes: dict[str, list[int]] = {}
230 for key in keys:
231 indexes.setdefault(key[_SOURCE_ID], []).append(key[_CHUNK_INDEX])
232 where = " OR ".join(
233 f"({_SOURCE_ID} = '{escape_sql_string(source)}'"
234 f" AND {_CHUNK_INDEX} IN ({', '.join(map(str, chunks))}))"
235 for source, chunks in indexes.items()
236 )
237 found = {
238 (row[_SOURCE_ID], row[_CHUNK_INDEX]): row
239 for row in table.search().where(where).limit(None).to_list()
240 }
241 return [
242 {**found[place], score: key[score]}
243 for key in keys
244 if (place := (key[_SOURCE_ID], key[_CHUNK_INDEX])) in found
245 ]
248def _lexical_rows(
249 table: LanceTable,
250 query_text: str,
251 limit: int,
252 chunk_type: ChunkType | None,
253 column: str = _CHUNK_COLUMN,
254) -> list[dict[str, Any]]:
255 """BM25 rows for *query_text* over a single FTS *column*.
257 ``MatchQuery`` pins the column and matches plain terms, so an unpinned search
258 cannot widen to the title index and a quoted span cannot reach LanceDB as a
259 phrase (which the positionless index rejects). This is the one place FTS
260 queries are built; every arm goes through it.
261 """
262 from lancedb.query import MatchQuery
264 # lancedb's stubs omit FullTextQuery from the fts overload and read MatchQuery's
265 # kw-only defaults as required; the call is the documented FTS query form.
266 query: LanceFtsQueryBuilder = table.search( # type: ignore[call-overload]
267 MatchQuery(query_text, column), # type: ignore[call-arg]
268 query_type="fts",
269 )
270 if chunk_type:
271 query = query.where(_chunk_type_predicate(chunk_type))
272 return _stable_top(table, query, limit, _SCORE_COLUMN)
275# Vector ANN index. IVF_PQ compresses vectors so search scales to millions;
276# refine_factor re-ranks the PQ candidates against full vectors to recover recall.
277# The index type is carried by the lancedb IvfPq config at build time.
278_VECTOR_METRIC: Final = "cosine"
279# The scalar index kinds lilbee builds on the columns its queries filter by.
280ScalarIndexType = Literal["BTREE", "BITMAP"]
281_ANN_NPROBES_FLOOR = 20
282# Fraction of IVF partitions probed per query. 0.05 was the "fast" end of the
283# recall/latency curve and measurably cost recall at scale: on the 8.8M-passage
284# MS MARCO index it gave recall@100 of 67%, where probing more partitions
285# reached 87% (docs/benchmarks/retrieval-msmarco.md). 0.15 is the "balanced /
286# high-recall" range for IVF and recovers most of that gap; refine_factor below
287# still re-ranks the survivors against full vectors.
288_ANN_NPROBES_PARTITION_FRACTION = 0.15
289_ANN_REFINE_FACTOR = 10
291# Stat columns of ``_sources``; mirrors the field names in ``schema._sources_schema``
292# and ``types.SourceRecord``. Legacy tables that predate these columns are migrated
293# in place with the SOURCE_STAT_UNKNOWN sentinel.
294_SOURCE_STAT_COLUMNS = ("size_bytes", "mtime_ns", "stat_captured_ns")
296# Extraction-metadata columns of ``_sources``; nullable strings, so legacy
297# tables migrate in place with NULL (meaning "extractor reported nothing").
298_SOURCE_META_COLUMNS = ("title", "authors", "created_at")
300# Document-title column of the chunks table; nullable so pre-title rows and
301# writers that carry no title (wiki pages) read as NULL.
302_TITLE_COLUMN = "title"
304# The column that names a source in every table; ``_sources`` alone maps it to a path.
305_SOURCE_ID = "source_id"
306_CHUNK_INDEX = "chunk_index"
307# The columns LanceDB adds to a keyword row and to a vector row.
308_SCORE_COLUMN = "_score"
309_DISTANCE_COLUMN = "_distance"
310# The name of each thread an upgrade adds a table's id column on.
311_UPGRADE_THREAD_PREFIX = "library-upgrade"
312# The ``_sources`` column that holds a source's path.
313_PATH = "path"
314# The ``_sources`` column that names why a sync holds a row's path out (``HeldOut``).
315_HELD_OUT = "held_out"
316# The content hash a sync holds the path out at, and the reason shown for it.
317_HELD_HASH = "held_hash"
318_HELD_REASON = "held_reason"
319_HELD_COLUMNS = (_HELD_OUT, _HELD_HASH, _HELD_REASON)
321# Tables that hold a source's rows, deleted with the source. The concept
322# nodes/edges tables carry no source column (corpus-level aggregates).
323_PER_SOURCE_TABLES = (
324 CHUNKS_TABLE,
325 PAGE_TEXTS_TABLE,
326 CHUNK_CONCEPTS_TABLE,
327 ENTITIES_TABLE,
328 WIKI_MENTIONS_TABLE,
329)
331# Each table that named a source by its path before ``source_id``, and that column.
332_PATH_KEY_COLUMNS = (
333 (CHUNKS_TABLE, "source"),
334 (PAGE_TEXTS_TABLE, "source"),
335 (CHUNK_CONCEPTS_TABLE, "chunk_source"),
336 (ENTITIES_TABLE, "source"),
337 (CITATIONS_TABLE, "source_filename"),
338 (WIKI_MENTIONS_TABLE, "source"),
339)
340_PATH_KEY_BY_TABLE = dict(_PATH_KEY_COLUMNS)
341# The tables an upgrade visits, in its order: ``_sources`` first.
342_UPGRADE_ORDER = (SOURCES_TABLE, *_PATH_KEY_BY_TABLE)
343# The name an upgrade shows each table by.
344_UPGRADE_LABELS = {
345 SOURCES_TABLE: "source list",
346 CHUNKS_TABLE: "search chunks",
347 PAGE_TEXTS_TABLE: "page text",
348 CHUNK_CONCEPTS_TABLE: "concept links",
349 ENTITIES_TABLE: "entities",
350 CITATIONS_TABLE: "citations",
351 WIKI_MENTIONS_TABLE: "wiki mentions",
352}
353# The path column of ``_sources`` before it was renamed.
354_OLD_PATH = "filename"
355# The ``_meta`` rows an upgrade has yet to stamp; a store stamped 1 keeps 1.
356_UNSTAMPED = f"schema_version >= 2 AND schema_version < {META_SCHEMA_VERSION}"
358# The type column of ``_sources``, and of the chunks table.
359_SOURCE_TYPE = "source_type"
360_CHUNK_TYPE = "chunk_type"
361# The source rows a listing shows: each that holds a document, failed at a later edit or not.
362_LISTED = (
363 f"{_SOURCE_TYPE} != '{SourceType.WIKI}' AND ({_HELD_OUT} IS NULL"
364 f" OR {_HELD_OUT} = '{HeldOut.FAILED}' AND file_hash IS NOT NULL)"
365)
366# The source rows stored at their location; every other source is stored by its key.
367_DOCUMENT = f"{_SOURCE_TYPE} = '{SourceType.DOCUMENT}'"
368# The source rows of wiki pages, which share no row with a source of another type.
369_WIKI = f"{_SOURCE_TYPE} = '{SourceType.WIKI}'"
370# The chunks of wiki pages; an upgrade gives them the id of a page, never of a document.
371_WIKI_CHUNK = f"{_CHUNK_TYPE} = '{ChunkType.WIKI}'"
372# Every other chunk, whose key names a document.
373_NOT_WIKI_CHUNK = f"({_CHUNK_TYPE} IS NULL OR NOT ({_WIKI_CHUNK}))"
374# What a wiki page's key is joined with before its md5 gives the page its id at an upgrade.
375_WIKI_ID_SUFFIX = "wiki"
376_WIKI_ID_SQL = "md5(concat({key}, chr(0), '" + _WIKI_ID_SUFFIX + "'))"
377# The source rows a move found by a sync carries: the listed ones and every failure.
378_MOVABLE = (
379 f"{_SOURCE_TYPE} != '{SourceType.WIKI}'"
380 f" AND ({_HELD_OUT} IS NULL OR {_HELD_OUT} = '{HeldOut.FAILED}')"
381)
382# The source rows an upgrade compares a document key with: any row whose type is not wiki.
383_NOT_WIKI_SOURCE = f"({_SOURCE_TYPE} IS NULL OR NOT ({_WIKI}))"
385# Sentinel: relocation must leave the stored title untouched (extraction-derived).
386_KEEP_TITLE = "\x00keep"
388# The key columns a batched flush deletes, looks up or merges by in each table it writes.
389_FLUSH_KEY_COLUMNS = (
390 (CHUNKS_TABLE, _SOURCE_ID),
391 (PAGE_TEXTS_TABLE, _SOURCE_ID),
392 (SOURCES_TABLE, _PATH),
393 (SOURCES_TABLE, _SOURCE_ID),
394)
396# A filtered delete scans every row its key index does not cover. The index is
397# rebuilt once this many flushes of rows sit past it.
398_KEY_INDEX_STALE_FLUSHES = 32
400# Source rows per locked write or per ``IN`` list: the first sync after a
401# stat-column upgrade backfills every source, and an unchunked replace would
402# join millions of paths into one predicate.
403_SOURCE_STAT_BATCH_ROWS = 2000
405# One hold of the write lock moves at most this many sources, so a writer that
406# waits is not kept out for the length of a whole plan batch.
407_RELOCATE_HOLD_MOVES = 2000
409_NULL_TEXT_SQL = "CAST(NULL AS STRING)"
411# What a dropped source row a citation names keeps: its path, and no record a sync reads.
412_CITED_ONLY_SQL = {
413 _HELD_OUT: f"'{HeldOut.REMOVED}'",
414 "chunk_count": "0",
415 _HELD_HASH: _NULL_TEXT_SQL,
416 _HELD_REASON: _NULL_TEXT_SQL,
417}
419# Rows a location fill reads and writes back in one batch: bounds its memory at any row count.
420_FILL_BATCH_ROWS = 8192
421# The address LanceDB gives a row in one version of its table.
422_ROW_ID = "_rowid"
423# The scratch table a fill pairs each row already at a new place with the row moving onto it in.
424_FILL_PAIRS_TABLE = "_fill_pairs"
425_LANDING_ID = "landing_id"
426_MOVED_ID = "moved_id"
427_FILL_PAIRS_SCHEMA = pa.schema(
428 [pa.field(_PATH, pa.utf8()), pa.field(_LANDING_ID, pa.utf8()), pa.field(_MOVED_ID, pa.utf8())]
429)
430# Ids one citation update re-points: its lookup list is built again for every row it reads.
431_CITATION_LOOKUP_IDS = 128
433# Beside the tables: the checkpoint up to which the wiki's page index covers the source rows.
434_WIKI_REFRESHED_NAME = ".wiki_refreshed"
435_WIKI_REFRESHED_UNREADABLE = "Cannot read %s (%s): the wiki index reads every source again"
437# Rows per Arrow batch when a filtered read walks the whole source table.
438_SOURCE_SCAN_BATCH_ROWS = 20_000
440_HEAD_COLUMNS = [_OLD_PATH, "file_hash", "chunk_count", "ingested_at"]
441_SUMMARY_COLUMNS = [*_HEAD_COLUMNS, "source_type"]
444def _source_heads(rows: pa.RecordBatch | pa.Table) -> list[SourceHead]:
445 """The listing columns of *rows* as dicts."""
446 return [
447 SourceHead(
448 filename=row[_OLD_PATH],
449 file_hash=row["file_hash"],
450 chunk_count=row["chunk_count"],
451 ingested_at=row["ingested_at"],
452 )
453 for row in rows.select(_HEAD_COLUMNS).to_pylist()
454 ]
457# Rows per Arrow batch when the aggregate scan walks the whole chunks table;
458# bounds the decoded-text working set while the scan stays columnar.
459_TERM_SCAN_BATCH_ROWS = 20_000
461# The title arm collapses each matched document to one row, so it over-fetches
462# to gather enough distinct documents before deduping. Bounded so a title that
463# hits a huge document can't scan the whole corpus.
464_TITLE_FETCH_FACTOR = 20
465_TITLE_MIN_FETCH = 200
466_TITLE_FETCH_CEILING = 4096
468# Bounds the fetch that settles a tie across a search limit, so a term in every
469# chunk of a huge corpus cannot make one query read all of them.
470_TIE_FETCH_CEILING = 4096
471# A widening reads keys only and costs 11 to 22 ms at any width up to the ceiling, so
472# few wide steps beat many narrow ones: a full tie takes three steps instead of nine.
473_TIE_WIDEN_FACTOR = 8
476def _ann_nprobes(row_count: int) -> int:
477 """Partitions to probe: a fixed fraction of the IVF partition count (~sqrt(N)), floored."""
478 partitions = math.isqrt(max(row_count, 0))
479 return max(_ANN_NPROBES_FLOOR, math.ceil(partitions * _ANN_NPROBES_PARTITION_FRACTION))
482def _check_vector_dims(records: list[dict], embedding_dim: int) -> None:
483 """Raise ``ValueError`` when any record's vector is not *embedding_dim* wide."""
484 for rec in records:
485 vec = rec.get("vector", [])
486 if len(vec) != embedding_dim:
487 raise ValueError(
488 f"Vector dimension mismatch: expected {embedding_dim}, "
489 f"got {len(vec)} (source={rec.get('source', '?')})"
490 )
493def _citations_for_wiki_predicate(wiki_source: str) -> str:
494 """SQL predicate selecting every citation row belonging to *wiki_source*."""
495 return f"wiki_source = '{escape_sql_string(wiki_source)}'"
498def _swap_head_sql(column: str, old_head: str, new_head: str) -> str:
499 """The SQL value of *column* with *new_head* in place of the characters of *old_head*."""
500 return f"concat('{escape_sql_string(new_head)}', substr({column}, {len(old_head) + 1}))"
503def _in_list(names: Iterable[str]) -> str:
504 """The body of a SQL ``IN`` list that holds *names*."""
505 return ", ".join(f"'{escape_sql_string(name)}'" for name in names)
508def _claimed_names(moves: Sequence[SourceMove], unclaimed: set[str]) -> list[str | None]:
509 """The old name each of *moves* takes out of *unclaimed*: its first candidate still there.
511 Moves that share one candidates tuple share one cursor into it, so files
512 of one content cost one pass over their old names and not one for each file.
513 """
514 cursors: dict[int, int] = {}
515 claimed: list[str | None] = []
516 for move in moves:
517 pool = move.candidates
518 at = cursors.get(id(pool), 0)
519 while at < len(pool) and pool[at] not in unclaimed:
520 at += 1
521 cursors[id(pool)] = at
522 if at == len(pool):
523 claimed.append(None)
524 continue
525 unclaimed.discard(pool[at])
526 claimed.append(pool[at])
527 return claimed
530def _located_moves(
531 moves: Sequence[SourceMove], locations: SourceLocations
532) -> tuple[list[SourceMove], dict[str, str]]:
533 """*moves* with every name as its location, and the name of each location.
535 Moves that share one candidates tuple still share one, which the claim reads.
536 """
537 names: dict[str, str] = {}
539 def _at(name: str) -> str:
540 names[location := locations.location(name)] = name
541 return location
543 pools = {id(move.candidates): tuple(map(_at, move.candidates)) for move in moves}
544 located = [SourceMove(pools[id(move.candidates)], _at(move.new), move.stat) for move in moves]
545 return located, names
548def _title_after_move(
549 row: Mapping[str, Any], old: str, new: str, derive: Callable[[str], str]
550) -> str | None:
551 """The title the source *row* takes when *old* becomes *new*, or ``_KEEP_TITLE``.
553 A title that came from the old name follows the new one; any other title stays,
554 and so does a title the new name derives again, so that no row is written for it.
555 """
556 stored = row[_TITLE_COLUMN] or ""
557 if stored != (derive(old) or "") or stored == (derive(new) or ""):
558 return _KEEP_TITLE
559 return derive(new) or None
562def _moved_source_row(row: dict[str, Any], move: SourceMove, title: str | None) -> dict[str, Any]:
563 """The source *row* under *move*'s new name, with *title* and the stat of the file."""
564 moved = {**row, _PATH: move.new}
565 if title is not _KEEP_TITLE:
566 moved[_TITLE_COLUMN] = title
567 if move.stat is not None:
568 moved["size_bytes"] = move.stat.size_bytes
569 moved["mtime_ns"] = move.stat.mtime_ns
570 moved["stat_captured_ns"] = move.stat.captured_ns
571 return moved
574def _member_rows(rows: Iterable[dict[str, Any]], moves: Mapping[str, str]) -> list[dict[str, Any]]:
575 """The *rows* below a moved archive, each under the archive's new name."""
576 members = _members_by_archive(moves, (row[_PATH] for row in rows))
577 paths = {
578 path: moves[archive] + path[len(archive) :]
579 for archive, found in members.items()
580 for path in found
581 }
582 return [{**row, _PATH: paths[row[_PATH]]} for row in rows if row[_PATH] in paths]
585def _new_source_id() -> str:
586 """An id for a source no row holds: random, so a new file at a vacated path takes none."""
587 return uuid.uuid4().hex
590def _keyed(rows: Iterable[Mapping[str, Any]], key: str, ids: Mapping[str, str]) -> list[dict]:
591 """*rows* with the id of the path under *key* in its place."""
592 return [
593 {**{name: value for name, value in row.items() if name != key}, _SOURCE_ID: ids[row[key]]}
594 for row in rows
595 ]
598def _pathed(rows: Iterable[Mapping[str, Any]], key: str, paths: Mapping[str, str]) -> list[dict]:
599 """*rows* with the path of their source id under *key*; a row of no source row drops."""
600 return [
601 {
602 **{name: value for name, value in row.items() if name != _SOURCE_ID},
603 key: paths[row[_SOURCE_ID]],
604 }
605 for row in rows
606 if row[_SOURCE_ID] in paths
607 ]
610def _id_chunks(rows: list[dict[str, Any]]) -> list[SearchChunk]:
611 """Search results of chunk *rows* that name their source by id until the search ends."""
612 return [
613 SearchChunk(
614 **{name: value for name, value in row.items() if name != _SOURCE_ID},
615 source=row[_SOURCE_ID],
616 )
617 for row in rows
618 ]
621def _source_record(row: Mapping[str, Any], locations: SourceLocations) -> SourceRecord:
622 """The ``SourceRecord`` a ``_sources`` row shows: its shown name as ``filename``, no id."""
623 shown = {name: value for name, value in row.items() if name not in (_SOURCE_ID, *_HELD_COLUMNS)}
624 shown[_OLD_PATH] = locations.name(shown.pop(_PATH))
625 return cast("SourceRecord", shown)
628class _SourceWrite(NamedTuple):
629 """One source row a write puts in ``_sources``, by the name and type of its source.
631 With *fields* the row holds them; without, a found row keeps its own and a new
632 row only names the source. A *record* holds the source out of a sync.
633 """
635 name: str
636 source_type: SourceType
637 fields: Mapping[str, Any] | None = None
638 record: HeldRecord | None = None
639 new_id: str | None = None
642class _TableStep(NamedTuple):
643 """What an upgrade does to the table in progress, and how many tables of all are done."""
645 work: TableWork
646 table: str
647 rows: int
648 tables_done: int
649 tables_total: int
652def _documents(names: Iterable[str]) -> Iterator[tuple[str, SourceType]]:
653 """Each of *names* as the name of a document."""
654 return ((name, SourceType.DOCUMENT) for name in names)
657def _source_type_of(chunk_type: str) -> SourceType:
658 """The type of the source a chunk of *chunk_type* is of: a wiki page, else a document."""
659 return SourceType.WIKI if chunk_type == ChunkType.WIKI else SourceType.DOCUMENT
662# The order a lookup by a name with no type tries the types in.
663_READ_ORDER = (SourceType.WIKI, SourceType.IMPORTED, SourceType.DOCUMENT)
666def _is_wiki(source_type: str) -> bool:
667 """Whether *source_type* is a wiki page's, the one type no other source shares a row with."""
668 return source_type == SourceType.WIKI
671def _write_paths(write: _SourceWrite, locations: SourceLocations) -> tuple[str, ...]:
672 """Where the row *write* finds can be: its key, then, for a non-wiki source, its other form."""
673 if _is_wiki(write.source_type):
674 return (write.name,)
675 return tuple(
676 dict.fromkeys(
677 (
678 locations.key(write.name, write.source_type),
679 write.name,
680 locations.location(write.name),
681 )
682 )
683 )
686def _found_rows(
687 sources: LanceTable, writes: Sequence[_SourceWrite], locations: SourceLocations
688) -> dict[str, dict[str, Any]]:
689 """The row of its type class each of *writes* finds, by name; the key is looked at first."""
690 paths = sorted({path for write in writes for path in _write_paths(write, locations)})
691 held: dict[tuple[str, bool], dict[str, Any]] = {}
692 for start in range(0, len(paths), _SOURCE_STAT_BATCH_ROWS):
693 where = f"{_PATH} IN ({_in_list(paths[start : start + _SOURCE_STAT_BATCH_ROWS])})"
694 for row in sources.search().where(where).limit(None).to_list():
695 held[row[_PATH], _is_wiki(row[_SOURCE_TYPE])] = row
696 found: dict[str, dict[str, Any]] = {}
697 for write in writes:
698 wiki = _is_wiki(write.source_type)
699 rows = [held[path, wiki] for path in _write_paths(write, locations) if (path, wiki) in held]
700 if rows:
701 found[write.name] = rows[0]
702 return found
705def _stat_write(backfill: SourceStatBackfill) -> _SourceWrite:
706 """The write that gives a tracked source its verified stat; a record names no held-out one."""
707 record, stat = backfill
708 fields = {name: value for name, value in record.items() if name != _OLD_PATH}
709 fields |= {
710 "size_bytes": stat.size_bytes,
711 "mtime_ns": stat.mtime_ns,
712 "stat_captured_ns": stat.captured_ns,
713 }
714 return _SourceWrite(record["filename"], SourceType(record["source_type"]), fields)
717def _at_or_below(path: str, roots: set[str]) -> bool:
718 """Whether *path* is one of *roots* or lies below one, on whole path segments."""
719 if path in roots:
720 return True
721 cut = path.find("/")
722 while cut != -1:
723 if path[:cut] in roots:
724 return True
725 cut = path.find("/", cut + 1)
726 return False
729def _members_by_archive(names: Iterable[str], filenames: Iterable[str]) -> dict[str, list[str]]:
730 """The *filenames* under ``name/`` for each of *names*, in one pass over *filenames*."""
731 wanted = set(names)
732 members: dict[str, list[str]] = {}
733 for filename in filenames:
734 cut = filename.find("/")
735 while cut != -1:
736 if filename[:cut] in wanted:
737 members.setdefault(filename[:cut], []).append(filename)
738 cut = filename.find("/", cut + 1)
739 return members
742def _held_sources(sources: LanceTable, names: Iterable[str]) -> set[str]:
743 """The *names* that hold a row a move carries in the source table."""
744 wanted = sorted(names)
745 held: set[str] = set()
746 for start in range(0, len(wanted), _SOURCE_STAT_BATCH_ROWS):
747 quoted = _in_list(wanted[start : start + _SOURCE_STAT_BATCH_ROWS])
748 query = sources.search().where(f"{_PATH} IN ({quoted}) AND {_MOVABLE}").select([_PATH])
749 held.update(row[_PATH] for row in query.limit(None).to_list())
750 return held
753def _source_rows(sources: LanceTable, names: list[str]) -> dict[str, dict[str, Any]]:
754 """The row a move carries that the source table holds for each of *names*, by path."""
755 if not names:
756 return {}
757 query = sources.search().where(f"{_PATH} IN ({_in_list(names)}) AND {_MOVABLE}")
758 return {row[_PATH]: row for row in query.limit(None).to_list()}
761def _listed_matching(search: str | None, locations: SourceLocations) -> str:
762 """The listed source rows whose shown name or title holds *search*."""
763 if not search:
764 return _LISTED
765 return f"{_LISTED} AND {locations.search_filter(search, _PATH, _TITLE_COLUMN)}"
768def _md5(text: str) -> str:
769 """The id the migration gives the source at *text*, as LanceDB's ``md5`` computes it."""
770 return hashlib.md5(text.encode("utf-8"), usedforsecurity=False).hexdigest()
773def _wiki_id(key: str) -> str:
774 """The id the migration gives the wiki page at *key*, as ``_WIKI_ID_SQL`` computes it."""
775 return _md5(f"{key}\0{_WIKI_ID_SUFFIX}")
778def _holds_document(row: Mapping[str, Any]) -> bool:
779 """Whether a document source *row* holds an indexed document, as ``_LISTED`` reads it."""
780 kind = row[_HELD_OUT]
781 return kind is None or (kind == HeldOut.FAILED and row["file_hash"] is not None)
784def _held_row(row: dict[str, Any], record: HeldRecord, *, holds_document: bool) -> dict[str, Any]:
785 """*row* with *record*; a failure keeps the document the row holds, a removal holds none."""
786 held = {
787 **row,
788 _HELD_OUT: record.kind.value,
789 _HELD_HASH: record.file_hash,
790 _HELD_REASON: record.reason,
791 }
792 if record.kind is HeldOut.FAILED and holds_document:
793 return held
794 return {**held, "file_hash": None, "chunk_count": 0}
797def _keyed_by_path(db: LanceDBConnection) -> bool:
798 """Whether a table of *db* still names a source by path, or ``_sources`` lacks a column."""
799 sources = open_committed(db, SOURCES_TABLE)
800 if sources is not None:
801 added = (_SOURCE_ID, _PATH, *_SOURCE_STAT_COLUMNS, *_SOURCE_META_COLUMNS, _HELD_OUT)
802 if not set(sources.schema.names).issuperset(added):
803 return True
804 for name, column in _PATH_KEY_COLUMNS:
805 table = open_committed(db, name)
806 if table is not None and column in table.schema.names:
807 return True
808 return False
811def _upgrade_pending(db: LanceDBConnection) -> bool:
812 """Whether *db* is keyed by path, or ``_meta`` still waits for the upgrade's stamp."""
813 if _keyed_by_path(db):
814 return True
815 meta = open_committed(db, META_TABLE)
816 return meta is not None and meta.count_rows(_UNSTAMPED) > 0
819def _upgrade_steps(db: LanceDBConnection, names: list[str]) -> list[tuple[str, LanceTable]]:
820 """The tables of *db* an upgrade visits: ``_sources``, then each still keyed by path."""
821 tables = [(name, db.open_table(name)) for name in _UPGRADE_ORDER if name in names]
822 return [
823 (name, table)
824 for name, table in tables
825 if name == SOURCES_TABLE or _PATH_KEY_BY_TABLE[name] in table.schema.names
826 ]
829def _held_unmigrated(db: LanceDBConnection, data_root: Path) -> bool:
830 """Whether ``_sources`` lacks a held-out record column, or a record file is at *data_root*."""
831 if legacy_files(data_root):
832 return True
833 sources = open_committed(db, SOURCES_TABLE)
834 return sources is not None and not {_HELD_HASH, _HELD_REASON} <= set(sources.schema.names)
837def _add_held_columns(db: LanceDBConnection) -> None:
838 """Give an existing ``_sources`` the held-out record columns it lacks. Lock held."""
839 sources = open_committed(db, SOURCES_TABLE)
840 if sources is None:
841 return
842 missing = [name for name in (_HELD_HASH, _HELD_REASON) if name not in sources.schema.names]
843 if missing:
844 sources.add_columns(dict.fromkeys(missing, _NULL_TEXT_SQL))
847def _index_names(table: LanceTable) -> dict[str, list[str]]:
848 """The columns of each index of *table*, by index name; none when the listing fails."""
849 return {index.name: index.columns for index in _listed_indices(table) or []}
852def _key_column(table: LanceTable, column: str, where: str | None) -> pa.ChunkedArray:
853 """*column* of each row of *table* that *where* selects, or of every row, read as Arrow."""
854 query = table.search() if where is None else table.search().where(where)
855 return query.select([column]).limit(None).to_arrow().column(column)
858def _keys_missing_from(
859 sources: LanceTable, table: LanceTable, column: str
860) -> list[tuple[str, SourceType]]:
861 """Each distinct key in *column* of *table* that no ``_sources`` row of its class holds.
863 A wiki chunk names a wiki page; every other key names a document. Each class
864 reads only its own keys and its own source paths, and Arrow compares them,
865 so no row becomes a Python object and one key column is held at a time.
866 """
867 typed = table.name == CHUNKS_TABLE
868 classes = [(SourceType.DOCUMENT, _NOT_WIKI_CHUNK if typed else None, _NOT_WIKI_SOURCE)]
869 if typed:
870 classes.insert(0, (SourceType.WIKI, _WIKI_CHUNK, _WIKI))
871 missing: list[tuple[str, SourceType]] = []
872 for source_type, mine, theirs in classes:
873 keys = _key_column(table, column, mine)
874 held = _key_column(sources, _PATH, theirs).cast(keys.type)
875 absent = pc.filter(keys, pc.invert(pc.is_in(keys, value_set=held)))
876 missing += [(key, source_type) for key in sorted(pc.unique(absent).to_pylist())]
877 return missing
880def _add_source_id(db: LanceDBConnection, name: str, column: str) -> None:
881 """Give table *name* its ``source_id`` from the path in *column*, indexed as the path is."""
882 from lancedb.index import BTree
884 table = db.open_table(name)
885 if _SOURCE_ID not in table.schema.names:
886 table.add_columns({_SOURCE_ID: f"md5({column})"})
887 if name == CHUNKS_TABLE:
888 table.update(where=_WIKI_CHUNK, values_sql={_SOURCE_ID: _WIKI_ID_SQL.format(key=column)})
889 indexes = _index_names(table)
890 if [column] in indexes.values() and [_SOURCE_ID] not in indexes.values():
891 table.create_index(_SOURCE_ID, config=BTree())
894def _migrate_sources_table(table: LanceTable) -> None:
895 """Give ``_sources`` an id per row, its path column, and every column it lacks. Lock held."""
896 from lancedb.index import BTree
898 if _SOURCE_ID not in table.schema.names:
899 table.add_columns({_SOURCE_ID: f"md5({_OLD_PATH})"})
900 if _OLD_PATH in table.schema.names:
901 # lancedb's stub types each alteration as an iterable; it takes one dict.
902 table.alter_columns({"path": _OLD_PATH, "rename": _PATH}) # type: ignore[dict-item]
903 if f"{_OLD_PATH}_idx" in _index_names(table):
904 # A key index keeps its name through the rename; the flush refreshes it by the new one.
905 table.create_index(_PATH, config=BTree(), replace=True)
906 table.create_index(_SOURCE_ID, config=BTree(), replace=True)
907 table.drop_index(f"{_OLD_PATH}_idx")
908 defaults = {name: f"CAST({SOURCE_STAT_UNKNOWN} AS BIGINT)" for name in _SOURCE_STAT_COLUMNS}
909 defaults |= {name: _NULL_TEXT_SQL for name in (*_SOURCE_META_COLUMNS, _HELD_OUT)}
910 missing = {name: sql for name, sql in defaults.items() if name not in table.schema.names}
911 if missing:
912 table.add_columns(missing)
915def _label_keyed(locations: SourceLocations) -> str:
916 """The document rows still named by a registered label in place of their location."""
917 labels = " OR ".join(under_sql(_PATH, label) for label in locations.roots)
918 return f"{_DOCUMENT} AND ({labels})"
921def _holds_label_keys(db: LanceDBConnection, locations: SourceLocations) -> bool:
922 """Whether a source row of *db* is still named by a registered label."""
923 sources = open_committed(db, SOURCES_TABLE)
924 if sources is None or not locations.roots or _PATH not in sources.schema.names:
925 return False
926 return bool(sources.count_rows(_label_keyed(locations)))
929def _pinned(db: LanceDBConnection, table: LanceTable) -> LanceTable:
930 """*table* held at its version now, read while later commits go to *table*."""
931 pinned: LanceTable = db.open_table(table.name)
932 pinned.checkout(table.version)
933 return pinned
936def _rekey_prefix(
937 db: LanceDBConnection, sources: LanceTable, where: str, old: str, new: str
938) -> None:
939 """Put each source row *where* selects at the place below *new* it has below *old*. Lock held.
941 A row whose new place a row it does not move holds folds into that row. A
942 selection of no row writes nothing. Rows stream in batches of
943 ``_FILL_BATCH_ROWS``, so memory does not follow the row count.
944 """
945 if not sources.count_rows(where):
946 return
947 _fold_landings(db, sources, where, old, new)
948 _rewrite_prefix(db, sources, where, old, new)
951def _rewrite_prefix(
952 db: LanceDBConnection, sources: LanceTable, where: str, old: str, new: str
953) -> None:
954 """Give each row *where* selects *new* in place of its head *old*: one update a batch.
956 Each update names its rows by their address in the version read, which a row
957 keeps until it is rewritten, so no update reads past its own rows.
958 """
959 query = _pinned(db, sources).search().where(where).select([_SOURCE_ID]).with_row_id(True)
960 values = {_PATH: _swap_head_sql(_PATH, old, new)}
961 for batch in query.limit(None).to_batches(_FILL_BATCH_ROWS):
962 rows = ", ".join(map(str, batch.column(_ROW_ID).to_pylist()))
963 sources.update(where=f"{_ROW_ID} IN ({rows})", values_sql=values)
966def _fold_landings(
967 db: LanceDBConnection, sources: LanceTable, where: str, old: str, new: str
968) -> None:
969 """Fold each row *where* selects into the non-wiki row not moving at its place below *new*.
971 The citations of the folded row cite the row it lands on; its own rows go.
972 """
973 landed = f"NOT ({_WIKI}) AND {under_sql(_PATH, new)} AND NOT ({where})"
974 if not sources.count_rows(landed):
975 return
976 pairs = _fill_pairs(db, sources, where, landed, (old, new))
977 found = pairs.search().where(f"{_MOVED_ID} IS NOT NULL").select([_MOVED_ID, _LANDING_ID])
978 for batch in found.limit(None).to_batches(_FILL_BATCH_ROWS):
979 moved, landing = batch.column(_MOVED_ID), batch.column(_LANDING_ID)
980 _fold_batch(db, sources, moved.to_pylist(), landing.to_pylist())
981 db.drop_table(_FILL_PAIRS_TABLE)
984def _fill_pairs(
985 db: LanceDBConnection, sources: LanceTable, where: str, landed: str, heads: tuple[str, str]
986) -> LanceTable:
987 """A scratch table of each *landed* row under its *old* path, with the row moving onto it.
989 It replaces any table of its name, so one a killed fold left is made again.
990 """
991 from lancedb.index import BTree
993 old, new = heads
994 pairs: LanceTable = db.create_table(
995 _FILL_PAIRS_TABLE, schema=_FILL_PAIRS_SCHEMA, mode="overwrite"
996 )
997 labelled: dict[str, str | Expr] = {
998 _PATH: _swap_head_sql(_PATH, new, old),
999 _LANDING_ID: _SOURCE_ID,
1000 }
1001 for batch in (
1002 sources.search().where(landed).select(labelled).limit(None).to_batches(_FILL_BATCH_ROWS)
1003 ):
1004 empty = pa.nulls(batch.num_rows, pa.utf8())
1005 pairs.add(
1006 pa.table(
1007 [batch.column(_PATH), batch.column(_LANDING_ID), empty], schema=_FILL_PAIRS_SCHEMA
1008 )
1009 )
1010 pairs.create_index(_PATH, config=BTree())
1011 moving = (
1012 _pinned(db, sources).search().where(where).select({_PATH: _PATH, _MOVED_ID: _SOURCE_ID})
1013 )
1014 for batch in moving.limit(None).to_batches(_FILL_BATCH_ROWS):
1015 pairs.merge_insert(_PATH).when_matched_update_all().execute(pa.Table.from_batches([batch]))
1016 return pairs
1019def _fold_batch(
1020 db: LanceDBConnection, sources: LanceTable, moved: list[str], landing: list[str]
1021) -> None:
1022 """Fold each of the ids *moved* into the id *landing* holds beside it. Lock held."""
1023 folded = f"{_SOURCE_ID} IN ({_in_list(moved)})"
1024 citations = open_committed(db, CITATIONS_TABLE)
1025 if citations is not None:
1026 _repoint_citations(citations, moved, landing)
1027 for name in _PER_SOURCE_TABLES:
1028 table = open_committed(db, name)
1029 if table is not None:
1030 table.delete(folded)
1031 _safe_delete_unlocked(sources, folded)
1034def _repoint_citations(citations: LanceTable, moved: list[str], landing: list[str]) -> None:
1035 """Make each citation of an id of *moved* cite the id of *landing* beside it. Lock held."""
1036 for start in range(0, len(moved), _CITATION_LOOKUP_IDS):
1037 keys = moved[start : start + _CITATION_LOOKUP_IDS]
1038 values = landing[start : start + _CITATION_LOOKUP_IDS]
1039 where = f"{_SOURCE_ID} IN ({_in_list(keys)})"
1040 citations.update(where=where, values_sql={_SOURCE_ID: _looked_up_sql(keys, values)})
1043def _looked_up_sql(keys: list[str], values: list[str]) -> str:
1044 """The SQL value of the item of *values* at the place of ``source_id`` in *keys*."""
1045 at = f"CAST(array_position(make_array({_in_list(keys)}), {_SOURCE_ID}) AS BIGINT)"
1046 return f"array_element(make_array({_in_list(values)}), {at})"
1049def _fill_locations(db: LanceDBConnection, locations: SourceLocations) -> None:
1050 """Give each source row a registered label names its location in place of that key. Lock held.
1052 A deeper root goes first, so a row an outer root keeps for a file of a root
1053 inside it folds into the inner root's row. Each batch is one commit, and a
1054 filled row matches no label, so a killed run is finished by the next.
1055 """
1056 sources = open_committed(db, SOURCES_TABLE)
1057 if sources is None or not _holds_label_keys(db, locations):
1058 return
1059 for root, label in locations.deepest_first:
1060 _rekey_prefix(db, sources, f"{_DOCUMENT} AND {under_sql(_PATH, label)}", label, root)
1061 log.info("Stored where each source of the library at %s is", db.uri)
1064def _refresh_key_index_unlocked(table: LanceTable, column: str, flush_rows: int) -> None:
1065 """Rebuild the BTREE index on *column* once enough appended rows sit past it.
1067 Caller holds ``write_lock()``. A delete filtered on *column* probes the index
1068 and scans only the rows past it, so its cost follows the flush size and not
1069 the table size.
1070 """
1071 from lancedb.index import BTree
1073 stats = table.index_stats(f"{column}_idx")
1074 unindexed = table.count_rows() if stats is None else stats.num_unindexed_rows
1075 if unindexed >= _KEY_INDEX_STALE_FLUSHES * flush_rows:
1076 table.create_index(column, config=BTree(), replace=True)
1079def _holds_rows(table: LanceTable) -> bool:
1080 """Whether *table* holds a row; a table that cannot be counted is taken to hold one."""
1081 try:
1082 return bool(table.count_rows())
1083 except Exception:
1084 return True
1087def _get_distance(chunk: SearchChunk) -> float:
1088 """Extract distance as a sortable float (inf for None)."""
1089 return chunk.distance if chunk.distance is not None else float("inf")
1092def _count_within_threshold(sorted_results: list[SearchChunk], threshold: float) -> int:
1093 """Count results whose distance is within the given threshold."""
1094 for i, r in enumerate(sorted_results):
1095 if _get_distance(r) > threshold:
1096 return i
1097 return len(sorted_results)
1100class Store:
1101 """LanceDB vector store: wraps all DB operations with config-driven defaults."""
1103 def __init__(
1104 self, config: Config, on_upgrade: DetailedProgressCallback = noop_callback
1105 ) -> None:
1106 self._config = config
1107 self._on_upgrade = on_upgrade
1108 # The latest event of an upgrade that runs, and when it began; None when none runs.
1109 self._upgrade: LibraryUpgradeEvent | None = None
1110 self._upgrade_started = 0.0
1111 # One open at a time: a caller that finds an open running waits for it.
1112 self._open_lock = threading.Lock()
1113 self._fts_ready: bool = False
1114 self._title_fts_ready: bool = False
1115 self._doc_prefix_warned: bool = False
1116 # Degradations the health endpoint reports; set where they are detected.
1117 self._fts_degraded: bool = False
1118 self._doc_prefix_mismatch: bool = False
1119 # Scalar indexes (source/chunk_type) are built at ingest; a serve-only
1120 # store builds them lazily from the search path.
1121 self._scalar_ready: bool = False
1122 # Tables whose failed key index refresh was already warned about.
1123 self._key_index_warned: set[str] = set()
1124 self._db: LanceDBConnection | None = None
1125 # Cache of {filename: ingested_at} rebuilt only when sources
1126 # mutate; callers (temporal filter) hit it per-query.
1127 self._source_ingested_cache: dict[str, str] | None = None
1128 # The source registry names are shown by.
1129 self._registry = SourceLocations({})
1130 self.crawl_pages = CrawlPages(self.get_db, config.lancedb_dir)
1132 def _fts_config(self) -> FTS:
1133 """Shared FTS index config: positionless (with_position=True overflows
1134 LanceDB's list encoding on optimize()) and stemmed for the configured
1135 corpus language."""
1136 from lancedb.index import FTS
1138 return FTS(with_position=False, language=self._config.fts_language)
1140 def _index_build_lock(self, blocking: bool) -> AbstractContextManager[None]:
1141 """Write lock for index builds: full budget from ingest, short from the read path."""
1142 return self._write_lock() if blocking else self._write_lock(_READ_LOCK_TIMEOUT)
1144 def _write_lock(self, timeout: float = LOCK_TIMEOUT) -> AbstractContextManager[None]:
1145 """Acquire the write lock keyed on *this* store's data directory.
1147 A per-instance ``Lilbee`` writes to its own ``lancedb_dir``; locking the
1148 global ``cfg`` dir instead would leave those writes uncoordinated across
1149 processes.
1150 """
1151 self.get_db() # opens the library, and so migrates it, outside the lock
1152 return write_lock(self._config.lancedb_dir, timeout)
1154 def _locations(self) -> SourceLocations:
1155 """The rule sources are named by: the registry, read again when config.toml changes."""
1156 if settings.refresh_registry(self._config):
1157 self._take_registry()
1158 return self._registry
1160 def _take_registry(self) -> None:
1161 """Name sources by the registry the config holds."""
1162 if self._config.linked_roots != self._registry.roots:
1163 self._registry = SourceLocations(self._config.linked_roots)
1164 self._invalidate_source_cache()
1166 @contextmanager
1167 def registry_pinned(self) -> Iterator[None]:
1168 """Name sources by the registry the config holds now, whatever changes it in the block."""
1169 self._take_registry()
1170 with settings.pin_registry(self._config):
1171 yield
1173 def _invalidate_source_cache(self) -> None:
1174 """Drop the cached {filename: ingested_at} map."""
1175 self._source_ingested_cache = None
1177 def source_ingested_at_map(self) -> dict[str, str]:
1178 """Return {filename: ingested_at} for every source, cached until mutation.
1180 Best-effort: a reader racing a concurrent invalidation can store a
1181 pre-mutation snapshot. The only consumer (temporal query filter)
1182 treats a missing/stale entry as "do not filter," so staleness
1183 degrades ranking precision, never correctness.
1184 """
1185 if self._source_ingested_cache is not None:
1186 return self._source_ingested_cache
1187 mapping = {s["filename"]: s.get("ingested_at", "") for s in self.get_sources()}
1188 self._source_ingested_cache = mapping
1189 return mapping
1191 def _chunks_schema(self) -> pa.Schema:
1192 return pa.schema(
1193 [
1194 pa.field(_SOURCE_ID, pa.utf8()),
1195 pa.field("content_type", pa.utf8()),
1196 pa.field(_CHUNK_TYPE, pa.utf8()),
1197 pa.field("page_start", pa.int32()),
1198 pa.field("page_end", pa.int32()),
1199 pa.field("line_start", pa.int32()),
1200 pa.field("line_end", pa.int32()),
1201 pa.field("chunk", pa.utf8()),
1202 pa.field(_CHUNK_INDEX, pa.int32()),
1203 pa.field(_TITLE_COLUMN, pa.utf8()),
1204 pa.field("vector", pa.list_(pa.float32(), self._config.embedding_dim)),
1205 ]
1206 )
1208 def _chunks_table(self) -> LanceTable:
1209 """Open/create the chunks table, adding the title column to pre-title tables."""
1210 table = ensure_table(self.get_db(), CHUNKS_TABLE, self._chunks_schema())
1211 if _TITLE_COLUMN not in table.schema.names:
1212 table.add_columns({_TITLE_COLUMN: _NULL_TEXT_SQL})
1213 self._backfill_stem_titles_unlocked(table)
1214 return table
1216 def _backfill_stem_titles_unlocked(self, table: LanceTable) -> None:
1217 """Backfill filename-stem titles for pre-upgrade rows. Caller holds ``write_lock()``.
1219 Without this the title arm only matches documents ingested after the
1220 upgrade. Extracted (H1/EXIF) titles still need ``lilbee rebuild``;
1221 failure leaves NULLs, the pre-backfill behavior.
1222 """
1223 from lilbee.data.title import derive_title # circular at module scope
1225 try:
1226 rows = table.search().select([_SOURCE_ID]).limit(None).to_list()
1227 sources = self.source_paths({r[_SOURCE_ID] for r in rows})
1228 filled = 0
1229 for source_id, source in sorted(sources.items(), key=lambda item: item[1]):
1230 title = derive_title(source)
1231 if not title:
1232 continue
1233 table.update(where=f"{_SOURCE_ID} = '{source_id}'", values={_TITLE_COLUMN: title})
1234 filled += 1
1235 log.info(
1236 "Backfilled filename titles for %d of %d existing sources; run "
1237 "`lilbee rebuild` to derive titles from document content",
1238 filled,
1239 len(sources),
1240 )
1241 except Exception:
1242 log.warning(
1243 "Title backfill failed; pre-upgrade rows keep NULL titles until `lilbee rebuild`",
1244 exc_info=True,
1245 )
1247 def get_meta(self) -> StoreMeta | None:
1248 """Return the persisted store metadata row, or ``None`` if unset."""
1249 table = self.open_table(META_TABLE)
1250 if table is None:
1251 return None
1252 rows = table.search().limit(None).to_list()
1253 if not rows:
1254 return None
1255 # _meta is meant to hold one row, but a swallowed delete on rewrite could
1256 # leave a stale one behind; take the newest so identity reads stay
1257 # deterministic rather than returning an arbitrary row.
1258 row = max(rows, key=lambda r: r["updated_at"])
1259 return StoreMeta(
1260 embedding_model=row["embedding_model"],
1261 embedding_dim=int(row["embedding_dim"]),
1262 schema_version=int(row["schema_version"]),
1263 updated_at=row["updated_at"],
1264 )
1266 def _write_meta_unlocked(self, *, embedding_model: str, embedding_dim: int) -> None:
1267 """Overwrite the single ``_meta`` row with the supplied identity.
1269 Caller must hold ``write_lock()``. Args are passed explicitly rather than
1270 re-read from ``self._config`` so the caller can snapshot cfg at a coherent
1271 instant and not race with a concurrent ``set_embedding_model``.
1272 """
1273 db = self.get_db()
1274 table = ensure_table(db, META_TABLE, _meta_schema())
1275 _safe_delete_unlocked(table, META_DELETE_ALL_PREDICATE)
1276 table.add(
1277 [
1278 {
1279 "embedding_model": embedding_model,
1280 "embedding_dim": embedding_dim,
1281 "schema_version": META_SCHEMA_VERSION,
1282 "updated_at": datetime.now(UTC).isoformat(),
1283 }
1284 ]
1285 )
1286 # A rebuild is the remedy the prefix warning names, so the next query
1287 # re-decides instead of logging a mismatch this write just resolved.
1288 self._doc_prefix_warned = False
1290 def restamp_meta(self) -> None:
1291 """Record the current embedding identity and schema version after a full rebuild."""
1292 with self._write_lock():
1293 self._write_meta_unlocked(
1294 embedding_model=self._config.embedding_model,
1295 embedding_dim=self._config.embedding_dim,
1296 )
1298 def _library_absent(self) -> bool:
1299 """True when the library is not open, has no folder and no older record file waits.
1301 An older record file is imported by the open, so a data dir that holds one is not absent.
1302 """
1303 return (
1304 self._db is None
1305 and not self._config.lancedb_dir.exists()
1306 and not legacy_files(self._config.data_root)
1307 )
1309 def _has_chunks(self) -> bool:
1310 """Return True when the chunks table exists and has at least one row.
1312 A data dir without a library folder holds none, and the answer leaves it uncreated.
1313 """
1314 if self._library_absent():
1315 return False
1316 chunks = self.open_table(CHUNKS_TABLE)
1317 return chunks is not None and chunks.count_rows() > 0
1319 def has_chunks(self) -> bool:
1320 """Public predicate: True iff the store currently holds at least one chunk."""
1321 return self._has_chunks()
1323 def initialize_meta_if_legacy(self) -> bool:
1324 """Pin a legacy store's identity to the current cfg if not already set.
1326 Returns ``True`` when a meta row was just written. No-op when meta already
1327 exists or no chunks are present. This is the path that converts a
1328 pre-upgrade store (chunks present, no ``_meta``) into a gated store. It
1329 snapshots cfg under the write lock to keep the recorded identity coherent
1330 with what the gate is comparing against.
1331 """
1332 if self.get_meta() is not None:
1333 return False
1334 if not self._has_chunks():
1335 return False
1336 embedding_model = self._config.embedding_model
1337 embedding_dim = self._config.embedding_dim
1338 with self._write_lock():
1339 # Re-check under the lock so two callers do not both warn-and-write.
1340 if self.get_meta() is not None:
1341 return False
1342 log.warning(
1343 "Legacy store has chunks but no _meta row. Initializing _meta from "
1344 "the current configuration (embedding_model=%s, embedding_dim=%d). "
1345 "If you changed embedding_model before upgrading, run `lilbee rebuild` "
1346 "to ensure the store is consistent.",
1347 embedding_model,
1348 embedding_dim,
1349 )
1350 self._write_meta_unlocked(embedding_model=embedding_model, embedding_dim=embedding_dim)
1351 return True
1353 def index_mismatch(
1354 self, model: str | None = None, dim: int | None = None
1355 ) -> EmbeddingModelMismatchError | None:
1356 """The drift between the persisted embedding identity and cfg, or None.
1358 *model* and *dim* judge a candidate embedder (a picked ref) in place of
1359 cfg's; each defaults to cfg's value.
1361 Pure check, no side effects, and the one verdict search, ingest, the
1362 embedding setter, sync and health all read. Migration of legacy stores
1363 (chunks present, no ``_meta``) is the caller's responsibility via
1364 ``initialize_meta_if_legacy``; rewriting a legacy bare-repo ``_meta`` row
1365 to the canonical full ref is the caller's responsibility via
1366 ``canonicalize_meta_if_legacy``. Safe to call from inside an existing
1367 ``write_lock()`` (no recursive lock attempt). cfg fields are snapshotted
1368 at entry so the comparison is coherent even if another thread mutates
1369 them mid-call.
1370 """
1371 current_model = self._config.embedding_model if model is None else model
1372 current_dim = self._config.embedding_dim if dim is None else dim
1373 meta = self.get_meta()
1374 if meta is None:
1375 return None
1376 if refs_compatible(
1377 meta["embedding_model"], current_model, meta["embedding_dim"], current_dim
1378 ):
1379 return None
1380 return EmbeddingModelMismatchError(
1381 persisted_model=meta["embedding_model"],
1382 persisted_dim=meta["embedding_dim"],
1383 current_model=current_model,
1384 current_dim=current_dim,
1385 )
1387 def _ensure_embedding_compat(self) -> None:
1388 """Raise when the persisted embedding identity drifts from cfg."""
1389 mismatch = self.index_mismatch()
1390 if mismatch is not None:
1391 raise mismatch
1393 def _doc_prefix_is_stale(self) -> bool:
1394 """Whether the stored documents predate the embedding family's document prefix."""
1395 meta = self.get_meta()
1396 if meta is None:
1397 return False
1398 profile = resolve_embedding_profile(self._config.embedding_model)
1399 return bool(profile.doc_prefix) and meta["schema_version"] < profile.doc_prefix_since
1401 def _warn_stale_doc_prefix(self) -> None:
1402 """Warn once when the embedding family's document prefix postdates this store.
1404 Queries would carry the family prefix while stored documents do not;
1405 the mismatch degrades quality silently until a rebuild re-embeds. The
1406 check reads the metadata row, which the query path cannot pay per search.
1407 """
1408 if self._doc_prefix_warned:
1409 return
1410 self._doc_prefix_warned = True
1411 if not self._doc_prefix_is_stale():
1412 return
1413 self._doc_prefix_mismatch = True
1414 log.warning(
1415 "This index predates '%s' document prefixes: queries are prefixed "
1416 "but stored documents are not. Run `lilbee rebuild` to re-embed.",
1417 self._config.embedding_model,
1418 )
1420 def upgrade_progress(self) -> LibraryUpgradeEvent | None:
1421 """The upgrade that runs, timed to now, or None when none does."""
1422 upgrade = self._upgrade
1423 if upgrade is None:
1424 return None
1425 return upgrade.model_copy(update={"elapsed_s": time.monotonic() - self._upgrade_started})
1427 def _upgrade_warnings(self) -> list[HealthWarning]:
1428 """The warning of an upgrade that runs, or none."""
1429 upgrade = self.upgrade_progress()
1430 if upgrade is None:
1431 return []
1432 message = f"{UPGRADE_HEADLINE} {upgrade.text}"
1433 remedy = "Server requests answer 503 after 10 s. Try again when it ends."
1434 return [HealthWarning(code=WarningCode.LIBRARY_UPGRADING, message=message, remedy=remedy)]
1436 def _index_mismatch_warning(self) -> HealthWarning | None:
1437 """The stale-index degradation, from the same verdict search refuses on."""
1438 mismatch = self.index_mismatch()
1439 if mismatch is None:
1440 return None
1441 return HealthWarning(
1442 code=WarningCode.INDEX_EMBEDDING_MISMATCH,
1443 message=(
1444 f"The index was built with embedding model '{mismatch.persisted_model}' "
1445 f"({mismatch.persisted_dim} dims), but '{mismatch.current_model}' "
1446 f"({mismatch.current_dim} dims) is configured, so search refuses it."
1447 ),
1448 remedy=(
1449 f"Rebuild the index under '{mismatch.current_model}', or switch the "
1450 f"embedding model back to '{mismatch.persisted_model}'."
1451 ),
1452 )
1454 def health_warnings(self) -> list[HealthWarning]:
1455 """Retrieval degradations a client should know about before it searches.
1457 A library that is not open reports only an upgrade that runs: opening it here
1458 would wait for that upgrade.
1459 """
1460 if self._db is None:
1461 return self._upgrade_warnings()
1462 warnings: list[HealthWarning] = []
1463 if (stale := self._index_mismatch_warning()) is not None:
1464 warnings.append(stale)
1465 if self._fts_degraded:
1466 warnings.append(
1467 HealthWarning(
1468 code=WarningCode.FTS_UNAVAILABLE,
1469 message=(
1470 "Keyword search is unavailable, so answers are drawn from "
1471 "vector similarity alone and exact terms may be missed."
1472 ),
1473 remedy="Run `lilbee rebuild` to rebuild the search index.",
1474 )
1475 )
1476 table = self.open_table(CHUNKS_TABLE)
1477 indices = _index_registry(table, self._config.lancedb_dir) if table is not None else None
1478 if table is not None and indices is not None:
1479 dangling_scalar = _scalar_index_dangling(table, indices, self._config.lancedb_dir)
1480 if dangling_scalar:
1481 warnings.append(
1482 HealthWarning(
1483 code=WarningCode.SCALAR_INDEX_UNAVAILABLE,
1484 message=(
1485 "The scalar index on "
1486 + ", ".join(dangling_scalar)
1487 + " is missing its files, so filtered search is degraded."
1488 ),
1489 remedy="Run `lilbee rebuild` to rebuild the search index.",
1490 )
1491 )
1492 if self._doc_prefix_mismatch and not self._doc_prefix_is_stale():
1493 # The remedy is `lilbee rebuild`, which usually runs in another
1494 # process, so the query path's cached verdict outlives the fix.
1495 self._doc_prefix_mismatch = False
1496 if self._doc_prefix_mismatch:
1497 warnings.append(
1498 HealthWarning(
1499 code=WarningCode.EMBEDDING_PREFIX_MISMATCH,
1500 message=(
1501 f"This index predates '{self._config.embedding_model}' document "
1502 "prefixes: queries are prefixed but stored documents are not, "
1503 "so retrieval is less accurate."
1504 ),
1505 remedy="Run `lilbee rebuild` to re-embed the documents.",
1506 )
1507 )
1508 return warnings
1510 def assert_embedding_compatible(self) -> None:
1511 """Run the full embedding-identity gate (legacy init, canonicalize, check).
1513 Mirrors the gate ``search`` applies. Callers that write under a fresh
1514 embedder (import) use this to fail before any destructive work when the
1515 store was built by a different model.
1516 """
1517 self.initialize_meta_if_legacy()
1518 self.canonicalize_meta_if_legacy()
1519 self._ensure_embedding_compat()
1521 def _needs_canonical_meta_rewrite(
1522 self, meta: StoreMeta | None, current_model: str, current_dim: int
1523 ) -> bool:
1524 """True iff *meta* is the legacy form and refs-compatible with current cfg."""
1525 if meta is None or meta["embedding_model"] == current_model:
1526 return False
1527 return refs_compatible(
1528 meta["embedding_model"], current_model, meta["embedding_dim"], current_dim
1529 )
1531 def canonicalize_meta_if_legacy(self) -> bool:
1532 """Rewrite a legacy bare-repo ``_meta`` row to the canonical full ref.
1534 Pre-canonical lilbee persisted only ``<org>/<repo>`` in
1535 ``_meta.embedding_model``. The current code persists the full
1536 ``<org>/<repo>/<filename>.gguf``. When the two refer to the same
1537 model under :func:`refs_compatible` but differ as raw strings, the
1538 meta row is rewritten so the legacy name never surfaces. Returns
1539 ``True`` on write; ``False`` when missing, already canonical, or
1540 incompatible (the gate handles incompatibility).
1541 """
1542 current_model = self._config.embedding_model
1543 current_dim = self._config.embedding_dim
1544 if not self._needs_canonical_meta_rewrite(self.get_meta(), current_model, current_dim):
1545 return False
1546 with self._write_lock():
1547 meta = self.get_meta() # re-read under the lock for racing callers
1548 if not self._needs_canonical_meta_rewrite(meta, current_model, current_dim):
1549 return False
1550 assert meta is not None # filtered above # noqa: S101
1551 log.info(
1552 "Migrating legacy embedding ref in store meta: %r -> %r",
1553 meta["embedding_model"],
1554 current_model,
1555 )
1556 self._write_meta_unlocked(embedding_model=current_model, embedding_dim=current_dim)
1557 return True
1559 def get_db(self) -> LanceDBConnection:
1560 """The open library; the first call opens it, and upgrades it when it is older."""
1561 db = self._db
1562 if db is None:
1563 with self._open_lock:
1564 db = self._open_held()
1565 return db
1567 @property
1568 def is_open(self) -> bool:
1569 """Whether the library is open, so a call needs no open or upgrade first."""
1570 return self._db is not None
1572 def open_in_background(self) -> None:
1573 """Open the library on a daemon thread; ``wait_for_open`` waits from this call on.
1575 An open that already runs is left to finish: its waiters wait for it.
1576 """
1577 if not self._open_lock.acquire(blocking=False):
1578 return
1579 threading.Thread(target=self._open_then_release, name="library-open", daemon=True).start()
1581 def wait_for_open(self, timeout: float) -> bool:
1582 """Wait up to *timeout* seconds for another thread to open the library.
1584 Return False when the open is still running at the deadline.
1585 """
1586 if not self._open_lock.acquire(timeout=timeout):
1587 return False
1588 self._open_lock.release()
1589 return True
1591 def _open_then_release(self) -> None:
1592 """Open the library, then free the open lock ``open_in_background`` took."""
1593 try:
1594 self._open_held()
1595 finally:
1596 self._open_lock.release()
1598 def _open_held(self) -> LanceDBConnection:
1599 """The library, opened unless it is open. Caller holds the open lock."""
1600 if self._db is not None:
1601 return self._db
1602 from lancedb.db import LanceDBConnection
1604 self._config.lancedb_dir.mkdir(parents=True, exist_ok=True)
1605 db = LanceDBConnection(
1606 str(self._config.lancedb_dir), read_consistency_interval=READ_CONSISTENCY_INTERVAL
1607 )
1608 if self._upgrade_needed(db):
1609 self._upgrade_library(db)
1610 self._db = db
1611 return db
1613 def _upgrade_needed(self, db: LanceDBConnection) -> bool:
1614 """Whether *db* is keyed by path, keeps the older record files, or keys a row by label."""
1615 return (
1616 _upgrade_pending(db)
1617 or _held_unmigrated(db, self._config.data_root)
1618 or _holds_label_keys(db, self._locations())
1619 or bool(legacy_crawl_files(self._config.data_dir))
1620 )
1622 def _upgrade_library(self, db: LanceDBConnection) -> None:
1623 """Upgrade *db* under the write lock, waiting for it without a deadline, and report it."""
1624 self._upgrade_started = time.monotonic()
1625 try:
1626 with ExitStack() as held:
1627 try:
1628 held.enter_context(write_lock(self._config.lancedb_dir, timeout=0.0))
1629 except LockTimeoutError:
1630 self._report_upgrade(UpgradePhase.WAITING)
1631 held.enter_context(write_lock(self._config.lancedb_dir, timeout=None))
1632 # Another process can have upgraded it during the wait.
1633 if self._upgrade_needed(db):
1634 self._upgrade_unlocked(db)
1635 except BaseException:
1636 self._report_upgrade(UpgradePhase.FAILED)
1637 raise
1638 finally:
1639 self._upgrade = None
1640 self._report_upgrade(UpgradePhase.DONE)
1642 def _upgrade_unlocked(self, db: LanceDBConnection) -> None:
1643 """Run each step of the upgrade that *db* still needs. Caller holds the write lock."""
1644 held = _held_unmigrated(db, self._config.data_root)
1645 if held:
1646 _add_held_columns(db)
1647 if _upgrade_pending(db):
1648 self._migrate_unlocked(db)
1649 if held:
1650 self._import_held_unlocked(db)
1651 if legacy_crawl_files(self._config.data_dir):
1652 import_legacy_crawl_files(db, self._config.data_dir)
1653 _fill_locations(db, self._locations())
1655 def _report_upgrade(self, phase: UpgradePhase, step: _TableStep | None = None) -> None:
1656 """Pass *phase* to the reporter, timed from the upgrade's start; record it while it runs."""
1657 event = LibraryUpgradeEvent(
1658 phase=phase,
1659 elapsed_s=time.monotonic() - self._upgrade_started,
1660 **({} if step is None else step._asdict()),
1661 )
1662 self._upgrade = event if event.running else None
1663 self._on_upgrade(EventType.LIBRARY_UPGRADE, event)
1665 def _migrate_unlocked(self, db: LanceDBConnection) -> None:
1666 """Key every table by ``source_id``. Caller holds the write lock.
1668 Each step reads the table's columns, and the stamp reads ``_meta``, so a
1669 run that was killed is finished by the next one. An existing source takes
1670 the md5 of its path as its id. A store stamped 2 is stamped 3 at the end; a
1671 store stamped 1 predates the document prefix, and keeps 1 so that the
1672 stale-prefix warning still names it.
1673 """
1674 names = table_names(db)
1675 steps = _upgrade_steps(db, names)
1676 keyed = [(name, _PATH_KEY_BY_TABLE[name]) for name, _ in steps if name != SOURCES_TABLE]
1677 rows: dict[str, int] = {}
1678 done = 0
1679 # LanceDB adds a column one fragment at a time: each table's runs on its own thread
1680 # from the start, and its path column goes once its every path has a source row.
1681 # Only this thread reports, in table order, so the count of tables done only rises.
1682 with ThreadPoolExecutor(len(_PATH_KEY_COLUMNS), _UPGRADE_THREAD_PREFIX) as pool:
1683 added = [pool.submit(_add_source_id, db, name, column) for name, column in keyed]
1684 for name, table in steps:
1685 rows[name] = table.count_rows()
1686 work = TableWork.UPDATING if name == SOURCES_TABLE else TableWork.CHECKING
1687 step = _TableStep(work, _UPGRADE_LABELS[name], rows[name], done, len(steps))
1688 self._report_upgrade(UpgradePhase.TABLE, step)
1689 if name == SOURCES_TABLE:
1690 _migrate_sources_table(table)
1691 done += 1
1692 else:
1693 self._hold_missing_paths(db, table, _PATH_KEY_BY_TABLE[name])
1694 for (name, column), column_added in zip(keyed, added, strict=True):
1695 step = _TableStep(
1696 TableWork.ADDING_IDS, _UPGRADE_LABELS[name], rows[name], done, len(steps)
1697 )
1698 self._report_upgrade(UpgradePhase.TABLE, step)
1699 column_added.result()
1700 db.open_table(name).drop_columns([column])
1701 done += 1
1702 if META_TABLE in names:
1703 db.open_table(META_TABLE).update(
1704 where="schema_version >= 2", values={"schema_version": META_SCHEMA_VERSION}
1705 )
1706 log.info("Keyed the library at %s by source id", self._config.lancedb_dir)
1708 def _import_held_unlocked(self, db: LanceDBConnection) -> None:
1709 """Move the held-out records of the older files into ``_sources``, then delete them.
1711 Gives each recorded name its record in one commit; a row that holds an
1712 indexed document takes no removal. Lock held. The files go last; an open
1713 killed before then imports the same records again onto the same rows.
1714 """
1715 sources = ensure_table(db, SOURCES_TABLE, _sources_schema())
1716 root = self._config.data_root
1717 writes = [
1718 _SourceWrite(name, SourceType.DOCUMENT, record=record)
1719 for name, record in read_legacy_records(root).items()
1720 ]
1721 self._put_sources_unlocked(sources, writes, keep_documents=True)
1722 remove_legacy_files(root)
1723 log.info("Moved the held-out file records of %s into the index", root)
1725 def _hold_missing_paths(self, db: LanceDBConnection, table: LanceTable, column: str) -> None:
1726 """Give each key in *column* of *table* no source row of its class holds a row. Lock held.
1728 A wiki page, or a removed source a citation names, gets one under the id the
1729 column's key and type give, so no key is lost when the column goes.
1730 """
1731 sources = ensure_table(db, SOURCES_TABLE, _sources_schema())
1732 writes = [
1733 _SourceWrite(key, kind, new_id=_wiki_id(key) if _is_wiki(kind) else _md5(key))
1734 for key, kind in _keys_missing_from(sources, table, column)
1735 ]
1736 self._put_sources_unlocked(sources, writes)
1738 def open_table(self, name: str) -> LanceTable | None:
1739 """Open a table if it exists, otherwise return None."""
1740 return open_committed(self.get_db(), name)
1742 def ensure_fts_index(self, *, blocking: bool = True) -> None:
1743 """Create the chunks FTS index, or run ``optimize()`` once it exists.
1745 ``optimize()`` folds newly added rows into the FTS index and also
1746 runs LanceDB's default compaction + version pruning (default prune
1747 window: 7 days). Work scales with recent deltas rather than total
1748 chunk count, so large corpora no longer pay the full
1749 ``create_index(config=FTS(), replace=True)`` rebuild cost on every sync.
1751 ``blocking=False`` (the search path) marks an existing index ready
1752 without the lock and skips maintenance when another process holds it
1753 or the filesystem cannot lock, so a long concurrent ingest cannot stall
1754 or fail a query. An unusable
1755 index registry never counts as an existing index.
1756 """
1757 probe = self.open_table(CHUNKS_TABLE)
1758 if probe is None:
1759 return
1760 indices = _index_registry(probe, self._config.lancedb_dir)
1761 if indices is not None and _has_fts_index(indices):
1762 self._fts_ready = True
1763 try:
1764 with self._index_build_lock(blocking):
1765 self._ensure_fts_index_unlocked()
1766 except _UPKEEP_SKIPPED as exc:
1767 if blocking:
1768 raise
1769 log.debug("Skipped FTS index maintenance: %s", exc)
1771 def _ensure_fts_index_unlocked(self) -> None:
1772 """Body of ``ensure_fts_index``. Caller holds ``write_lock()``."""
1773 table = self.open_table(CHUNKS_TABLE)
1774 if table is None:
1775 return
1776 try:
1777 indices = _index_registry(table, self._config.lancedb_dir)
1778 if indices is None:
1779 self._repair_fts_registry(table)
1780 return
1781 if _has_fts_index(indices):
1782 self._fts_ready = True
1783 try:
1784 # One optimize folds new rows into every index on the table.
1785 table.optimize()
1786 log.debug("FTS index optimized on '%s'", CHUNKS_TABLE)
1787 if _index_registry(table, self._config.lancedb_dir) is None:
1788 # Suspected producer of an index without files; unconfirmed.
1789 self._repair_fts_registry(table)
1790 return
1791 except Exception as exc:
1792 if _is_fts_position_overflow(exc):
1793 # Positional indexes overflow on optimize(); rebuild
1794 # positionless once.
1795 self._rebuild_fts(table, "a positional index overflowed on optimize()")
1796 else:
1797 log.warning(
1798 "FTS optimize() failed; the existing index still serves hybrid search",
1799 exc_info=True,
1800 )
1801 else:
1802 # Positionless: with_position=True overflows LanceDB's list
1803 # encoding on optimize(), and nothing issues phrase queries.
1804 table.create_index(_CHUNK_COLUMN, config=self._fts_config(), replace=False)
1805 self._fts_ready = True
1806 log.debug("FTS index created on '%s'", CHUNKS_TABLE)
1807 # Only the opt-in title arm needs the title index.
1808 if self._config.title_search:
1809 self._ensure_title_fts_unlocked(table)
1810 except Exception:
1811 log.debug("FTS index ensure failed (empty table?)", exc_info=True)
1813 def _ensure_title_fts_unlocked(self, table: LanceTable) -> None:
1814 """Create the title FTS index when the column exists. Caller holds ``write_lock()``.
1816 Failure never blocks the chunk index: the title arm feature-detects the
1817 index per query, so a store without it simply searches without titles.
1818 """
1819 indices = _index_registry(table, self._config.lancedb_dir)
1820 if indices is None:
1821 self._repair_fts_registry(table)
1822 return
1823 if _TITLE_COLUMN not in table.schema.names or _has_fts_index(indices, _TITLE_COLUMN):
1824 self._title_fts_ready = _has_fts_index(indices, _TITLE_COLUMN)
1825 return
1826 try:
1827 # Positionless for the same reason as the chunk index.
1828 table.create_index(_TITLE_COLUMN, config=self._fts_config(), replace=False)
1829 self._title_fts_ready = True
1830 log.debug("Title FTS index created on '%s'", CHUNKS_TABLE)
1831 except Exception:
1832 # Only reached with title_search enabled, so a silent failure means
1833 # the user's opted-in title arm quietly does nothing. Warn, don't hide.
1834 log.warning(
1835 "Title FTS index creation failed; the title-search arm will "
1836 "contribute nothing until it can be built",
1837 exc_info=True,
1838 )
1840 def ensure_title_fts_index(self, *, blocking: bool = True) -> None:
1841 """Build the title FTS index for a title_search toggle after startup.
1843 Without this, a process that latched ``_fts_ready`` before the toggle
1844 never builds the index and the title arm stays silently dead until the
1845 next ingest or restart.
1846 """
1847 table = self.open_table(CHUNKS_TABLE)
1848 if table is None:
1849 return
1850 indices = _index_registry(table, self._config.lancedb_dir)
1851 if indices is not None and _has_fts_index(indices, _TITLE_COLUMN):
1852 self._title_fts_ready = True
1853 return
1854 try:
1855 with self._index_build_lock(blocking):
1856 self._ensure_title_fts_unlocked(table)
1857 except _UPKEEP_SKIPPED as exc:
1858 if blocking:
1859 raise
1860 log.debug("Skipped title FTS build: %s", exc)
1862 def _rebuild_fts(self, table: LanceTable, reason: str) -> None:
1863 """Replace the FTS indexes with fresh positionless ones. Caller holds the lock.
1865 The one-shot remediation for an index built ``with_position=True`` that
1866 overflows on every ``optimize()``. The title index is rebuilt too when
1867 the title arm is enabled.
1868 """
1869 try:
1870 table.create_index(_CHUNK_COLUMN, config=self._fts_config(), replace=True)
1871 if self._config.title_search and _TITLE_COLUMN in table.schema.names:
1872 table.create_index(_TITLE_COLUMN, config=self._fts_config(), replace=True)
1873 log.warning("Rebuilt the FTS index because %s", reason)
1874 except Exception:
1875 log.warning(
1876 "Positionless FTS rebuild failed; the existing index still serves",
1877 exc_info=True,
1878 )
1880 def _repair_fts_registry(self, table: LanceTable) -> list[IndexConfig] | None:
1881 """Rebuild the chunks FTS indexes when the registry is ``None``. Caller holds the lock.
1883 ``None`` means the listing raised or an FTS index lost its files.
1884 Each FTS index (chunk text, plus title when the title arm is on) is
1885 rebuilt with ``replace=True`` in its own try. The registry is then
1886 read once more; the directory check on a readable listing decides
1887 whether a scalar or vector index needs its own rebuild.
1889 Returns the registry after the rebuild, ``None`` (logged once) when
1890 it is still unusable.
1891 """
1892 columns = [_CHUNK_COLUMN]
1893 if self._config.title_search and _TITLE_COLUMN in table.schema.names:
1894 columns.append(_TITLE_COLUMN)
1895 for column in columns:
1896 try:
1897 table.create_index(column, config=self._fts_config(), replace=True)
1898 except Exception:
1899 log.warning(
1900 "Could not rebuild the FTS index on '%s.%s'", table.name, column, exc_info=True
1901 )
1902 continue
1903 log.warning(
1904 "Rebuilt the FTS index on '%s.%s' because an FTS index is unusable",
1905 table.name,
1906 column,
1907 )
1908 if column == _CHUNK_COLUMN:
1909 self._fts_ready = True
1910 else:
1911 self._title_fts_ready = True
1912 indices = _index_registry(table, self._config.lancedb_dir)
1913 if indices is None:
1914 log.warning(
1915 "The index registry on '%s' is still unusable after the FTS rebuild; "
1916 "run `lilbee rebuild` to rebuild the search index",
1917 table.name,
1918 )
1919 return indices
1921 # Tables and (column, index_type) pairs the query path filters by.
1922 # chunk_concepts serves the concept boost (ConceptGraph._chunk_concepts_from);
1923 # without its index every boosted query full-scans the table.
1924 _SCALAR_TARGETS: tuple[tuple[str, tuple[tuple[str, ScalarIndexType], ...]], ...] = (
1925 (CHUNKS_TABLE, ((_SOURCE_ID, "BTREE"), (_CHUNK_TYPE, "BITMAP"))),
1926 (CHUNK_CONCEPTS_TABLE, ((_SOURCE_ID, "BTREE"),)),
1927 )
1929 def ensure_scalar_indexes(self, *, blocking: bool = True) -> None:
1930 """Build scalar indexes on the columns lilbee filters by.
1932 ``source`` and ``chunk_type`` predicates run as prefilters (LanceDB's
1933 default), but without an index each is a full-table scan. Readiness
1934 latches only when every target table exists and is covered, so a table
1935 created later (chunk_concepts under serve ordering) still gets its
1936 index on a following call. A table whose scalar index lost its files
1937 counts as pending, so the build step replaces it. A table whose index
1938 registry is unusable keeps readiness unlatched until the FTS path
1939 repairs it. The lock is taken only
1940 when there is something to build; ``blocking=False`` (the search path)
1941 skips the build when another process holds it or the filesystem cannot
1942 lock, instead of stalling or failing the query.
1943 """
1944 pending = []
1945 complete = True
1946 for name, columns in self._SCALAR_TARGETS:
1947 table = self.open_table(name)
1948 if table is None:
1949 complete = False
1950 continue
1951 indices = _index_registry(table, self._config.lancedb_dir)
1952 if indices is None:
1953 # The FTS path repairs an unusable registry; check again then.
1954 complete = False
1955 continue
1956 names = table.schema.names
1957 dangling = _scalar_index_dangling(table, indices, self._config.lancedb_dir)
1958 needs_build = any(
1959 c in names and (not _has_scalar_index(indices, c) or c in dangling)
1960 for c, _ in columns
1961 )
1962 if needs_build:
1963 pending.append((name, columns))
1964 if not pending:
1965 self._scalar_ready = complete
1966 return
1967 try:
1968 with self._index_build_lock(blocking):
1969 for name, columns in pending:
1970 self._ensure_scalar_index_on(name, columns)
1971 self._scalar_ready = complete
1972 except _UPKEEP_SKIPPED as exc:
1973 if blocking:
1974 raise
1975 log.debug("Skipped scalar index build: %s", exc)
1977 def _ensure_scalar_index_on(
1978 self, table_name: str, columns: tuple[tuple[str, ScalarIndexType], ...]
1979 ) -> None:
1980 """Build the given (column, index_type) scalar indexes on *table_name*.
1982 Caller holds ``write_lock()``. A registered index whose files are gone
1983 is replaced; nothing is built while the index registry is unusable.
1984 Each column gets its own try so one failure does not skip the rest; a
1985 failure on a populated table warns (the prefilter speedup is silently
1986 lost) while an empty table's is debug.
1987 """
1988 table = self.open_table(table_name)
1989 if table is None:
1990 return
1991 indices = _index_registry(table, self._config.lancedb_dir)
1992 if indices is None:
1993 return
1994 names = table.schema.names
1995 fail_level = logging.WARNING if table.count_rows() > 0 else logging.DEBUG
1996 dangling = _scalar_index_dangling(table, indices, self._config.lancedb_dir)
1997 for column, index_type in columns:
1998 if column not in names or (
1999 _has_scalar_index(indices, column) and column not in dangling
2000 ):
2001 continue
2002 try:
2003 table.create_scalar_index(column, index_type=index_type, replace=column in dangling)
2004 log.debug("Scalar (%s) index created on '%s.%s'", index_type, table_name, column)
2005 except Exception:
2006 log.log(
2007 fail_level,
2008 "Scalar index create failed on '%s.%s'",
2009 table_name,
2010 column,
2011 exc_info=True,
2012 )
2014 def _create_vector_index(self, table: LanceTable) -> None:
2015 """Build the IVF_PQ index on the vector column, replacing any existing one."""
2016 from lancedb.index import IvfPq
2018 table.create_index("vector", config=IvfPq(distance_type=_VECTOR_METRIC), replace=True)
2020 def ensure_vector_index(self, *, force: bool = False) -> bool:
2021 """Build or refresh the ANN vector index when the corpus is large enough.
2023 Below ``cfg.ann_index_threshold`` (or when it is 0) the store keeps exact
2024 flat search, which is faster and exact for small vaults and is all a
2025 laptop needs. Once an index exists, ``optimize()`` folds new rows in.
2026 Pass ``force=True`` to build regardless of the threshold (publish flow).
2027 An unusable index registry gets its FTS repair first, and a registered
2028 index whose files are gone is rebuilt instead of optimized. Returns
2029 True when an index was created, refreshed, or rebuilt.
2030 """
2031 threshold = self._config.ann_index_threshold
2032 with self._write_lock():
2033 table = self.open_table(CHUNKS_TABLE)
2034 if table is None:
2035 return False
2036 indices = _index_registry(table, self._config.lancedb_dir)
2037 if indices is None:
2038 indices = self._repair_fts_registry(table)
2039 if indices is None:
2040 return False
2041 if _vector_index_dangling(table, indices, self._config.lancedb_dir):
2042 return self._rebuild_dangling_vector_index(table)
2043 if _has_vector_index(indices):
2044 table.optimize()
2045 log.debug("Vector index optimized on '%s'", CHUNKS_TABLE)
2046 if _index_registry(table, self._config.lancedb_dir) is None:
2047 # Repair in the same step when the prune dropped FTS files.
2048 self._repair_fts_registry(table)
2049 return True
2050 if not force and (threshold <= 0 or table.count_rows() < threshold):
2051 return False
2052 return self._build_vector_index(table)
2054 def _rebuild_dangling_vector_index(self, table: LanceTable) -> bool:
2055 """Replace a registered vector index whose files are gone. Caller holds the lock."""
2056 try:
2057 self._create_vector_index(table)
2058 except Exception:
2059 log.warning(
2060 "Could not rebuild the vector index on '%s' whose files are missing",
2061 CHUNKS_TABLE,
2062 exc_info=True,
2063 )
2064 return False
2065 log.warning("Rebuilt the vector index on '%s' because its files are missing", CHUNKS_TABLE)
2066 return True
2068 def _build_vector_index(self, table: LanceTable) -> bool:
2069 """Build the first vector index, warning on failure. Caller holds the lock."""
2070 try:
2071 self._create_vector_index(table)
2072 log.info("Vector ANN index created on '%s'", CHUNKS_TABLE)
2073 return True
2074 except Exception:
2075 log.warning(
2076 "Vector ANN index build failed on '%s' at %d rows; search falls back "
2077 "to exact flat scan, which is slow at this scale. Free up memory/disk "
2078 "and re-run to rebuild the index.",
2079 CHUNKS_TABLE,
2080 table.count_rows(),
2081 exc_info=True,
2082 )
2083 return False
2085 def _add_chunks_unlocked(self, records: list[dict]) -> int:
2086 """Add chunk records and return the count. Caller must hold ``write_lock()``."""
2087 embedding_model = self._config.embedding_model
2088 embedding_dim = self._config.embedding_dim
2089 self._ensure_embedding_compat()
2090 self._fts_ready = False
2091 self._scalar_ready = False
2092 if not records:
2093 return 0
2094 _check_vector_dims(records, embedding_dim)
2095 ids = self._ensure_ids_unlocked(
2096 (record["source"], _source_type_of(record[_CHUNK_TYPE])) for record in records
2097 )
2098 table = self._chunks_table()
2099 table.add(_keyed(records, "source", ids))
2100 if self.get_meta() is None:
2101 self._write_meta_unlocked(embedding_model=embedding_model, embedding_dim=embedding_dim)
2102 return len(records)
2104 def add_chunks(self, records: list[dict]) -> int:
2105 """Add chunk records to the store. Returns count added.
2107 Raises ``EmbeddingModelMismatchError`` if the persisted ``_meta`` row was
2108 written under a different embedding model than the current ``cfg``. On the
2109 first write to a fresh store, ``_meta`` is initialized from the current cfg.
2111 The gate runs inside the write lock and uses a single cfg snapshot so a
2112 concurrent ``set_embedding_model`` cannot slip a write in past a stale
2113 compatibility check.
2114 """
2115 with self._write_lock():
2116 return self._add_chunks_unlocked(records)
2118 def replace_chunks(self, source: str, chunk_type: ChunkType, records: list[dict]) -> int:
2119 """Replace *source*'s chunk rows of *chunk_type* with *records* under one write lock.
2121 Same compatibility and dimension gates as :meth:`add_chunks`, both run
2122 before the delete so a rejected write leaves the existing rows in place.
2123 The lock serializes writers; delete and add are separate commits, so a
2124 concurrent reader can briefly see the rows absent, and callers retry on
2125 a crash between the two. A delete failure propagates, following the same
2126 rule as :meth:`_delete_rows_of_unlocked`: swallowed, the caller would
2127 read a success-shaped result over rows still describing the old body.
2128 Returns the count added.
2129 """
2130 with self._write_lock():
2131 self._ensure_embedding_compat()
2132 _check_vector_dims(records, self._config.embedding_dim)
2133 source_id = self._ensure_ids_unlocked([(source, _source_type_of(chunk_type))])[source]
2134 table = self._chunks_table()
2135 table.delete(f"{_SOURCE_ID} = '{source_id}' AND {_chunk_type_predicate(chunk_type)}")
2136 return self._add_chunks_unlocked(records)
2138 def bm25_probe(
2139 self, query_text: str, top_k: int = 5, chunk_type: ChunkType | None = None
2140 ) -> list[SearchChunk]:
2141 """Quick BM25-only search for confidence checking. Returns up to top_k results.
2143 When *chunk_type* is set, only chunks of that type ("raw" or "wiki") are returned.
2144 """
2145 table = self.open_table(CHUNKS_TABLE)
2146 if table is None:
2147 return []
2148 if not self._fts_ready:
2149 self.ensure_fts_index(blocking=False)
2150 if not self._fts_ready:
2151 return []
2152 try:
2153 results = self._with_paths(
2154 _id_chunks(_lexical_rows(table, query_text, top_k, chunk_type))
2155 )
2156 norms = normalized_bm25([r.bm25_score or 0.0 for r in results])
2157 return [
2158 r.model_copy(update={"score": norm}) for r, norm in zip(results, norms, strict=True)
2159 ]
2160 except Exception:
2161 log.debug("BM25 probe failed", exc_info=True)
2162 return []
2164 def search(
2165 self,
2166 query_vector: Vector,
2167 top_k: int | None = None,
2168 max_distance: float | None = None,
2169 query_text: str | None = None,
2170 chunk_type: ChunkType | None = None,
2171 ) -> list[SearchChunk]:
2172 """Search for similar chunks. Hybrid when FTS available, else vector-only.
2174 Results with distance > max_distance are filtered out (vector-only path).
2175 Pass max_distance=0 to disable filtering.
2176 When *chunk_type* is set, only chunks of that type ("raw" or "wiki") are returned.
2178 Raises ``EmbeddingModelMismatchError`` if the persisted ``_meta`` row was
2179 written under a different embedding model than the current ``cfg``.
2180 """
2181 if top_k is None:
2182 top_k = self._config.top_k
2183 if max_distance is None:
2184 max_distance = self._config.max_distance
2185 table = self.open_table(CHUNKS_TABLE)
2186 if table is None:
2187 return []
2188 self.initialize_meta_if_legacy()
2189 self.canonicalize_meta_if_legacy()
2190 self._ensure_embedding_compat()
2191 self._warn_stale_doc_prefix()
2193 if not self._scalar_ready:
2194 # A serve-only store never ran ingest, where scalar indexes are
2195 # built; without them the source/chunk_type prefilters full-scan.
2196 self.ensure_scalar_indexes(blocking=False)
2197 # A repair replaces indexes; the handle must see the new versions.
2198 table.checkout_latest()
2200 if query_text:
2201 hits = self._keyword_arm(
2202 table, query_text, query_vector, top_k, max_distance, chunk_type
2203 )
2204 if hits is not None:
2205 return self._with_paths(hits)
2207 try:
2208 rows = self._vector_arm(
2209 table, query_vector, top_k * self._config.candidate_multiplier, chunk_type
2210 )
2211 except Exception:
2212 # A dead scalar index fails the prefilter; the next query re-verifies.
2213 self._scalar_ready = False
2214 raise
2215 log.debug(
2216 "Vector search: query=%r, candidates=%d, max_distance=%.2f",
2217 query_text or "vector-only",
2218 len(rows),
2219 max_distance,
2220 )
2221 if rows:
2222 log.debug("Top 5 distances: %s", [r.distance for r in rows[:5]])
2223 results = self._filter_and_rerank(rows, query_vector, top_k, max_distance)
2224 return [
2225 r.model_copy(
2226 update={"score": vector_similarity(r.distance) if r.distance is not None else 0.0}
2227 )
2228 for r in self._with_paths(results)
2229 ]
2231 def _keyword_arm(
2232 self,
2233 table: LanceTable,
2234 query_text: str,
2235 query_vector: Vector,
2236 top_k: int,
2237 max_distance: float,
2238 chunk_type: ChunkType | None,
2239 ) -> list[SearchChunk] | None:
2240 """Hybrid results, or None when keyword search cannot serve this query and
2241 the caller must fall back to vector-only.
2243 Owns the keyword-search health the endpoint reports. An index that will
2244 not build and an index that fails mid-query drop every query to vector
2245 recall alike, so both stamp the flag, and a working index clears it.
2246 """
2247 if not self._fts_ready:
2248 self.ensure_fts_index(blocking=False)
2249 table.checkout_latest()
2250 if self._config.title_search and not self._title_fts_ready:
2251 self.ensure_title_fts_index(blocking=False)
2252 if not self._fts_ready:
2253 # An empty corpus has nothing to index and nothing to degrade.
2254 self._fts_degraded = self._has_chunks()
2255 return None
2256 try:
2257 hits = self._hybrid_search(
2258 table, query_text, query_vector, top_k, max_distance, chunk_type
2259 )
2260 except Exception:
2261 # Falling back changes recall characteristics for the query;
2262 # a corpus-wide FTS breakage must not present as silence.
2263 self._fts_degraded = True
2264 # A dead scalar index fails the same way; the next query re-verifies.
2265 self._scalar_ready = False
2266 log.warning("Hybrid search failed, falling back to vector-only", exc_info=True)
2267 return None
2268 self._fts_degraded = False
2269 return hits
2271 def _vector_arm(
2272 self,
2273 table: LanceTable,
2274 query_vector: Vector,
2275 limit: int,
2276 chunk_type: ChunkType | None,
2277 ) -> list[SearchChunk]:
2278 """Vector-arm candidates with the ANN recall recovery applied.
2280 When ``chunk_type`` is set, the predicate is pushed into the query so
2281 the limit applies *after* the type filter; post-filtering would
2282 silently starve wiki-only queries whose matches live past the window.
2283 """
2284 query = table.search(query_vector).distance_type(_VECTOR_METRIC)
2285 indices = _index_registry(table, self._config.lancedb_dir)
2286 if indices is None or _has_vector_index(indices):
2287 # IVF_PQ is lossy; probe more partitions and refine against full
2288 # vectors so recall stays close to the exact flat scan. An unusable
2289 # registry may hide the index, and the tuning is harmless on a flat scan.
2290 query = query.nprobes(_ann_nprobes(table.count_rows()))
2291 query = query.refine_factor(_ANN_REFINE_FACTOR)
2292 if chunk_type:
2293 query = query.where(_chunk_type_predicate(chunk_type))
2294 return _id_chunks(_stable_top(table, query, limit, _DISTANCE_COLUMN, ascending=True))
2296 def _with_paths(self, results: list[SearchChunk]) -> list[SearchChunk]:
2297 """*results* with the path of each source in place of its id, in one lookup."""
2298 paths = self.source_paths({r.source for r in results})
2299 return [
2300 r.model_copy(update={"source": paths[r.source]}) for r in results if r.source in paths
2301 ]
2303 def _fts_arm(
2304 self,
2305 table: LanceTable,
2306 query_text: str,
2307 limit: int,
2308 chunk_type: ChunkType | None,
2309 ) -> list[SearchChunk]:
2310 """BM25-arm candidates over the chunk text."""
2311 return _id_chunks(_lexical_rows(table, query_text, limit, chunk_type))
2313 def _title_arm(
2314 self,
2315 table: LanceTable,
2316 query_text: str,
2317 limit: int,
2318 chunk_type: ChunkType | None,
2319 ) -> list[SearchChunk]:
2320 """One BM25 row per document whose title matches, in title-relevance order.
2322 Every chunk of a document carries the same title, so all of its chunks
2323 tie on BM25 and a plain ``limit`` would return an arbitrary tie-ordered
2324 subset of a single document. Instead this over-fetches, collapses each
2325 source to one deterministic representative (its first chunk), and returns
2326 the top *limit* documents ordered by title score -- so "a query naming a
2327 document by title surfaces its chunks" holds as one stable row per doc.
2329 Empty when the store predates the title column or its FTS index (old
2330 indexes keep working), when the index registry is unusable, and on
2331 any query-time failure: the optional
2332 title arm must never take down the healthy chunk arm, so its failure
2333 degrades to no-titles, mirroring ``bm25_probe``.
2334 """
2335 indices = _index_registry(table, self._config.lancedb_dir)
2336 if indices is None or not _has_fts_index(indices, _TITLE_COLUMN):
2337 return []
2338 # Every chunk of one document ties on title BM25, so a fixed window can
2339 # fill up with a single long document's chunks and starve every other
2340 # title-matching document. Widen the fetch until enough distinct
2341 # documents surface, the matches run out, or the ceiling is hit.
2342 fetch = max(limit * _TITLE_FETCH_FACTOR, _TITLE_MIN_FETCH)
2343 while True:
2344 try:
2345 hits = _lexical_rows(table, query_text, fetch, chunk_type, column=_TITLE_COLUMN)
2346 except Exception:
2347 log.debug("Title arm search failed; contributing no title rows", exc_info=True)
2348 return []
2349 rows = _id_chunks(hits)
2350 best: dict[str, SearchChunk] = {}
2351 for row in rows:
2352 seen = best.get(row.source)
2353 if seen is None or row.chunk_index < seen.chunk_index:
2354 best[row.source] = row
2355 if len(best) >= limit or len(hits) < fetch or fetch >= _TITLE_FETCH_CEILING:
2356 break
2357 fetch = min(fetch * 4, _TITLE_FETCH_CEILING)
2358 ordered = sorted(best.values(), key=lambda r: (-(r.bm25_score or 0.0), r.source))
2359 return ordered[:limit]
2361 def _hybrid_search(
2362 self,
2363 table: LanceTable,
2364 query_text: str,
2365 query_vector: Vector,
2366 top_k: int,
2367 max_distance: float,
2368 chunk_type: ChunkType | None = None,
2369 ) -> list[SearchChunk]:
2370 """Multi-arm retrieval fused by weighted reciprocal rank; the fused ordering is final.
2372 A vector arm and a chunk-BM25 arm always run; a title-BM25 arm joins
2373 when ``cfg.title_search`` is on. Each row's fused score is the
2374 weight-normalized sum of its arm contributions: the vector arm has
2375 weight 1.0, the lexical arm ``cfg.lexical_fusion_weight`` (scaled per
2376 query when ``cfg.adaptive_fusion`` is on), the title arm
2377 ``cfg.title_search_weight``. So a row a single peaked arm is certain
2378 about scores that arm's share of the total weight, not a fixed 0.5.
2380 Each arm fetches exactly ``top_k`` rows. Deeper pools measurably hurt
2381 rank fusion by flooding the fused top-k with both-arm mediocrity and
2382 burying single-arm certainty (lexical identifier hits above all). No
2383 MMR runs here: lexical passages are often mutually similar, which MMR
2384 penalizes, trading relevant hits for off-topic neighbors.
2386 Title rows carry ``bm25_score``, so a title match counts as lexical
2387 support for the distance exemption like any other lexical hit.
2388 """
2389 title_rows: list[SearchChunk] = []
2390 if self._config.title_search:
2391 title_rows = self._title_arm(table, query_text, top_k, chunk_type)
2392 vector_rows = self._vector_arm(table, query_vector, top_k, chunk_type)
2393 base_lexical_weight = self._config.lexical_fusion_weight
2394 base_title_weight = self._config.title_search_weight
2395 lexical_weight = base_lexical_weight
2396 title_weight = base_title_weight
2397 if self._config.adaptive_fusion:
2398 # Quiet the lexical arms per query by vector confidence. The title
2399 # arm is lexical too, so the same factor scales it.
2400 scale = adaptive_weight_scale(vector_rows, self._config.adaptive_fusion_margin)
2401 lexical_weight = base_lexical_weight * scale
2402 title_weight = base_title_weight * scale
2403 fused = fuse_arms(
2404 vector_rows,
2405 self._fts_arm(table, query_text, top_k, chunk_type),
2406 title_rows,
2407 lexical_weight=lexical_weight,
2408 title_weight=title_weight,
2409 )
2410 fused = _drop_unsupported_far_rows(fused, max_distance)
2411 return fused[:top_k]
2413 def _filter_and_rerank(
2414 self,
2415 results: list[SearchChunk],
2416 query_vector: Vector,
2417 top_k: int,
2418 max_distance: float,
2419 ) -> list[SearchChunk]:
2420 """Apply the configured distance filter, then MMR-rerank down to top_k."""
2421 if max_distance > 0:
2422 before = len(results)
2423 if self._config.adaptive_threshold:
2424 results = self._adaptive_filter(results, top_k, max_distance)
2425 filter_name = "adaptive"
2426 else:
2427 results = self._fixed_filter(results, max_distance)
2428 filter_name = "fixed"
2429 log.debug(
2430 "After %s filter: %d/%d results, threshold=%.2f",
2431 filter_name,
2432 len(results),
2433 before,
2434 max_distance,
2435 )
2436 if len(results) > top_k:
2437 results = mmr_rerank(query_vector, results, top_k, self._config.mmr_lambda)
2438 return results
2440 def _adaptive_filter(
2441 self, results: list[SearchChunk], top_k: int, initial_threshold: float
2442 ) -> list[SearchChunk]:
2443 """Widen cosine distance threshold when too few results.
2444 Inspired by grantflow's (grantflow-ai/grantflow) adaptive retrieval
2445 pattern which widens thresholds on recursive retry. Step size and
2446 cap are configurable via ``self._config.adaptive_threshold_step``.
2448 Pre-sorts results by distance for a single-pass cutoff search.
2449 Step size is ``self._config.adaptive_threshold_step`` (default 0.2).
2450 """
2451 cap = max(initial_threshold, _MAX_THRESHOLD)
2452 step = self._config.adaptive_threshold_step
2454 sorted_results = sorted(results, key=_get_distance)
2456 threshold = initial_threshold
2457 for _ in range(_MAX_FILTER_ITERATIONS):
2458 if threshold > cap:
2459 break
2460 cutoff = _count_within_threshold(sorted_results, threshold)
2461 if cutoff >= top_k:
2462 return sorted_results[:cutoff]
2463 threshold += step
2464 # Final pass at cap
2465 cutoff = _count_within_threshold(sorted_results, cap)
2466 return sorted_results[:cutoff]
2468 def _fixed_filter(self, results: list[SearchChunk], threshold: float) -> list[SearchChunk]:
2469 """Simple fixed threshold filter - keep only results within distance threshold."""
2470 return [r for r in results if _get_distance(r) <= threshold]
2472 def add_entities(self, records: list[dict]) -> int:
2473 """Append typed entity rows; creates the table on first write.
2475 Additive to the store: existing tables and schemas are untouched, and
2476 stores without this table behave as if nothing was ever extracted.
2477 """
2478 if not records:
2479 return 0
2480 from lilbee.retrieval.entities.schema import _entities_schema
2482 with self._write_lock():
2483 ids = self._ensure_ids_unlocked(_documents(record["source"] for record in records))
2484 table = ensure_table(self.get_db(), ENTITIES_TABLE, _entities_schema())
2485 table.add(_keyed(records, "source", ids))
2486 return len(records)
2488 def entity_schema_state(self) -> EntitySchemaState | None:
2489 """The persisted entity schema row, or ``None`` when never induced.
2491 The schema is machine state induced from the corpus and lives inside
2492 the index, so it travels with the data.
2493 """
2494 table = self.open_table(ENTITY_SCHEMA_TABLE)
2495 if table is None:
2496 return None
2497 rows = table.search().limit(None).to_list()
2498 if not rows:
2499 return None
2500 # One row by contract; take the newest if a rewrite ever left a stale one.
2501 row = max(rows, key=lambda r: r["updated_at"])
2502 return EntitySchemaState(
2503 schema_json=str(row["schema_json"]),
2504 applied=bool(row["applied"]),
2505 source_count=int(row["source_count"]),
2506 updated_at=str(row["updated_at"]),
2507 )
2509 def save_entity_schema(self, schema_json: str, *, applied: bool, source_count: int) -> None:
2510 """Overwrite the single persisted entity schema row."""
2511 with self._write_lock():
2512 db = self.get_db()
2513 table = ensure_table(db, ENTITY_SCHEMA_TABLE, _entity_schema_state_schema())
2514 _safe_delete_unlocked(table, ENTITY_SCHEMA_DELETE_ALL_PREDICATE)
2515 table.add(
2516 [
2517 {
2518 "schema_json": schema_json,
2519 "applied": applied,
2520 "source_count": source_count,
2521 "updated_at": datetime.now(UTC).isoformat(),
2522 }
2523 ]
2524 )
2526 def mark_entity_schema_applied(self) -> None:
2527 """Record that a full extraction pass completed under the stored schema."""
2528 state = self.entity_schema_state()
2529 if state is None:
2530 return
2531 self.save_entity_schema(
2532 state["schema_json"], applied=True, source_count=state["source_count"]
2533 )
2535 def entity_value_counts(self, entity_type: str) -> tuple[int, int]:
2536 """(mentions, distinct normalized values) for one entity type.
2538 Full scan by design: a count is a corpus property. Streaming batches
2539 keep memory flat at any corpus size.
2540 """
2541 table = self.open_table(ENTITIES_TABLE)
2542 if table is None:
2543 return 0, 0
2544 mentions = 0
2545 values: set[str] = set()
2546 arrow = table.to_arrow().select(["type", "normalized_value"])
2547 for batch in arrow.to_batches(max_chunksize=_TERM_SCAN_BATCH_ROWS):
2548 types = batch.column("type").to_pylist()
2549 vals = batch.column("normalized_value").to_pylist()
2550 for t_, v in zip(types, vals, strict=True):
2551 if t_ == entity_type:
2552 mentions += 1
2553 values.add(v)
2554 return mentions, len(values)
2556 def entity_association_counts(self, counted: str, grouped_by: str) -> dict[str, int]:
2557 """Distinct *counted*-type values co-occurring with each *grouped_by* value.
2559 Co-occurrence is per chunk: two entities extracted from the same
2560 ``(source, chunk_index)`` are associated. This is the GROUP BY that
2561 answers "how many X is each Y associated with".
2562 """
2563 table = self.open_table(ENTITIES_TABLE)
2564 if table is None:
2565 return {}
2566 per_chunk: dict[tuple[str, int], tuple[set[str], set[str]]] = {}
2567 arrow = table.to_arrow().select(["type", "normalized_value", _SOURCE_ID, _CHUNK_INDEX])
2568 for batch in arrow.to_batches(max_chunksize=_TERM_SCAN_BATCH_ROWS):
2569 rows = zip(
2570 batch.column("type").to_pylist(),
2571 batch.column("normalized_value").to_pylist(),
2572 batch.column(_SOURCE_ID).to_pylist(),
2573 batch.column(_CHUNK_INDEX).to_pylist(),
2574 strict=True,
2575 )
2576 for t_, v, src, idx in rows:
2577 if t_ not in (counted, grouped_by):
2578 continue
2579 counted_vals, group_vals = per_chunk.setdefault((src, idx), (set(), set()))
2580 (counted_vals if t_ == counted else group_vals).add(v)
2581 associations: dict[str, set[str]] = {}
2582 for counted_vals, group_vals in per_chunk.values():
2583 for group_value in group_vals:
2584 associations.setdefault(group_value, set()).update(counted_vals)
2585 return {k: len(v) for k, v in sorted(associations.items())}
2587 def count_term_mentions(self, term: str) -> tuple[int, int]:
2588 """(matching chunks, distinct matching sources) for a case-insensitive
2589 substring scan of the WHOLE chunks table.
2591 This is deliberately a full scan, not a top-k search: a count is a
2592 corpus property, and any retrieval shortcut undercounts it. Streaming
2593 Arrow batches keeps the working set to one batch of text at a time,
2594 so cost is linear in corpus size and memory stays flat.
2595 """
2596 table = self.open_table(CHUNKS_TABLE)
2597 if table is None:
2598 return 0, 0
2599 needle = term.lower()
2600 chunk_hits = 0
2601 sources: set[str] = set()
2602 arrow = table.to_arrow().select([_SOURCE_ID, "chunk"])
2603 for batch in arrow.to_batches(max_chunksize=_TERM_SCAN_BATCH_ROWS):
2604 texts = batch.column("chunk").to_pylist()
2605 names = batch.column(_SOURCE_ID).to_pylist()
2606 for name, text in zip(names, texts, strict=True):
2607 if text and needle in text.lower():
2608 chunk_hits += 1
2609 sources.add(name)
2610 return chunk_hits, len(sources)
2612 def chunk_batches(
2613 self, columns: list[str], batch_rows: int
2614 ) -> Iterator[list[Mapping[str, Any]]]:
2615 """Every chunk row in batches of *batch_rows*, for *columns*; ``source`` is the path."""
2616 table = self.open_table(CHUNKS_TABLE)
2617 if table is None:
2618 return
2619 paths = self._all_source_paths()
2620 read = [_SOURCE_ID if column == "source" else column for column in columns]
2621 arrow = table.search().select(read).limit(None).to_arrow()
2622 for batch in arrow.to_batches(max_chunksize=batch_rows):
2623 rows: list[Mapping[str, Any]] = list(_pathed(batch.to_pylist(), "source", paths))
2624 yield rows
2626 def _all_source_paths(self) -> dict[str, str]:
2627 """The path of every source row, by id."""
2628 table = self.open_table(SOURCES_TABLE)
2629 if table is None:
2630 return {}
2631 rows = table.search().select([_SOURCE_ID, _PATH]).limit(None).to_list()
2632 name = self._locations().name
2633 return {row[_SOURCE_ID]: name(row[_PATH]) for row in rows}
2635 def count_chunks(self) -> int:
2636 """Total chunks in the store."""
2637 table = self.open_table(CHUNKS_TABLE)
2638 return table.count_rows() if table is not None else 0
2640 def get_chunks_by_source(self, source: str) -> list[SearchChunk]:
2641 """Return every chunk whose ``source`` equals *source*.
2643 The database does the filtering, so only the matching rows are read.
2644 A query failure raises rather than falling back to a whole-table scan:
2645 a document's chunks are a bounded read, and the scan that would rescue
2646 it costs the entire index, vectors included, in memory.
2647 """
2648 return self._chunks_of(source, "1 = 1")
2650 def _chunks_of(self, source: str, predicate: str) -> list[SearchChunk]:
2651 """*source*'s chunk rows that *predicate* also matches."""
2652 table = self.open_table(CHUNKS_TABLE)
2653 source_id = self.source_ids([source], None).get(source)
2654 if table is None or source_id is None:
2655 return []
2656 query = table.search().where(f"{_SOURCE_ID} = '{source_id}' AND {predicate}")
2657 rows = _pathed(query.limit(None).to_list(), "source", {source_id: source})
2658 return [SearchChunk(**row) for row in rows]
2660 def get_chunks_by_indices(self, source: str, indices: Sequence[int]) -> list[SearchChunk]:
2661 """Return *source*'s chunks whose ``chunk_index`` is in *indices*.
2663 Rows come back in ``chunk_index`` order; indices past either end of
2664 the document are simply absent from the result. Filtering happens in
2665 the database for the same reason as :meth:`get_chunks_by_source`:
2666 neighbor expansion runs once per hit source per query, so a
2667 whole-table rescue would spike memory on the hottest path there is.
2668 """
2669 if not indices:
2670 return []
2671 wanted = ", ".join(str(int(i)) for i in indices)
2672 rows = self._chunks_of(source, f"chunk_index IN ({wanted})")
2673 return sorted(rows, key=lambda c: c.chunk_index)
2675 def _delete_rows_of_unlocked(self, source_ids: Iterable[str]) -> None:
2676 """Delete the rows the sources *source_ids* hold in every per-source table.
2678 Caller must hold ``write_lock()``. One ``IN`` delete per table covers
2679 every source, so a batched flush pays a constant number of predicate
2680 deletes instead of one set per document. A delete failure propagates:
2681 swallowed, it would leave every flushed file silently stale; raised,
2682 the flush fails and the files replan on the next sync.
2683 """
2684 quoted = _in_list(source_ids)
2685 if not quoted:
2686 return
2687 for name in _PER_SOURCE_TABLES:
2688 table = self.open_table(name)
2689 if table is not None:
2690 table.delete(f"{_SOURCE_ID} IN ({quoted})")
2692 def delete_by_source(self, source: str, source_type: SourceType) -> None:
2693 """Delete the chunks and page texts of *source* of *source_type*."""
2694 with self._write_lock():
2695 self._delete_rows_of_unlocked(self.source_ids([source], source_type).values())
2696 self._invalidate_source_cache()
2698 def add_page_texts(self, records: list[dict]) -> int:
2699 """Add per-page text rows (no vectors). Returns count added."""
2700 if not records:
2701 return 0
2702 with self._write_lock():
2703 ids = self._ensure_ids_unlocked(_documents(record["source"] for record in records))
2704 table = ensure_table(self.get_db(), PAGE_TEXTS_TABLE, _page_texts_schema())
2705 table.add(_keyed(records, "source", ids))
2706 return len(records)
2708 def get_page_texts(self, source: str | None = None) -> list[PageTextRecord]:
2709 """Return per-page text rows, all or for a single *source*."""
2710 arrow = self.page_texts_arrow(source)
2711 return cast("list[PageTextRecord]", arrow.to_pylist())
2713 def page_texts_arrow(self, source: str | None = None) -> pa.Table:
2714 """Return per-page text rows as an Arrow table in a single scan.
2716 The columnar sibling of :meth:`get_page_texts`: the export path keeps the
2717 whole set in Arrow (no per-row Python objects) from read through file
2718 write. Empty with the canonical schema when the table or *source* is empty.
2719 """
2721 table = self.open_table(PAGE_TEXTS_TABLE)
2722 source_id = None
2723 if source is not None:
2724 source_id = self.source_ids([source], SourceType.DOCUMENT).get(source, "")
2725 if table is None:
2726 arrow = _page_texts_schema().empty_table()
2727 else:
2728 query = table.search().select(_page_texts_schema().names)
2729 if source_id is not None:
2730 query = query.where(f"{_SOURCE_ID} = '{escape_sql_string(source_id)}'")
2731 arrow = query.limit(None).to_arrow()
2732 ids = arrow.column(_SOURCE_ID).to_pylist()
2733 paths = self.source_paths(ids)
2734 shown = pa.array([paths.get(key) for key in ids], pa.utf8())
2735 return arrow.set_column(0, "source", shown).filter(pc.is_valid(shown))
2737 def sources_arrow(self) -> pa.Table:
2738 """Return each tracked source's extraction metadata as an Arrow table.
2740 The columnar sibling of :meth:`get_sources`, keyed by ``source`` so it
2741 joins straight onto the page-text table. ``get_sources`` builds a dict per
2742 source, which on a corpus of millions of single-page documents costs more
2743 than the text being exported.
2744 """
2745 import pyarrow as pa
2747 table = self.open_table(SOURCES_TABLE)
2748 columns = ["source", *SourceMeta._fields]
2749 if table is None:
2750 return pa.schema([pa.field(name, pa.utf8()) for name in columns]).empty_table()
2751 query = table.search().where(_LISTED).select([_PATH, *SourceMeta._fields])
2752 arrow = query.limit(None).to_arrow().rename_columns(columns)
2753 name = self._locations().name
2754 shown = pa.array([name(path) for path in arrow.column("source").to_pylist()], pa.utf8())
2755 arrow = arrow.set_column(0, "source", shown)
2756 if pc.count_distinct(arrow.column("source")).as_py() == arrow.num_rows:
2757 return arrow
2758 # One row per source, so a caller joining on it cannot fan out. A doubled
2759 # row would otherwise multiply every page it owns. Last wins; single-threaded
2760 # because that is the only execution mode with an ordered aggregate.
2761 grouped = arrow.group_by("source", use_threads=False).aggregate(
2762 [(name, "last") for name in SourceMeta._fields]
2763 )
2764 renamed = grouped.rename_columns(
2765 [name.removesuffix("_last") for name in grouped.schema.names]
2766 )
2767 return renamed.select(columns)
2769 def wiki_chunk_sources(self) -> set[str]:
2770 """Return the distinct sources of the chunk rows written by the wiki layer."""
2771 table = self.open_table(CHUNKS_TABLE)
2772 if table is None:
2773 return set()
2774 rows = table.search().where(_WIKI_CHUNK).select([_SOURCE_ID]).limit(None).to_list()
2775 return set(self.source_paths(row[_SOURCE_ID] for row in rows).values())
2777 def wiki_pages_citing(self, paths: Iterable[str]) -> set[str]:
2778 """The wiki pages whose citations cite a source at or below one of *paths*."""
2779 table = self.open_table(CITATIONS_TABLE)
2780 roots = set(paths)
2781 if table is None or not roots:
2782 return set()
2783 rows = table.search().select([_SOURCE_ID, "wiki_source"]).limit(None).to_list()
2784 cited = self.source_paths(row[_SOURCE_ID] for row in rows)
2785 return {
2786 row["wiki_source"]
2787 for row in rows
2788 if row[_SOURCE_ID] in cited and _at_or_below(cited[row[_SOURCE_ID]], roots)
2789 }
2791 def wiki_citations_by_page(self) -> dict[str, list[CitationRecord]]:
2792 """Every citation, naming its source by path, under the wiki page that holds it."""
2793 table = self.open_table(CITATIONS_TABLE)
2794 if table is None:
2795 return {}
2796 by_page: dict[str, list[CitationRecord]] = {}
2797 for record in self._shown_citations(table.search().limit(None).to_list()):
2798 by_page.setdefault(record["wiki_source"], []).append(record)
2799 return by_page
2801 def wiki_citation_sources(self) -> set[str]:
2802 """Return the distinct wiki_source values present in the citations table."""
2803 table = self.open_table(CITATIONS_TABLE)
2804 if table is None:
2805 return set()
2806 rows = table.search().select(["wiki_source"]).limit(None).to_list()
2807 return {row["wiki_source"] for row in rows}
2809 def get_sources(
2810 self,
2811 *,
2812 search: str | None = None,
2813 limit: int | None = None,
2814 offset: int = 0,
2815 ) -> list[SourceRecord]:
2816 """Return source records, filtered by *search* and sliced by offset/limit."""
2817 table = self.open_table(SOURCES_TABLE)
2818 if table is None:
2819 return []
2820 locations = self._locations()
2821 query = table.search().where(_listed_matching(search, locations))
2822 if offset:
2823 query = query.offset(offset)
2824 return [_source_record(row, locations) for row in query.limit(limit).to_list()]
2826 def source_summary(self, newest: int) -> SourceSummary:
2827 """Totals over every listed source and the *newest* most recently ingested, in one pass.
2829 Reads five columns in batches and keeps only running totals and the
2830 current newest rows, so memory stays flat however many sources exist.
2831 """
2832 total_chunks = 0
2833 documents = 0
2834 kinds: Counter[str] = Counter()
2835 latest: list[SourceHead] = []
2836 sort_keys = [("ingested_at", "descending"), ("filename", "ascending")]
2837 for table in self._iter_source_batches(_SUMMARY_COLUMNS):
2838 documents += table.num_rows
2839 total_chunks += pc.sum(table.column("chunk_count")).as_py() or 0
2840 for entry in pc.value_counts(table.column("source_type")).to_pylist():
2841 kinds[entry["values"]] += entry["counts"]
2842 top = table.take(pc.select_k_unstable(table, newest, sort_keys))
2843 latest.extend(_source_heads(top))
2844 latest.sort(key=lambda head: head["filename"])
2845 latest.sort(key=lambda head: head["ingested_at"], reverse=True)
2846 del latest[newest:]
2847 return SourceSummary(documents, total_chunks, dict(kinds), latest)
2849 def iter_source_heads(self) -> Iterator[SourceHead]:
2850 """Every listed source, streamed in batches so no more than one batch is held."""
2851 for table in self._iter_source_batches(_HEAD_COLUMNS):
2852 yield from _source_heads(table)
2854 def _iter_source_batches(self, columns: list[str]) -> Iterator[pa.Table]:
2855 """The listed source rows in batches, each row's shown name as ``filename``."""
2856 table = self.open_table(SOURCES_TABLE)
2857 if table is None:
2858 return
2859 name = self._locations().name
2860 read = [_PATH if column == _OLD_PATH else column for column in columns]
2861 query = table.search().where(_LISTED).select(read).limit(None)
2862 for batch in query.to_batches(_SOURCE_SCAN_BATCH_ROWS):
2863 rows = pa.Table.from_batches([batch])
2864 at = rows.schema.get_field_index(_PATH)
2865 shown = pa.array([name(path) for path in rows.column(at).to_pylist()], pa.utf8())
2866 yield rows.set_column(at, _OLD_PATH, shown)
2868 def get_sources_where(self, keep: Callable[[str], bool]) -> list[SourceRecord]:
2869 """Source records whose filename *keep* accepts, read in batches.
2871 A row *keep* refuses never becomes a Python object, so a caller that wants
2872 one slice of a large table holds only that slice.
2873 """
2874 import pyarrow as pa
2876 table = self.open_table(SOURCES_TABLE)
2877 if table is None:
2878 return []
2879 kept: list[SourceRecord] = []
2880 locations = self._locations()
2881 query = table.search().where(_LISTED).limit(None)
2882 for batch in query.to_batches(_SOURCE_SCAN_BATCH_ROWS):
2883 names = [locations.name(path) for path in batch.column(_PATH).to_pylist()]
2884 mask = pa.array([keep(name) for name in names])
2885 kept.extend(_source_record(row, locations) for row in batch.filter(mask).to_pylist())
2886 return kept
2888 def missing_sources(self, names: Iterable[str]) -> set[str]:
2889 """The *names* that hold no source row, from one pass over the name column."""
2890 missing = set(names)
2891 table = self.open_table(SOURCES_TABLE)
2892 if table is None or not missing:
2893 return missing
2894 name = self._locations().name
2895 scan = table.search().where(_LISTED).select([_PATH]).limit(None)
2896 for batch in scan.to_batches(_SOURCE_SCAN_BATCH_ROWS):
2897 missing.difference_update(name(path) for path in batch.column(_PATH).to_pylist())
2898 return missing
2900 def sources_version(self) -> int | None:
2901 """The version of the source table now, or None when there is no table."""
2902 table = self.open_table(SOURCES_TABLE)
2903 return None if table is None else int(table.version)
2905 def sources_by_hash(self, hashes: Sequence[str], *, version: int) -> dict[str, list[str]]:
2906 """Document sources at *version* of the source table, by content hash, for *hashes*.
2908 A failed file counts at the hash it failed at.
2909 """
2910 found: dict[str, list[str]] = {}
2911 table = self.open_table(SOURCES_TABLE)
2912 if table is None:
2913 return found
2914 table.checkout(version)
2915 name = self._locations().name
2916 documents = f"{_SOURCE_TYPE} != '{SourceType.IMPORTED}' AND {_LISTED}"
2917 failed = f"{_HELD_OUT} = '{HeldOut.FAILED}'"
2918 for start in range(0, len(hashes), _SOURCE_STAT_BATCH_ROWS):
2919 quoted = _in_list(hashes[start : start + _SOURCE_STAT_BATCH_ROWS])
2920 for column, where in (("file_hash", documents), (_HELD_HASH, failed)):
2921 query = table.search().where(f"{column} IN ({quoted}) AND {where}")
2922 for row in query.select([_PATH, column]).limit(None).to_list():
2923 found.setdefault(row[column], []).append(name(row[_PATH]))
2924 return found
2926 def sources_checkpoint(self) -> SourcesCheckpoint:
2927 """The time and the version of the source table, read under the write lock.
2929 A source row is stamped under that lock, so no row stamped before the
2930 time is committed later.
2931 """
2932 with self._write_lock():
2933 return SourcesCheckpoint(datetime.now(UTC).isoformat(), self.sources_version())
2935 def wiki_refreshed(self) -> SourcesCheckpoint | None:
2936 """The checkpoint up to which the wiki's page index covers the source rows.
2938 None with no record or an unreadable one: every source row is owed.
2939 An unreadable record is reported and deleted, so the next read is silent.
2940 """
2941 record = self._config.lancedb_dir / _WIKI_REFRESHED_NAME
2942 try:
2943 fields = json.loads(record.read_text(encoding="utf-8"))
2944 version = fields["version"]
2945 return SourcesCheckpoint(str(fields["time"]), None if version is None else int(version))
2946 except FileNotFoundError:
2947 return None
2948 except (OSError, ValueError, KeyError, TypeError) as exc:
2949 log.warning(_WIKI_REFRESHED_UNREADABLE, record, exc)
2950 # A record that cannot be deleted stays, and the next read reports it again.
2951 with suppress(OSError):
2952 self.forget_wiki_refreshed()
2953 return None
2955 def record_wiki_refreshed(self, checkpoint: SourcesCheckpoint) -> None:
2956 """Record that the wiki's page index covers the source rows up to *checkpoint*.
2958 A write that fails leaves the record as it was.
2959 """
2960 record = self._config.lancedb_dir / _WIKI_REFRESHED_NAME
2961 write_private_text(record, json.dumps(checkpoint._asdict()))
2963 def forget_wiki_refreshed(self) -> None:
2964 """Delete the record: the wiki's page index covers no source row."""
2965 (self._config.lancedb_dir / _WIKI_REFRESHED_NAME).unlink(missing_ok=True)
2967 def sources_ingested_since(self, since: str) -> list[str]:
2968 """The filename of each source row written at or after *since*, an ISO time in UTC."""
2969 table = self.open_table(SOURCES_TABLE)
2970 if table is None:
2971 return []
2972 written = f"ingested_at >= '{escape_sql_string(since)}' AND {_LISTED}"
2973 query = table.search().where(written).select([_PATH])
2974 name = self._locations().name
2975 return [name(row[_PATH]) for row in query.limit(None).to_list()]
2977 def count_sources(self, *, search: str | None = None) -> int:
2978 """Count tracked sources matching *search* without materializing rows."""
2979 table = self.open_table(SOURCES_TABLE)
2980 if table is None:
2981 return 0
2982 count: int = table.count_rows(filter=_listed_matching(search, self._locations()))
2983 return count
2985 def _source_row(
2986 self,
2987 file_hash: str,
2988 chunk_count: int,
2989 source_type: SourceType,
2990 stat: SourceStat | None = None,
2991 meta: SourceMeta | None = None,
2992 ) -> dict:
2993 """The fields of one ``_sources`` row but its id and path; no stat reads as unknown.
2995 Absent extraction metadata persists as NULL, matching rows written
2996 before the metadata columns existed.
2997 """
2998 meta = meta or SourceMeta()
2999 return {
3000 "file_hash": file_hash,
3001 # Stamped under the write lock that commits the row: sources_checkpoint relies on it.
3002 "ingested_at": datetime.now(UTC).isoformat(),
3003 "chunk_count": chunk_count,
3004 _SOURCE_TYPE: source_type,
3005 "size_bytes": stat.size_bytes if stat else SOURCE_STAT_UNKNOWN,
3006 "mtime_ns": stat.mtime_ns if stat else SOURCE_STAT_UNKNOWN,
3007 "stat_captured_ns": stat.captured_ns if stat else SOURCE_STAT_UNKNOWN,
3008 "title": meta.title or None,
3009 "authors": meta.authors or None,
3010 "created_at": meta.created_at or None,
3011 **dict.fromkeys(_HELD_COLUMNS),
3012 }
3014 def _placeholder_row(self, source_type: SourceType) -> dict:
3015 """The fields of a row that only names a source: a wiki page, or a source held out."""
3016 row = self._source_row("", 0, source_type)
3017 return row if _is_wiki(source_type) else {**row, _HELD_OUT: HeldOut.REMOVED}
3019 def _put_sources_unlocked(
3020 self, sources: LanceTable, writes: Iterable[_SourceWrite], *, keep_documents: bool = False
3021 ) -> dict[str, str]:
3022 """Put each of *writes* in *sources* in one merge on id; returns each name's id. Lock held.
3024 A write keeps the id of the row of its type class it finds at its key or the
3025 other form of its name, else takes its new id or a random one, and its row is
3026 stored at its key. With *keep_documents* a removal leaves a row that holds a
3027 document as it is.
3028 """
3029 writes = list(
3030 {(write.name, _is_wiki(write.source_type)): write for write in writes}.values()
3031 )
3032 locations = self._locations()
3033 found = _found_rows(sources, writes, locations) if writes else {}
3034 ids: dict[str, str] = {}
3035 rows: dict[str, dict[str, Any]] = {}
3036 for write in writes:
3037 held = found.get(write.name)
3038 ids[write.name] = held[_SOURCE_ID] if held else write.new_id or _new_source_id()
3039 row = self._written_row(write, held, keep_documents=keep_documents)
3040 if row is not None:
3041 path = locations.key(write.name, write.source_type)
3042 rows[ids[write.name]] = {**row, _SOURCE_ID: ids[write.name], _PATH: path}
3043 if rows:
3044 merge = sources.merge_insert(_SOURCE_ID).when_matched_update_all()
3045 merge = merge.when_not_matched_insert_all()
3046 merge.execute(pa.Table.from_pylist(list(rows.values()), schema=sources.schema))
3047 return ids
3049 def _written_row(
3050 self, write: _SourceWrite, found: dict[str, Any] | None, *, keep_documents: bool
3051 ) -> dict[str, Any] | None:
3052 """The fields *write* puts over the row it *found*; None when it changes nothing."""
3053 if write.record is None and write.fields is None and found is not None:
3054 return None
3055 base = write.fields if write.fields is not None else found
3056 row = dict(base) if base is not None else self._placeholder_row(write.source_type)
3057 if write.record is None:
3058 return row
3059 holds = found is not None and _holds_document(found)
3060 if keep_documents and holds and write.record.kind is HeldOut.REMOVED:
3061 return None
3062 return _held_row(row, write.record, holds_document=holds)
3064 def _sources_table(self) -> LanceTable:
3065 """Open or create ``_sources``.
3067 A table that names sources by id without a ``_sources`` beside it is what a
3068 rebuild that stopped after its first drop leaves: no path names its rows,
3069 so it goes before the new ``_sources`` comes.
3070 """
3071 db = self.get_db()
3072 names = table_names(db)
3073 if SOURCES_TABLE not in names:
3074 for name in (*_PER_SOURCE_TABLES, CITATIONS_TABLE):
3075 if name in names:
3076 db.drop_table(name)
3077 return ensure_table(db, SOURCES_TABLE, _sources_schema())
3079 def _lookup(self, column: str, values: Iterable[str], other: str) -> dict[str, str]:
3080 """``{value: other}`` for the source rows whose *column* holds one of *values*."""
3081 return {row[column]: row[other] for row in self._rows_where(column, values, [other])}
3083 def _rows_where(
3084 self, column: str, values: Iterable[str], others: list[str]
3085 ) -> Iterator[dict[str, Any]]:
3086 """*column* and *others* of each source row whose *column* holds one of *values*."""
3087 table = self.open_table(SOURCES_TABLE)
3088 wanted = sorted(set(values))
3089 if table is None:
3090 return
3091 for start in range(0, len(wanted), _SOURCE_STAT_BATCH_ROWS):
3092 batch = _in_list(wanted[start : start + _SOURCE_STAT_BATCH_ROWS])
3093 query = table.search().where(f"{column} IN ({batch})").select([column, *others])
3094 yield from query.limit(None).to_list()
3096 def source_ids(self, paths: Iterable[str], source_type: SourceType | None) -> dict[str, str]:
3097 """The id of each of the names *paths* that a source row of *source_type* holds.
3099 A name finds its source at the key of each type of the class of *source_type*
3100 (a wiki page, or not), the class a write keeps one row per name in. With
3101 *source_type* None, a name finds a wiki page first.
3102 """
3103 names = list(paths)
3104 types = [
3105 kind
3106 for kind in _READ_ORDER
3107 if source_type is None or _is_wiki(kind) == _is_wiki(source_type)
3108 ]
3109 key = self._locations().key
3110 wanted = [((key(name, kind), kind), name) for kind in types for name in names]
3111 rows = self._rows_where(_PATH, {at[0] for at, _ in wanted}, [_SOURCE_TYPE, _SOURCE_ID])
3112 held = {(row[_PATH], row[_SOURCE_TYPE]): row[_SOURCE_ID] for row in rows}
3113 found: dict[str, str] = {}
3114 for at, name in wanted:
3115 if at in held:
3116 found.setdefault(name, held[at])
3117 return found
3119 def source_paths(self, source_ids: Iterable[str]) -> dict[str, str]:
3120 """The shown name of each of *source_ids* that a source row holds."""
3121 name = self._locations().name
3122 found = self._lookup(_SOURCE_ID, source_ids, _PATH)
3123 return {source_id: name(path) for source_id, path in found.items()}
3125 def web_titles(self, names: Iterable[str]) -> dict[str, str]:
3126 """The stored title of each crawled page among *names* that has one."""
3127 key = self._locations().key
3128 wanted = {key(name, SourceType.DOCUMENT): name for name in names if is_web_source(name)}
3129 if not wanted:
3130 return {}
3131 rows = self._rows_where(_PATH, wanted, [_SOURCE_TYPE, _TITLE_COLUMN])
3132 return {
3133 wanted[row[_PATH]]: row[_TITLE_COLUMN]
3134 for row in rows
3135 if row[_SOURCE_TYPE] == SourceType.DOCUMENT and row[_TITLE_COLUMN]
3136 }
3138 def keyed_source_ids(self, paths: Iterable[str]) -> dict[str, str]:
3139 """The id the keying of an older library gave each of *paths*, when a row holds it."""
3140 keyed = {path: _md5(path) for path in paths}
3141 held = self.source_paths(keyed.values())
3142 return {path: source_id for path, source_id in keyed.items() if source_id in held}
3144 def _ensure_ids_unlocked(self, sources: Iterable[tuple[str, SourceType]]) -> dict[str, str]:
3145 """The id of each named source of its type, adding a placeholder row where none is.
3147 Caller holds the lock.
3148 """
3149 writes = [_SourceWrite(name, source_type) for name, source_type in sources]
3150 return self._put_sources_unlocked(self._sources_table(), writes) if writes else {}
3152 def upsert_source(
3153 self,
3154 filename: str,
3155 file_hash: str,
3156 chunk_count: int,
3157 source_type: SourceType = SourceType.DOCUMENT,
3158 stat: SourceStat | None = None,
3159 meta: SourceMeta | None = None,
3160 ) -> None:
3161 """Add or update a source tracking record."""
3162 with self._write_lock():
3163 fields = self._source_row(file_hash, chunk_count, source_type, stat, meta)
3164 write = _SourceWrite(filename, source_type, fields)
3165 self._put_sources_unlocked(self._sources_table(), [write])
3166 self._invalidate_source_cache()
3168 def update_source_stats(self, backfills: list[SourceStatBackfill]) -> None:
3169 """Record size/mtime for already-tracked sources in batched locked writes."""
3170 if not backfills:
3171 return
3172 for start in range(0, len(backfills), _SOURCE_STAT_BATCH_ROWS):
3173 batch = backfills[start : start + _SOURCE_STAT_BATCH_ROWS]
3174 with self._write_lock():
3175 writes = [_stat_write(backfill) for backfill in batch]
3176 self._put_sources_unlocked(self._sources_table(), writes)
3177 self._invalidate_source_cache()
3179 def optimize_sources(self) -> None:
3180 """Compact the sources table; per-flush upserts otherwise accrete tiny versions."""
3181 with self._write_lock():
3182 table = self.open_table(SOURCES_TABLE)
3183 if table is None:
3184 return
3185 try:
3186 table.optimize()
3187 except Exception:
3188 log.debug("Sources table optimize failed", exc_info=True)
3190 def write_chunks_batch(
3191 self, items: list[ChunkWrite], present: Callable[[str], bool] | None = None
3192 ) -> int:
3193 """Write several documents' chunks in one locked transaction. Returns chunks added.
3195 One ``write_lock`` acquisition covers the batch's cleanup deletes, page
3196 texts, chunk add, and source upserts, so a reader never observes a
3197 half-applied batch. Page texts land after the cleanup and before the
3198 source rows, so a page-text failure leaves the rows stale and the files
3199 replan next sync; a document with no chunks still persists its page
3200 texts and source row. The embedding-identity gate and per-vector
3201 dimension check mirror ``add_chunks``; a dimension mismatch raises and
3202 the whole batch is rejected. Under the lock, an item whose source
3203 *present* says is gone is left out: a crawl that retired the file after
3204 the sync read it removes its source under this lock, so the source never
3205 comes back.
3206 """
3207 if not items:
3208 return 0
3209 with self._write_lock(timeout=BATCH_LOCK_TIMEOUT):
3210 if present is not None:
3211 items = [it for it in items if present(it.source)]
3212 if not items:
3213 return 0
3214 embedding_model = self._config.embedding_model
3215 embedding_dim = self._config.embedding_dim
3216 self._ensure_embedding_compat()
3217 self._fts_ready = False
3218 self._scalar_ready = False
3219 all_records = [rec for it in items for rec in it.records]
3220 _check_vector_dims(all_records, embedding_dim)
3221 db = self.get_db()
3222 self._refresh_key_indexes_unlocked(items, len(all_records))
3223 # A new source's row exists before its chunks, so a flush that stops
3224 # between them leaves rows the next flush of the file finds and deletes.
3225 ids = self._ensure_ids_unlocked((it.source, it.source_type) for it in items)
3226 self._delete_rows_of_unlocked([ids[it.source] for it in items if it.needs_cleanup])
3227 page_rows = [row for it in items for row in (it.page_texts or [])]
3228 if page_rows:
3229 table = ensure_table(db, PAGE_TEXTS_TABLE, _page_texts_schema())
3230 table.add(_keyed(page_rows, "source", ids))
3231 keyed = _keyed(all_records, "source", ids)
3232 self._add_chunk_records_unlocked(keyed, embedding_model, embedding_dim)
3233 self._put_sources_unlocked(self._sources_table(), self._batch_source_writes(items))
3234 self._invalidate_source_cache()
3235 return len(all_records)
3237 def _refresh_key_indexes_unlocked(self, items: list[ChunkWrite], chunk_rows: int) -> None:
3238 """Keep the key index of each table this flush appends to current. Caller holds the lock."""
3239 flush_rows = {
3240 CHUNKS_TABLE: chunk_rows,
3241 PAGE_TEXTS_TABLE: sum(len(it.page_texts or []) for it in items),
3242 SOURCES_TABLE: len(items),
3243 }
3244 for name, column in _FLUSH_KEY_COLUMNS:
3245 table = self.open_table(name)
3246 if table is None or not flush_rows[name]:
3247 continue
3248 try:
3249 _refresh_key_index_unlocked(table, column, flush_rows[name])
3250 except Exception:
3251 self._report_key_index_failure(table, name, column)
3253 def _report_key_index_failure(self, table: LanceTable, name: str, column: str) -> None:
3254 """Warn once for a populated table whose key index cannot be refreshed, then at debug.
3256 The delete stays correct and scans the table, so every flush costs more
3257 as the table grows.
3258 """
3259 first = name not in self._key_index_warned and _holds_rows(table)
3260 if first:
3261 self._key_index_warned.add(name)
3262 log.log(
3263 logging.WARNING if first else logging.DEBUG,
3264 "Key index refresh failed on '%s.%s'; each write scans the table until it succeeds",
3265 name,
3266 column,
3267 exc_info=True,
3268 )
3270 def _add_chunk_records_unlocked(
3271 self,
3272 all_records: list[dict],
3273 embedding_model: str,
3274 embedding_dim: int,
3275 ) -> None:
3276 """Add the batch's chunk rows, writing meta on first use. Caller holds ``write_lock()``."""
3277 if not all_records:
3278 return
3279 self._chunks_table().add(all_records)
3280 if self.get_meta() is None:
3281 self._write_meta_unlocked(embedding_model=embedding_model, embedding_dim=embedding_dim)
3283 def _batch_source_writes(self, items: list[ChunkWrite]) -> list[_SourceWrite]:
3284 """One ``_sources`` row per batched document, of the type its item states."""
3285 return [
3286 _SourceWrite(
3287 it.source,
3288 it.source_type,
3289 self._source_row(it.file_hash, len(it.records), it.source_type, it.stat, it.meta),
3290 )
3291 for it in items
3292 ]
3294 def _drop_source_rows_unlocked(self, source_ids: list[str]) -> None:
3295 """Delete the source rows of *source_ids*, holding out each that a citation names.
3297 A held-out row keeps the path of a removed source for the citations of it,
3298 and holds nothing out of a sync.
3299 """
3300 if not source_ids:
3301 return
3302 table = self._sources_table()
3303 cited = self._cited_ids(source_ids)
3304 if cited:
3305 where = f"{_SOURCE_ID} IN ({_in_list(cited)})"
3306 table.update(where=where, values_sql=_CITED_ONLY_SQL)
3307 dropped = [source_id for source_id in source_ids if source_id not in cited]
3308 if dropped:
3309 _safe_delete_unlocked(table, f"{_SOURCE_ID} IN ({_in_list(dropped)})")
3311 def _cited_ids(self, source_ids: list[str]) -> set[str]:
3312 """The *source_ids* a citation row names."""
3313 table = self.open_table(CITATIONS_TABLE)
3314 if table is None:
3315 return set()
3316 query = table.search().where(f"{_SOURCE_ID} IN ({_in_list(source_ids)})")
3317 return {row[_SOURCE_ID] for row in query.select([_SOURCE_ID]).limit(None).to_list()}
3319 def delete_source(self, filename: str) -> None:
3320 """Remove a source file tracking record."""
3321 with self._write_lock():
3322 ids = self.source_ids([filename], SourceType.DOCUMENT)
3323 self._drop_source_rows_unlocked(list(ids.values()))
3324 self._invalidate_source_cache()
3326 def _remove_many_unlocked(
3327 self, names: list[str], held: Mapping[str, HeldRecord] | None = None
3328 ) -> None:
3329 """Delete the documents' rows and source records together; hold out each path in *held*.
3331 All deletes run under the caller's single ``write_lock()`` so no
3332 reader can observe chunks whose source record is already gone; one
3333 ``IN`` delete per table covers the whole set.
3334 """
3335 held = held or {}
3336 ids = self.source_ids(names, SourceType.DOCUMENT)
3337 self._delete_rows_of_unlocked(list(ids.values()))
3338 self._drop_source_rows_unlocked([ids[name] for name in ids if name not in held])
3339 self._hold_out_unlocked(held)
3341 def held_records(self) -> dict[str, HeldRecord]:
3342 """The record of every path a sync holds out, by path."""
3343 if self._library_absent():
3344 return {}
3345 table = self.open_table(SOURCES_TABLE)
3346 if table is None:
3347 return {}
3348 query = table.search().where(f"{_HELD_HASH} IS NOT NULL")
3349 rows = query.select([_PATH, *_HELD_COLUMNS]).limit(None).to_list()
3350 name = self._locations().name
3351 return {
3352 name(row[_PATH]): HeldRecord(
3353 HeldOut(row[_HELD_OUT]), row[_HELD_HASH], row[_HELD_REASON]
3354 )
3355 for row in rows
3356 }
3358 def hold_failed(self, failed: Mapping[str, HeldRecord]) -> None:
3359 """Hold each path of *failed* out of later syncs at its hash; a row keeps its document."""
3360 if not failed:
3361 return
3362 with self._write_lock():
3363 self._hold_out_unlocked(failed)
3364 self._invalidate_source_cache()
3366 def _hold_out_unlocked(self, held: Mapping[str, HeldRecord]) -> None:
3367 """Give each name of *held* its record in one commit. Caller holds ``write_lock()``."""
3368 if not held:
3369 return
3370 writes = [
3371 _SourceWrite(name, SourceType.DOCUMENT, record=record) for name, record in held.items()
3372 ]
3373 self._put_sources_unlocked(self._sources_table(), writes)
3375 def release_held(self, paths: Iterable[str]) -> None:
3376 """Hold none of *paths* out of a sync any more."""
3377 key = self._locations().key
3378 wanted = sorted({key(path, SourceType.DOCUMENT) for path in paths})
3379 if not wanted:
3380 return
3381 with self._write_lock():
3382 for start in range(0, len(wanted), _SOURCE_STAT_BATCH_ROWS):
3383 batch = _in_list(wanted[start : start + _SOURCE_STAT_BATCH_ROWS])
3384 self._release_unlocked(f"{_PATH} IN ({batch})")
3385 self._invalidate_source_cache()
3387 def release_failed(self) -> list[str]:
3388 """Hold no failed file out of a sync any more; returns their paths, sorted."""
3389 with self._write_lock():
3390 released = self._release_unlocked(f"{_HELD_OUT} = '{HeldOut.FAILED}'")
3391 self._invalidate_source_cache()
3392 return released
3394 def _release_unlocked(self, where: str) -> list[str]:
3395 """Release the held rows *where* selects; a row that holds no document goes.
3397 Caller holds ``write_lock()``. Returns the released paths, sorted.
3398 """
3399 table = self.open_table(SOURCES_TABLE)
3400 if table is None:
3401 return []
3402 query = table.search().where(f"({where}) AND {_HELD_HASH} IS NOT NULL")
3403 rows = query.select([_SOURCE_ID, _PATH, _HELD_OUT, "file_hash"]).limit(None).to_list()
3404 kept = [row[_SOURCE_ID] for row in rows if _holds_document(row)]
3405 if kept:
3406 table.update(
3407 where=f"{_SOURCE_ID} IN ({_in_list(kept)})",
3408 values_sql=dict.fromkeys(_HELD_COLUMNS, _NULL_TEXT_SQL),
3409 )
3410 self._drop_source_rows_unlocked(
3411 [row[_SOURCE_ID] for row in rows if not _holds_document(row)]
3412 )
3413 name = self._locations().name
3414 return sorted(name(row[_PATH]) for row in rows)
3416 def relocate_sources(self, moves: Sequence[SourceMove]) -> SourceRelocation:
3417 """Move the first of *moves* to their new names in one hold of the write lock.
3419 Returns the old name each new name took and how many of *moves* the hold
3420 settled; the caller passes the rest again. A move takes the first of its
3421 candidates that holds a source row and that no earlier move took; a move
3422 with none left is settled and absent from the names taken. The choice and
3423 the move share the lock, so no two writers take one old name.
3425 A moved source keeps its id, so its rows in every other table follow it:
3426 only its path, a title that came from the old name, and its disk stat
3427 change, in one commit of the source table. The sources inside a moved
3428 archive move with it. A hold settles at most ``_RELOCATE_HOLD_MOVES`` moves.
3429 """
3430 if not moves:
3431 return SourceRelocation({}, 0)
3432 with self._write_lock():
3433 step = self._relocate_hold_unlocked(moves[:_RELOCATE_HOLD_MOVES])
3434 self._invalidate_source_cache()
3435 return step
3437 def _relocate_hold_unlocked(self, moves: Sequence[SourceMove]) -> SourceRelocation:
3438 """Claim and move *moves*. Caller holds ``write_lock()``."""
3439 from lilbee.data.title import derive_title # circular at module scope
3441 sources = self.open_table(SOURCES_TABLE)
3442 if sources is None:
3443 return SourceRelocation({}, len(moves))
3444 moves, names = _located_moves(moves, self._locations())
3445 pools = {id(move.candidates): move.candidates for move in moves}.values()
3446 unclaimed = _held_sources(sources, {name for pool in pools for name in pool})
3447 claimed = _claimed_names(moves, unclaimed)
3448 rows = _source_rows(sources, [old for old in claimed if old is not None])
3449 taken = {move.new: old for move, old in zip(moves, claimed, strict=True) if old}
3450 titles = {
3451 old: _title_after_move(rows[old], old, move.new, derive_title)
3452 for move, old in zip(moves, claimed, strict=True)
3453 if old
3454 }
3455 moved = [
3456 _moved_source_row(rows[old], move, titles[old])
3457 for move, old in zip(moves, claimed, strict=True)
3458 if old
3459 ]
3460 if moved:
3461 movable = sources.search().where(_MOVABLE).limit(None).to_list()
3462 moved += _member_rows(movable, {old: new for new, old in taken.items()})
3463 self._delete_rows_of_unlocked(self._displaced_ids(moved))
3464 merge = sources.merge_insert(_SOURCE_ID).when_matched_update_all()
3465 displaced = f"{_PATH} IN ({_in_list(row[_PATH] for row in moved)})"
3466 merge = merge.when_not_matched_by_source_delete(displaced)
3467 merge.execute(pa.Table.from_pylist(moved, schema=sources.schema))
3468 self._retitle_chunks_unlocked(
3469 {rows[old][_SOURCE_ID]: title for old, title in titles.items()}
3470 )
3471 return SourceRelocation({names[new]: names[old] for new, old in taken.items()}, len(moves))
3473 def _displaced_ids(self, moved: list[dict[str, Any]]) -> list[str]:
3474 """The ids of the rows at the new paths of *moved* that are not moving themselves."""
3475 moving = {row[_SOURCE_ID] for row in moved}
3476 held = self._lookup(_PATH, (row[_PATH] for row in moved), _SOURCE_ID).values()
3477 return [source_id for source_id in held if source_id not in moving]
3479 def _retitle_chunks_unlocked(self, titles: Mapping[str, str | None]) -> None:
3480 """Give the chunks of each source id its new title, one update a title. Lock held.
3482 A title of ``_KEEP_TITLE`` leaves the chunks as they are.
3483 """
3484 chunks = self.open_table(CHUNKS_TABLE)
3485 if chunks is None or _TITLE_COLUMN not in chunks.schema.names:
3486 return
3487 by_title: dict[str | None, list[str]] = {}
3488 for source_id, title in titles.items():
3489 if title is not _KEEP_TITLE:
3490 by_title.setdefault(title, []).append(source_id)
3491 for title, source_ids in by_title.items():
3492 where = f"{_SOURCE_ID} IN ({_in_list(source_ids)})"
3493 chunks.update(where=where, values={_TITLE_COLUMN: title})
3495 def relocate_root(self, old: str, new: str) -> None:
3496 """Give each source at or below the root folder *old* the same place below *new*."""
3497 old_at, new_at = root_location(old), root_location(new)
3498 with self._write_lock():
3499 sources = self.open_table(SOURCES_TABLE)
3500 if sources is not None and old_at != new_at:
3501 _rekey_prefix(self.get_db(), sources, under_sql(_PATH, old_at), old_at, new_at)
3502 self._invalidate_source_cache()
3504 def fill_locations(self) -> None:
3505 """Give each source row a registered label still names its location."""
3506 with self._write_lock():
3507 _fill_locations(self.get_db(), self._locations())
3508 self._invalidate_source_cache()
3510 def name_by_label(self, labels: Iterable[str]) -> None:
3511 """Store each document of the registered roots *labels* by its shown name.
3513 A root about to leave the registry calls it, so a row it keeps for a
3514 citation shows its name and not a path on disk, as a label with no root
3515 does. A document of a root below one of *labels* that stays keeps its place.
3516 """
3517 going = set(labels)
3518 locations = self._locations()
3519 with self._write_lock():
3520 sources = self.open_table(SOURCES_TABLE)
3521 for root, label in locations.deepest_first:
3522 if sources is not None and label in going:
3523 where = f"{_DOCUMENT} AND {locations.named_by_sql(_PATH, root)}"
3524 _rekey_prefix(self.get_db(), sources, where, root, label)
3525 self._invalidate_source_cache()
3527 def member_sources(self, name: str) -> list[str]:
3528 """Sources ingested out of the archive *name*: every path under ``name/``."""
3529 table = self.open_table(SOURCES_TABLE)
3530 if table is None:
3531 return []
3532 locations = self._locations()
3533 at = escape_sql_string(locations.location(name))
3534 under = f"starts_with({_PATH}, '{at}/') AND {_LISTED}"
3535 rows = table.search().where(under).select([_PATH]).limit(None).to_list()
3536 return [locations.name(row[_PATH]) for row in rows]
3538 def remove_documents(
3539 self, names: list[str], held: Mapping[str, HeldRecord] | None = None
3540 ) -> RemoveResult:
3541 """Remove documents from the knowledge base by source name.
3543 Looks up known sources and deletes their chunks and source records. Never
3544 touches files on disk: source bytes are the user's, and a linked-in corpus
3545 must never be deleted. Each path of *held* takes its record in the same
3546 write, so a later sync holds it out. Durable, file-aware removal lives in
3547 :func:`lilbee.app.ingest.remove_documents_durably`.
3549 Returns a RemoveResult with removed and not_found lists.
3550 """
3551 filenames = [s["filename"] for s in self.get_sources()]
3552 known = set(filenames)
3553 members = _members_by_archive(names, filenames)
3554 targets = [*names, *(m for name in names for m in members.get(name, []))]
3555 removed = [name for name in targets if name in known]
3556 not_found = [name for name in names if name not in known]
3558 if removed or held:
3559 # One lock acquisition and one IN-delete per table for the whole
3560 # set, mirroring the batched flush path, instead of a LanceDB
3561 # version commit per document.
3562 with self._write_lock():
3563 self._remove_many_unlocked(removed, held)
3564 self._invalidate_source_cache()
3566 return RemoveResult(removed=removed, not_found=not_found)
3568 def clear_table(self, name: str, predicate: str) -> bool:
3569 """Delete rows matching *predicate* from *name*. Acquires write lock.
3571 Returns whether the delete succeeded, so a caller recording the
3572 outcome (prune, the legacy migration) does not report success over a
3573 swallowed failure.
3574 """
3575 with self._write_lock():
3576 table = self.open_table(name)
3577 if table is None:
3578 return True
3579 return _safe_delete_unlocked(table, predicate)
3581 def clear_and_add(self, name: str, schema: pa.Schema, rows: list[dict], predicate: str) -> None:
3582 """Replace the rows matching *predicate* with *rows* in one locked write.
3584 Delete and add run under a single write lock, so a reader never observes
3585 the table emptied mid-rebuild. A delete failure propagates, following the
3586 same rule as :meth:`_delete_rows_of_unlocked`: adding over rows whose
3587 predecessors are still there duplicates them, and reporting success would
3588 leave the caller acting on state it thinks it replaced.
3589 """
3590 with self._write_lock():
3591 db = self.get_db()
3592 table = ensure_table(db, name, schema)
3593 table.delete(predicate)
3594 if rows:
3595 table.add(rows)
3597 def add_citations(self, records: list[CitationRecord]) -> int:
3598 """Add citation records to the store. Returns count added."""
3599 if not records:
3600 return 0
3601 with self._write_lock():
3602 table = ensure_table(self.get_db(), CITATIONS_TABLE, _citations_schema())
3603 table.add(self._keyed_citations_unlocked(records))
3604 return len(records)
3606 def _keyed_citations_unlocked(self, records: Iterable[CitationRecord]) -> list[dict]:
3607 """Citation rows that name each cited source by id. Caller holds the lock."""
3608 records = list(records)
3609 ids = self._ensure_ids_unlocked(_documents(record["source_filename"] for record in records))
3610 return _keyed(records, "source_filename", ids)
3612 def _shown_citations(self, rows: list[dict[str, Any]]) -> list[CitationRecord]:
3613 """Citation *rows* as records, each naming its source by path."""
3614 paths = self.source_paths(row[_SOURCE_ID] for row in rows)
3615 return cast("list[CitationRecord]", _pathed(rows, "source_filename", paths))
3617 def get_citations_for_wiki(self, wiki_source: str) -> list[CitationRecord]:
3618 """Get all citations for a wiki page."""
3619 table = self.open_table(CITATIONS_TABLE)
3620 if table is None:
3621 return []
3622 escaped = escape_sql_string(wiki_source)
3623 return self._shown_citations(table.search().where(f"wiki_source = '{escaped}'").to_list())
3625 def get_citations_for_source(self, source_filename: str) -> list[CitationRecord]:
3626 """Get all citations that reference a source document (reverse lookup)."""
3627 table = self.open_table(CITATIONS_TABLE)
3628 source_id = self.source_ids([source_filename], SourceType.DOCUMENT).get(source_filename)
3629 if table is None or source_id is None:
3630 return []
3631 rows = table.search().where(f"{_SOURCE_ID} = '{source_id}'").to_list()
3632 return cast(
3633 "list[CitationRecord]", _pathed(rows, "source_filename", {source_id: source_filename})
3634 )
3636 def delete_citations_for_wiki(self, wiki_source: str) -> bool:
3637 """Delete all citations for a wiki page. Returns whether the delete succeeded."""
3638 return self.clear_table(CITATIONS_TABLE, _citations_for_wiki_predicate(wiki_source))
3640 def delete_all_wiki_rows(self) -> bool:
3641 """Delete every wiki chunk row, every citation, and every mention.
3642 Returns whether all deletes succeeded, so a caller cannot report a wipe
3643 over a swallowed failure. Only the wiki layer writes citations and
3644 mentions, so wiping it empties those tables outright. All deletes run
3645 before the results combine.
3646 """
3647 chunks_cleared = self.clear_table(CHUNKS_TABLE, _WIKI_CHUNK)
3648 citations_cleared = self.clear_table(CITATIONS_TABLE, "1 = 1")
3649 mentions_cleared = self.clear_wiki_mentions()
3650 return chunks_cleared and citations_cleared and mentions_cleared
3652 def replace_citations_for_wiki(self, wiki_source: str, records: list[CitationRecord]) -> None:
3653 """Swap a wiki page's citations for *records* under one write lock.
3655 The lock keeps another writer from landing between the delete and the
3656 add; the two remain separate commits, so a crash in between leaves the
3657 rows absent until the page is regenerated or accepted again.
3658 """
3659 with self._write_lock():
3660 table = ensure_table(self.get_db(), CITATIONS_TABLE, _citations_schema())
3661 rows = self._keyed_citations_unlocked(records)
3662 table.delete(_citations_for_wiki_predicate(wiki_source))
3663 if rows:
3664 table.add(rows)
3666 def replace_wiki_mentions_for_source(self, source: str, rows: list[dict]) -> None:
3667 """Swap one source's wiki mention rows for *rows* under one write lock.
3669 A source contributes its whole mention set at once, so an incremental
3670 wiki refresh replaces exactly that source's rows and leaves every other
3671 source's evidence in place for the corpus-wide aggregate.
3672 """
3673 with self._write_lock():
3674 ids = self._ensure_ids_unlocked(_documents(row["source"] for row in rows))
3675 table = ensure_table(self.get_db(), WIKI_MENTIONS_TABLE, _wiki_mentions_schema())
3676 for source_id in self.source_ids([source], SourceType.DOCUMENT).values():
3677 table.delete(f"{_SOURCE_ID} = '{source_id}'")
3678 if rows:
3679 table.add(_keyed(rows, "source", ids))
3681 def wiki_mention_rows(self, slugs: Iterable[str] | None = None) -> list[dict]:
3682 """Every wiki mention row, or only those for *slugs*, with its source's id and path.
3684 The wiki aggregates these across sources to rebuild the stub index. A
3685 full refresh reads them all; an incremental one reads only the slugs it
3686 touched, to recompute their corpus-wide totals without a full scan.
3687 """
3688 table = self.open_table(WIKI_MENTIONS_TABLE)
3689 if table is None:
3690 return []
3691 query = table.search()
3692 if slugs is not None:
3693 wanted = list(slugs)
3694 if not wanted:
3695 return []
3696 joined = ", ".join(f"'{escape_sql_string(s)}'" for s in wanted)
3697 query = query.where(f"slug IN ({joined})")
3698 rows = query.limit(None).to_list()
3699 paths = self.source_paths(row[_SOURCE_ID] for row in rows)
3700 return [
3701 {**row, "source": paths[source_id]}
3702 for row in rows
3703 if (source_id := row[_SOURCE_ID]) in paths
3704 ]
3706 def clear_wiki_mentions(self) -> bool:
3707 """Drop every wiki mention row (a full rebuild starts from nothing)."""
3708 return self.clear_table(WIKI_MENTIONS_TABLE, "1 = 1")
3710 def has_wiki_mentions(self) -> bool:
3711 """Whether any wiki mention row exists.
3713 A cold store or one migrated from the file-only index has none, which
3714 forces a refresh to rebuild in full and seed the table before an
3715 incremental pass can aggregate over it.
3716 """
3717 table = self.open_table(WIKI_MENTIONS_TABLE)
3718 return table is not None and table.count_rows() > 0
3720 def _memories_schema(self) -> pa.Schema:
3721 return pa.schema(
3722 [
3723 pa.field("id", pa.utf8()),
3724 pa.field("owner", pa.utf8()),
3725 pa.field("shared", pa.bool_()),
3726 pa.field("kind", pa.utf8()),
3727 pa.field("source", pa.utf8()),
3728 pa.field("text", pa.utf8()),
3729 pa.field("vector", pa.list_(pa.float32(), self._config.embedding_dim)),
3730 pa.field("created_at", pa.utf8()),
3731 pa.field("updated_at", pa.utf8()),
3732 ]
3733 )
3735 def _duplicate_memory_id_unlocked(self, table: LanceTable, record: MemoryRow) -> str | None:
3736 """Return the id of a near-duplicate same-owner, same-kind memory, if any."""
3737 if table.count_rows() == 0:
3738 return None
3739 predicate = (
3740 f"owner = '{escape_sql_string(record.owner)}' "
3741 f"AND kind = '{escape_sql_string(record.kind)}'"
3742 )
3743 rows = (
3744 table.search(record.vector)
3745 .distance_type(_VECTOR_METRIC)
3746 .where(predicate)
3747 .limit(1)
3748 .to_list()
3749 )
3750 if rows and rows[0].get(_DISTANCE_COLUMN, 1.0) <= self._config.memory_dedup_distance:
3751 return str(rows[0]["id"])
3752 return None
3754 def _evict_overflow_unlocked(self, table: LanceTable, owner: str) -> None:
3755 """Delete oldest memories for *owner* so an incoming insert stays within the cap."""
3756 cap = self._config.memory_max_per_owner
3757 predicate = f"owner = '{escape_sql_string(owner)}'"
3758 rows = table.search().where(predicate).limit(None).to_list()
3759 if len(rows) < cap:
3760 return
3761 rows.sort(key=lambda r: r.get("created_at", ""))
3762 for row in rows[: len(rows) - (cap - 1)]:
3763 _safe_delete_unlocked(table, f"id = '{escape_sql_string(str(row['id']))}'")
3765 def add_memory(self, record: MemoryRow) -> str:
3766 """Insert *record*, or update the nearest same-owner duplicate in place.
3768 Returns the stored id. Raises ``EmbeddingModelMismatchError`` when the store
3769 was built under a different embedding model, and ``ValueError`` on a vector
3770 dimension mismatch.
3771 """
3772 if len(record.vector) != self._config.embedding_dim:
3773 raise ValueError(
3774 f"Memory vector dimension mismatch: expected "
3775 f"{self._config.embedding_dim}, got {len(record.vector)}"
3776 )
3777 with self._write_lock():
3778 embedding_model = self._config.embedding_model
3779 embedding_dim = self._config.embedding_dim
3780 self._ensure_embedding_compat()
3781 db = self.get_db()
3782 table = ensure_table(db, MEMORIES_TABLE, self._memories_schema())
3783 duplicate_id = self._duplicate_memory_id_unlocked(table, record)
3784 if duplicate_id is not None and _safe_delete_unlocked(
3785 table, f"id = '{escape_sql_string(duplicate_id)}'"
3786 ):
3787 # Only reuse the id once the old row is actually gone; a swallowed
3788 # delete failure would otherwise leave two rows with the same id.
3789 record.id = duplicate_id
3790 self._evict_overflow_unlocked(table, record.owner)
3791 table.add([record.model_dump(mode="json")])
3792 if self.get_meta() is None:
3793 self._write_meta_unlocked(
3794 embedding_model=embedding_model, embedding_dim=embedding_dim
3795 )
3796 return record.id
3798 def get_memories(
3799 self,
3800 *,
3801 owner_predicate: str,
3802 kind: MemoryKind | None = None,
3803 ) -> list[MemoryRow]:
3804 """Return memories matching *owner_predicate* and optional *kind*, newest first."""
3805 table = self.open_table(MEMORIES_TABLE)
3806 if table is None:
3807 return []
3808 clauses = [f"({owner_predicate})"]
3809 if kind is not None:
3810 clauses.append(f"kind = '{escape_sql_string(kind)}'")
3811 rows = table.search().where(" AND ".join(clauses)).limit(None).to_list()
3812 memories = [MemoryRow(**r) for r in rows]
3813 memories.sort(key=lambda m: m.created_at, reverse=True)
3814 return memories
3816 def search_memories(
3817 self,
3818 query_vector: Vector,
3819 *,
3820 owner_predicate: str,
3821 top_k: int,
3822 max_distance: float,
3823 ) -> list[MemoryRow]:
3824 """Vector-recall FACT memories within *max_distance*, best first."""
3825 table = self.open_table(MEMORIES_TABLE)
3826 if table is None or top_k <= 0:
3827 return []
3828 self._ensure_embedding_compat()
3829 predicate = f"({owner_predicate}) AND kind = '{MemoryKind.FACT}'"
3830 rows = (
3831 table.search(query_vector)
3832 .distance_type(_VECTOR_METRIC)
3833 .where(predicate)
3834 .limit(top_k)
3835 .to_list()
3836 )
3837 return [MemoryRow(**r) for r in rows if r.get(_DISTANCE_COLUMN, 1.0) <= max_distance]
3839 def update_memory(self, memory_id: str, *, shared: bool, owner: str) -> bool:
3840 """Set the *shared* flag on *owner*'s memory. Returns True when found and owned.
3842 The ``owner`` predicate scopes the mutation to the caller's namespace so an
3843 agent cannot flip another owner's (or the human's) memory.
3844 """
3845 with self._write_lock():
3846 table = self.open_table(MEMORIES_TABLE)
3847 if table is None:
3848 return False
3849 predicate = self._owned_memory_predicate(memory_id, owner)
3850 rows = table.search().where(predicate).limit(1).to_list()
3851 if not rows:
3852 return False
3853 record = MemoryRow(**rows[0])
3854 record.shared = shared
3855 record.updated_at = datetime.now(UTC).isoformat()
3856 # If the delete fails, do not add the modified copy: that would leave
3857 # two rows for one id. Report not-updated instead.
3858 if not _safe_delete_unlocked(table, predicate):
3859 return False
3860 table.add([record.model_dump(mode="json")])
3861 return True
3863 def delete_memory(self, memory_id: str, *, owner: str) -> bool:
3864 """Delete *owner*'s memory by id. Returns True when a matching row was deleted.
3866 The ``owner`` predicate scopes the delete to the caller's namespace so an
3867 agent cannot destroy another owner's (or the human's) memory.
3868 """
3869 with self._write_lock():
3870 table = self.open_table(MEMORIES_TABLE)
3871 if table is None:
3872 return False
3873 predicate = self._owned_memory_predicate(memory_id, owner)
3874 if not table.search().where(predicate).limit(1).to_list():
3875 return False
3876 # Report the real outcome: a swallowed delete failure must not be
3877 # reported as a successful forget.
3878 return _safe_delete_unlocked(table, predicate)
3880 @staticmethod
3881 def _owned_memory_predicate(memory_id: str, owner: str) -> str:
3882 """SQL predicate matching a single memory id within *owner*'s namespace."""
3883 return f"id = '{escape_sql_string(memory_id)}' AND owner = '{escape_sql_string(owner)}'"
3885 def rebuild_memory_embeddings(self, embed: Callable[[list[str]], list[Vector]]) -> int:
3886 """Re-embed every memory under the current model, recreating the table.
3888 The vector column dimension is immutable, so a different-dim model needs a
3889 fresh table; recreating unconditionally also covers the same-dim case. Memory
3890 text is human-authored and re-embeddable, so no data is lost. Returns the count.
3892 The snapshot, embed, and table rebuild all run under the write lock so a
3893 concurrent ``add_memory`` cannot commit into the read-then-drop window and
3894 be erased; it either lands before the snapshot or blocks until the rebuild
3895 finishes.
3896 """
3897 with self._write_lock():
3898 table = self.open_table(MEMORIES_TABLE)
3899 if table is None:
3900 return 0
3901 rows = table.search().limit(None).to_list()
3902 if not rows:
3903 return 0
3904 memories = [MemoryRow(**r) for r in rows]
3905 vectors = embed([m.text for m in memories])
3906 for memory, vector in zip(memories, vectors, strict=True):
3907 # MemoryRow serializes to JSON on write, which has no ndarray encoding.
3908 memory.vector = vector.tolist()
3909 db = self.get_db()
3910 db.drop_table(MEMORIES_TABLE)
3911 new_table = ensure_table(db, MEMORIES_TABLE, self._memories_schema())
3912 new_table.add([m.model_dump(mode="json") for m in memories])
3913 return len(memories)
3915 def close(self) -> None:
3916 """Release the database connection and reset state."""
3917 self._db = None
3918 self._fts_ready = False
3919 self._title_fts_ready = False
3920 self._scalar_ready = False
3922 def drop_all(self) -> None:
3923 """Drop every table except ``_memories`` and ``_crawl_pages`` -- used by rebuild.
3925 Memory is user-authored data with no on-disk source, not derived from
3926 documents, so a rebuild preserves it. The crawl page table says which files
3927 under ``_web`` a crawl owns, which no document holds. Only a factory reset
3928 (which deletes the data directory) clears them.
3929 """
3930 with self._write_lock():
3931 self._fts_ready = False
3932 self._title_fts_ready = False
3933 self._scalar_ready = False
3934 db = self.get_db()
3935 # Before the tables: a record that outlives them would describe the next ones.
3936 self.forget_wiki_refreshed()
3937 kept = {MEMORIES_TABLE, CRAWL_PAGES_TABLE}
3938 dropped = [name for name in table_names(db) if name not in kept]
3939 # The source rows go first: a drop that stops part way leaves no file unchanged.
3940 for name in sorted(dropped, key=lambda name: name != SOURCES_TABLE):
3941 db.drop_table(name)
3942 self._invalidate_source_cache()