diff --git a/.gitignore b/.gitignore index 3ae8ddff..98939eb0 100644 --- a/.gitignore +++ b/.gitignore @@ -75,3 +75,4 @@ benchmark/results/ # integration tests outputs tests/integration/logs/ +daily/ diff --git a/docs/en/memory_search.md b/docs/en/memory_search.md index 529035ae..9ee3605a 100644 --- a/docs/en/memory_search.md +++ b/docs/en/memory_search.md @@ -32,6 +32,8 @@ and `[[wikilinks]]`. JSONL uses the `default` chunker and creates overlapping ch ## How the Index Is Built +### Index Update + The background Job `index_update_loop` maintains the index using configuration from `reme/config/default.yaml`: ```yaml @@ -67,6 +69,22 @@ stable batch of changes. The Markdown chunker parses YAML frontmatter, heading structure, and `[[...]]` into `FileNode`, `FileChunk`, and `FileLink` objects. For detailed chunking rules, see [Memory as File](./memory_as_file.md#memory-chunking). +### Index Optimization + +Both BM25 and the FAISS HNSW vector index use tombstone markers instead of physical removal when deleting nodes; +too many tombstones degrade search performance. An idle-time optimization mechanism is built in—the `optimize_index_cron` +scheduled job compacts tombstones and rebuilds indexes during off-peak hours: + +```yaml +optimize_index_cron: + backend: cron + cron: "0 2 * * *" + steps: + - backend: optimize_index_step +``` + +By default it runs at 2:00 AM daily; adjust the cron expression to customize the schedule. + ## What file_store Contains The default `file_store.default` backend is `local`: @@ -90,7 +108,8 @@ It combines three kinds of capability: | `embedding_store` | Disabled | When enabled, generate embeddings for chunks and support vector recall. | Out of the box, search therefore uses primarily BM25 plus link expansion. After setting `embedding_store: default`, -`SearchStep` runs vector and keyword recall together. +`SearchStep` runs vector and keyword recall together. Additionally, switching the `file_store` `backend` from `local` to +`faiss` upgrades vector retrieval from a linear scan to a FAISS HNSW index, offering faster recall at scale. ## How to Search diff --git a/docs/zh/memory_search.md b/docs/zh/memory_search.md index ae4bf6aa..483bf0d4 100644 --- a/docs/zh/memory_search.md +++ b/docs/zh/memory_search.md @@ -30,6 +30,8 @@ workspace files ## 索引怎么构建 +### 索引更新 + 索引由后台 Job `index_update_loop` 维护,配置来自 `reme/config/default.yaml`: ```yaml @@ -63,6 +65,20 @@ index_update_loop: Markdown chunker 会解析 YAML frontmatter、标题结构和 `[[...]]`,产出 `FileNode`、`FileChunk` 和 `FileLink`。更细的分块规则见 [Memory as File](./memory_as_file.md#memory-chunking)。 +### 索引优化 + +BM25 和 FAISS HNSW 向量索引在删除节点时都采用墓碑(tombstone)标记而非物理移除,积累过多会拖慢搜索。为此内置了闲暇时间索引优化机制——`optimize_index_cron` 定时任务在低峰期压缩墓碑并重建索引: + +```yaml +optimize_index_cron: + backend: cron + cron: "0 2 * * *" + steps: + - backend: optimize_index_step +``` + +默认每天凌晨 2 点执行,调整 cron 表达式即可自定义调度时间。 + ## file_store 里有什么 默认 `file_store.default` 是 `local`: @@ -85,7 +101,7 @@ file_store: | `file_graph.default` | 启用 | 保存 `FileNode` 和 wikilink 边 | | `embedding_store` | 默认关闭 | 开启后为 chunk 生成 embedding,并支持向量召回 | -所以开箱搜索主要是 BM25 + 链接展开。把 `embedding_store: default` 打开后,`SearchStep` 会同时跑向量召回和关键词召回。 +所以开箱搜索主要是 BM25 + 链接展开。把 `embedding_store: default` 打开后,`SearchStep` 会同时跑向量召回和关键词召回。此时若将 `file_store` 的 `backend` 从 `local` 改为 `faiss`,向量检索会从线性扫描升级为 FAISS HNSW 索引,在大规模 chunk 场景下召回效率更高。 ## 怎么搜索 diff --git a/reme/components/file_store/base_file_store.py b/reme/components/file_store/base_file_store.py index afba93a3..b89dd183 100644 --- a/reme/components/file_store/base_file_store.py +++ b/reme/components/file_store/base_file_store.py @@ -63,3 +63,12 @@ class BaseFileStore(BaseComponent): @abstractmethod async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: """Full-text keyword search over chunk text.""" + + # -- maintenance ------------------------------------------------------------ + + async def optimize_index(self) -> None: + """Optional idle-time maintenance hook (e.g. compacting derived indexes). + + Meant to be invoked off the request path (cron / idle schedulers). + Backends without derived index state keep the default no-op. + """ diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index b8bfc1fa..5fdf7846 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -1,7 +1,9 @@ -"""FAISS-backed file store: chunk JSONL stays authoritative; FAISS replaces the linear vector scan.""" +"""FAISS-backed file store: chunk JSONL stays authoritative; FAISS HNSW replaces the linear vector scan.""" import asyncio +import hashlib import json +from contextlib import suppress from uuid import uuid4 import aiofiles @@ -14,13 +16,49 @@ from ...schema import FileChunk, FileNode @R.register("faiss") class FaissLocalFileStore(LocalFileStore): - """LocalFileStore variant whose vector_search is backed by a FAISS IndexFlatIP. + """LocalFileStore variant whose vector_search is backed by a FAISS IndexHNSWFlat. Chunk persistence is unchanged (JSONL, owned by the parent). FAISS state is stored alongside as a binary index plus an id-map sidecar. If either file is missing or stale, the index is rebuilt from ``self.file_chunks``, which remains the source of truth. + HNSW parameters (``hnsw_m``, ``hnsw_ef_construction``) control graph + connectivity and build-time quality. ``efSearch`` is not a stored property; + it is set to ``limit * 5`` at query time so the candidate pool scales with + the number of results requested. FAISS internally raises the beam width to + ``max(efSearch, k)`` when ``k`` exceeds this value during progressive recall. + + When the index is small (``ntotal < max(1, sqrt(max(limit, ef_search))) * + M``) the graph would visit nearly every node anyway, so ``vector_search`` + bypasses HNSW and does an exact brute-force scan via the underlying + ``IndexFlat`` storage. The threshold ties brute-force coverage to ``M`` + (graph degree): higher M means costlier graph traversal, so brute-force is + preferred for larger indexes. + + Tombstone-driven compaction rebuilds the index when deleted rows accumulate. + The threshold is dynamic: ``max(tombstone_compact_ratio * ntotal, 128)`` + (default ratio 30%), so larger indexes tolerate more tombstones before a + rebuild. ``max_tombstones`` may be set to an int to override this with a + fixed threshold (useful in tests). + + Rebuilding an HNSW graph is expensive. When ``async_reindex`` is enabled the + rebuild is moved off the request path as a submit-a-flag job: mutations set a + boolean request flag and a single long-lived worker coroutine consumes it, + building the new index in a worker thread (FAISS releases the GIL during + ``add``) from a snapshot while the current index keeps serving searches, then + atomically swaps it in. Exactly one worker runs, so only one reindex proceeds + at a time; repeated submissions coalesce into the flag, and writes that land + during a build are reconciled: the stale snapshot is discarded and the flag + re-armed so the next round snapshots the current state. The live index is + never overwritten by a stale snapshot, so ``close()`` cannot persist one. + ``async_reindex`` defaults to ``False`` so behavior is unchanged unless opted + in; ``load`` (no old index to serve) always rebuilds inline. Embedding + backfill adds new vectors to the live index incrementally (HNSW insertion is + incremental anyway, so a rebuild would redo identical work); only tombstone + pressure or, in async mode, a delta that rivals the live row count goes + through a full rebuild. + faiss is imported lazily inside ``__init__`` so that merely importing this module (e.g. via ``reme version``) does not trigger the SWIG bindings and their associated DeprecationWarnings. @@ -29,13 +67,21 @@ class FaissLocalFileStore(LocalFileStore): def __init__( self, normalize: bool = True, - max_tombstones: int = 1024, + max_tombstones: int | None = None, + tombstone_compact_ratio: float = 0.3, + hnsw_m: int = 32, + hnsw_ef_construction: int = 64, + async_reindex: bool = False, **kwargs, ): super().__init__(**kwargs) self._faiss = self._import_faiss() self.normalize = normalize self.max_tombstones = max_tombstones + self.tombstone_compact_ratio = tombstone_compact_ratio + self.hnsw_m = hnsw_m + self.hnsw_ef_construction = hnsw_ef_construction + self.async_reindex = async_reindex self.faiss_path = self.component_metadata_path / f"faiss_index_{self.name}_{self.store_version}.bin" self.faiss_idmap_path = self.component_metadata_path / f"faiss_idmap_{self.name}_{self.store_version}.json" self._faiss_index = None # faiss.Index | None @@ -43,6 +89,15 @@ class FaissLocalFileStore(LocalFileStore): self._id_to_row: dict[str, int] = {} # chunk_id -> row (live entries only) self._tombstones: set[int] = set() # rows whose chunk_id was deleted self._faiss_dump_lock = asyncio.Lock() + # Async reindex machinery (only exercised when async_reindex is True). + # Submitting a rebuild just sets ``_reindex_event``; a single long-lived + # worker coroutine consumes it, so at most one reindex ever runs at a time + # and repeated submissions coalesce into the boolean flag. + self._reindex_event = asyncio.Event() # set == "rebuild requested" + self._reindex_worker_task: asyncio.Task | None = None + self._reindex_busy = False # True while a build is in flight (single worker) + self._index_writes = 0 # bumped on every index mutation; used to re-arm the flag + self._closing = False # set during _close() to stop spawning background reindexes @staticmethod def _import_faiss(): @@ -61,10 +116,21 @@ class FaissLocalFileStore(LocalFileStore): return self.embedding_store.dimensions if self.embedding_store is not None else 0 def _new_index(self): - return self._faiss.IndexFlatIP(self._dim) + index = self._faiss.IndexHNSWFlat(self._dim, self.hnsw_m, self._faiss.METRIC_INNER_PRODUCT) + index.hnsw.efConstruction = self.hnsw_ef_construction + index.hnsw.efSearch = 64 # safe default; overwritten at query time by _set_ef_search + return index + + def _set_ef_search(self, index, limit: int) -> None: + """Set HNSW efSearch to ``limit * 5`` for a good recall/speed balance. + + FAISS internally uses ``max(efSearch, k)`` during search, so when the + over-fetch ``k`` exceeds this value the beam width is raised automatically. + """ + index.hnsw.efSearch = limit * 5 def _prepare(self, vec: np.ndarray) -> np.ndarray: - """Cast to float32 (FAISS requirement) and L2-normalize so IndexFlatIP gives cosine.""" + """Cast to float32 (FAISS requirement) and L2-normalize so inner product gives cosine.""" v = np.ascontiguousarray(vec, dtype=np.float32) if v.ndim == 1: v = v[None, :] @@ -87,11 +153,13 @@ class FaissLocalFileStore(LocalFileStore): self._tombstones.add(old_row) self._id_map.append(cid) self._id_to_row[cid] = row + self._index_writes += 1 def _tombstone(self, chunk_id: str) -> None: row = self._id_to_row.pop(chunk_id, None) if row is not None: self._tombstones.add(row) + self._index_writes += 1 def _rebuild_index(self) -> None: """Rebuild FAISS state from self.file_chunks (the source of truth).""" @@ -105,16 +173,228 @@ class FaissLocalFileStore(LocalFileStore): vectors = np.stack([c.embedding for c in chunks]) self._add_to_index([c.id for c in chunks], vectors) + def _tombstone_threshold(self, scale: float = 1.0) -> int: + """Tombstone count that warrants a rebuild. + + Dynamic: ``max(tombstone_compact_ratio * ntotal * scale, 128)`` so larger + indexes tolerate more tombstones. ``scale`` lets callers lower the bar + (idle-time optimize_index uses 0.5). ``max_tombstones`` (when set) overrides with + a fixed value regardless of scale. + """ + if self.max_tombstones is not None: + return self.max_tombstones + ntotal = self._faiss_index.ntotal if self._faiss_index else 0 + return max(int(self.tombstone_compact_ratio * ntotal * scale), 128) + def _compact_if_needed(self) -> None: - if len(self._tombstones) >= self.max_tombstones: + # Dynamic threshold: max(ratio * ntotal, 128). This scales with index + # size so larger indexes tolerate more tombstones before a rebuild. + # max_tombstones (when set) overrides with a fixed value for testing. + if len(self._tombstones) < self._tombstone_threshold(): + return + if self.async_reindex: + self._submit_reindex() # flag the worker; it coalesces repeated requests + else: + self._rebuild_index() + + async def optimize_index(self) -> None: + """Idle-time maintenance: compact tombstones at half the write-path bar. + + The write path only rebuilds once tombstones reach the full threshold; + running off the request path we can afford to compact earlier, so this + uses ``scale=0.5``. Whether the rebuild runs inline or through the + background worker follows ``async_reindex``, same as the write path. + """ + await super().optimize_index() + if self._faiss_index is None: + return + tombstones = len(self._tombstones) + threshold = self._tombstone_threshold(scale=0.5) + if tombstones <= threshold: + return + self.logger.info(f"{self.name}: optimize_index compacting {tombstones} tombstones (threshold {threshold})") + if self.async_reindex: + self._submit_reindex() + else: self._rebuild_index() async def _after_embedding_backfill(self) -> None: - """Make newly backfilled vectors visible to FAISS before persistence.""" - self._rebuild_index() + """Make newly backfilled vectors visible to FAISS. + + HNSW insertion is incremental by nature, so backfilled vectors are added + directly to the live index instead of rebuilding the whole graph: + existing rows keep their connections and a rebuild would redo identical + work. Backfill never invalidates indexed rows -- stale embeddings have a + different dimension and can never have entered this index -- so no new + tombstones arise here; accumulated tombstone pressure is still honored + via ``_compact_if_needed`` afterwards. + + With ``async_reindex`` enabled, a delta that rivals the live row count + (e.g. the initial mass backfill into an empty index) is routed through + the background worker instead: a full off-loop rebuild costs about the + same as the inline add but keeps the event loop free. + """ + if self.embedding_store is None or self._dim == 0: + return + if self._faiss_index is None: + self._rebuild_index() + return + to_add = [ + chunk + for cid, chunk in self.file_chunks.items() + if cid not in self._id_to_row and self._embedding_dim_matches(chunk.embedding) + ] + if to_add: + if self.async_reindex and len(to_add) >= max(len(self._id_to_row), 1): + self._submit_reindex() + return + vectors = np.stack([c.embedding for c in to_add]) + self._add_to_index([c.id for c in to_add], vectors) + self._compact_if_needed() + + # -- async reindex ---------------------------------------------------- + + def _submit_reindex(self) -> None: + """Submit a rebuild request: ensure the worker exists and raise the flag. + + Idempotent -- the flag is a boolean, so repeated submissions while a build + is running or pending collapse into (at most) one follow-up rebuild. + Setting the flag is the only way to submit; the worker owns execution. + """ + if self._closing: + return # do not spawn background work while shutting down + self._ensure_reindex_worker() + self._reindex_event.set() + + def _ensure_reindex_worker(self) -> None: + """Start the single long-lived reindex worker if it is not running.""" + if self._reindex_worker_task is None or self._reindex_worker_task.done(): + self._reindex_worker_task = asyncio.ensure_future(self._reindex_worker()) + + async def _reindex_worker(self) -> None: + """Single consumer of reindex requests. + + Because exactly one worker drains the flag, only one reindex ever runs at a + time. The flag is cleared before snapshotting, so any write that lands + during the build is detected by ``_reindex_async`` (which compares + ``_index_writes`` before and after the build). When drift is detected the + stale snapshot is discarded and the flag re-armed, so the system + converges without ever publishing a stale index. + """ + while not self._closing: + await self._reindex_event.wait() + # Mark busy before clearing the flag (no await in between) so an + # observer can never see "flag clear and not busy" mid-handoff. + self._reindex_busy = True + self._reindex_event.clear() + if self._closing: + self._reindex_busy = False + return + try: + await self._reindex_async() + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception as e: # pragma: no cover - defensive + self.logger.exception(f"{self.name}: async reindex failed: {e}") + finally: + self._reindex_busy = False + + async def _reindex_async(self) -> None: + """Build a new index from a snapshot without blocking searches, then swap it + in atomically. If writes landed on the live index during the build + (``_index_writes`` moved), the snapshot is stale: it is discarded and the + flag re-armed so the next round snapshots the current state. The live + index is never overwritten by a stale snapshot, so ``close()`` cannot + persist one. + """ + if self.embedding_store is None or self._dim == 0: + return + dim = self._dim + items = [ + (cid, chunk.embedding) + for cid, chunk in self.file_chunks.items() + if self._embedding_dim_matches(chunk.embedding) + ] + snapshot_ids = [cid for cid, _ in items] + vectors = np.stack([emb for _, emb in items]) if items else None + writes_at_snapshot = self._index_writes # captured with the snapshot (no await yet) + + new_index = await asyncio.to_thread(self._build_index_blocking, dim, vectors) + + # Reconcile: if writes landed on the live index during the build, the + # snapshot is stale. Discard it and re-arm the flag so the next round + # snapshots the current state. The live index keeps serving searches + # with the already-applied mutations; it is never overwritten by a stale + # snapshot, so close() cannot persist one. + if self._index_writes != writes_at_snapshot: + self._reindex_event.set() + self.logger.info( + f"Async reindex discarded stale snapshot " + f"(writes {writes_at_snapshot} -> {self._index_writes}); re-arming", + ) + return + + # Atomic swap into the snapshot state (no await between assignments) so a + # concurrent search never observes a torn index / id_map / tombstone triple. + self._faiss_index = new_index + self._id_map = list(snapshot_ids) + self._id_to_row = {cid: row for row, cid in enumerate(snapshot_ids)} + self._tombstones = set() + self.logger.info(f"Async reindex complete: {self._faiss_index.ntotal} rows, live={len(self._id_to_row)}") + + def _build_index_blocking(self, dim: int, vectors: "np.ndarray | None"): + """Build a fresh HNSW index off the event loop (worker thread). + + FAISS releases the GIL during ``add``, so the event loop keeps serving + searches on the current index while this runs. + """ + index = self._faiss.IndexHNSWFlat(dim, self.hnsw_m, self._faiss.METRIC_INNER_PRODUCT) + index.hnsw.efConstruction = self.hnsw_ef_construction + index.hnsw.efSearch = 64 + if vectors is not None and vectors.size: + index.add(self._prepare(vectors)) + return index + + async def _stop_reindex_worker(self) -> None: + """Stop the reindex worker; used by close() and clear(). + + Cancelling detaches the awaiting coroutine promptly. An orphaned build + thread (if any) finishes into a local index that is simply discarded, so no + shared state is corrupted. + """ + self._reindex_event.clear() + task = self._reindex_worker_task + self._reindex_worker_task = None + if task is not None and not task.done(): + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + self._reindex_busy = False # -- persistence ------------------------------------------------------ + def _chunks_embedding_digest(self) -> str: + """Order-independent digest of the live (chunk_id, embedding) set. + + Hashes float16 canonical bytes (the chunk JSONL serialization dtype) so + the value is identical whether an embedding sits fresh in memory + (possibly float32 from the provider; assignment bypasses the pydantic + validator) or has round-tripped through the JSONL. Written into the + idmap sidecar at dump time and recomputed from ``self.file_chunks`` at + load time: a mismatch means the sidecar vectors belong to a different + chunk generation than the authoritative JSONL, which the live-id set + check alone cannot detect for same-ID in-place updates. + """ + digest = hashlib.sha256() + eligible = sorted( + cid for cid, chunk in self.file_chunks.items() if self._embedding_dim_matches(chunk.embedding) + ) + for cid in eligible: + digest.update(cid.encode("utf-8")) + digest.update(b"\x00") + digest.update(np.asarray(self.file_chunks[cid].embedding, dtype=np.float16).tobytes()) + return digest.hexdigest() + async def load(self) -> None: """Load chunks via the parent, then attach FAISS state (sidecar or rebuild).""" await super().load() @@ -127,13 +407,43 @@ class FaissLocalFileStore(LocalFileStore): async def _try_load_sidecar(self) -> bool: """Read the binary index plus id-map sidecar. On any mismatch or read error, wipe the partial files so the caller can rebuild from chunks cleanly. + + HNSW construction parameters are validated against the active config: + + - ``M`` is structural: it determines the graph topology (level-0 + neighbors have capacity 2*M). A mismatch means the persisted graph + was built with different connectivity, so we raise to trigger a full + rebuild from the authoritative chunks. + - ``efConstruction`` only affects how edges are formed during ``add()``. + Already-inserted vectors keep their connections regardless, so a + mismatch does not warrant a rebuild. Instead we update the live + index's ``efConstruction`` in place so subsequent insertions honor the + new value. """ if not (self.faiss_path.exists() and self.faiss_idmap_path.exists()): return False try: index = self._faiss.read_index(str(self.faiss_path)) + # Reject legacy / incompatible index types (e.g. a pre-HNSW IndexFlatIP + # sidecar) up front: they load fine and pass the dim/id_map checks but + # lack the ``.hnsw`` attribute, so _set_ef_search would raise + # AttributeError on the first vector_search. Failing here triggers a + # clean rebuild into an IndexHNSWFlat instead. + if not isinstance(index, self._faiss.IndexHNSWFlat): + raise ValueError(f"FAISS index type {type(index).__name__} is not IndexHNSWFlat") if index.d != self._dim: raise ValueError(f"FAISS dim {index.d} != embedding dim {self._dim}") + # Validate HNSW M: it is baked into the serialized graph (level-0 + # neighbor capacity is 2*M, recorded in cum_nneighbor_per_level at + # construction time, so this holds even for an empty index). + # nb_neighbors(0) // 2 recovers the original M, and a mismatch means + # future insertions would extend a graph whose topology does not + # match the active configuration — rebuild from chunks. + persisted_m = index.hnsw.nb_neighbors(0) // 2 + if persisted_m != self.hnsw_m: + raise ValueError( + f"FAISS HNSW M mismatch: persisted={persisted_m}, configured={self.hnsw_m}", + ) async with aiofiles.open(self.faiss_idmap_path, encoding=self.encoding) as f: data = json.loads(await f.read()) id_map = list(data.get("id_map", [])) @@ -150,6 +460,25 @@ class FaissLocalFileStore(LocalFileStore): } if set(live_ids) != expected_ids: raise ValueError("FAISS sidecar live ids do not match persisted chunks") + # A matching live-id set does not imply matching vectors: same-ID + # in-place updates keep the set identical while the embeddings move + # to a new generation. The content digest catches a sidecar that + # fell behind the authoritative JSONL (crash between the two writes + # in dump(), a swallowed sidecar write failure, or externally mixed + # files). Pre-digest sidecars miss the key and rebuild once. + if data.get("digest") != self._chunks_embedding_digest(): + raise ValueError("FAISS sidecar embedding digest does not match persisted chunks") + # efConstruction only affects edge formation during add(); existing + # edges are immutable. Update the live value so new insertions use + # the configured quality without forcing a rebuild of old vectors. + persisted_ef = index.hnsw.efConstruction + if persisted_ef != self.hnsw_ef_construction: + self.logger.info( + f"FAISS efConstruction updated from {persisted_ef} to " + f"{self.hnsw_ef_construction}; existing vectors retain their " + f"connections, new additions will use the new value", + ) + index.hnsw.efConstruction = self.hnsw_ef_construction self._faiss_index = index self._id_map = id_map self._tombstones = tombstones @@ -179,7 +508,16 @@ class FaissLocalFileStore(LocalFileStore): token = uuid4().hex tmp_index = self.faiss_path.with_name(f".{self.faiss_path.name}.{token}.tmp") tmp_idmap = self.faiss_idmap_path.with_name(f".{self.faiss_idmap_path.name}.{token}.tmp") - payload = json.dumps({"id_map": list(self._id_map), "tombstones": sorted(self._tombstones)}) + # The digest binds this sidecar to the chunk generation the JSONL just + # persisted, so a stale sidecar cannot pass load-time validation on the + # live-id set alone. + payload = json.dumps( + { + "id_map": list(self._id_map), + "tombstones": sorted(self._tombstones), + "digest": self._chunks_embedding_digest(), + }, + ) try: self._faiss.write_index(self._faiss_index, str(tmp_index)) async with aiofiles.open(tmp_idmap, "w", encoding=self.encoding) as f: @@ -233,10 +571,11 @@ class FaissLocalFileStore(LocalFileStore): continue if cid in existing and old_text_by_id.get(cid) == chunk.text: continue + # Reaching here means the chunk is new or its text changed; the old + # row (if any) is tombstoned and the fresh vector is re-added. if cid in existing: self._tombstone(cid) - if cid not in existing or old_text_by_id.get(cid) != chunk.text: - to_add.append(chunk) + to_add.append(chunk) if to_add: vectors = np.stack([c.embedding for c in to_add]) @@ -248,55 +587,108 @@ class FaissLocalFileStore(LocalFileStore): paths = [path] if isinstance(path, str) else path nodes = await self.file_graph.get_nodes(paths) deleted_ids = [cid for n in nodes for cid in n.chunk_ids] - await super().delete(path) + await self._delete_nodes(nodes) # reuse resolved nodes; avoids a second get_nodes if self._faiss_index is None: return for cid in deleted_ids: self._tombstone(cid) self._compact_if_needed() + async def _close(self) -> None: + """Stop any in-flight reindex before the parent persists and tears down. + + ``_closing`` is set first so the parent's final ``dump`` cannot submit a new + background reindex (which would leak as an orphan task). + """ + self._closing = True + await self._stop_reindex_worker() + await super()._close() + async def clear(self) -> None: - await super().clear() - self._faiss_index = self._new_index() if self.embedding_store is not None else None - self._id_map = [] - self._id_to_row = {} - self._tombstones.clear() - self.faiss_path.unlink(missing_ok=True) - self.faiss_idmap_path.unlink(missing_ok=True) + # Serialize with dump so a concurrent _write_sidecar cannot re-create the + # sidecar files we are about to unlink, or persist a half-reset index. + # Stop the reindex *inside* the lock: a dump holding the lock can submit a + # fresh reindex via _compact_if_needed, so stopping before acquiring the + # lock would let that worker run against the just-cleared state. The worker + # never takes _faiss_dump_lock, so stopping it while holding the lock cannot + # deadlock. + async with self._faiss_dump_lock: + await self._stop_reindex_worker() + await super().clear() + self._faiss_index = self._new_index() if self.embedding_store is not None else None + self._id_map = [] + self._id_to_row = {} + self._tombstones.clear() + self.faiss_path.unlink(missing_ok=True) + self.faiss_idmap_path.unlink(missing_ok=True) # -- search ----------------------------------------------------------- async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: - if self.embedding_store is None or not query or self._faiss_index is None or self._faiss_index.ntotal == 0: + if ( + self.embedding_store is None + or not query + or limit <= 0 + or self._faiss_index is None + or self._faiss_index.ntotal == 0 + ): return [] + query_embedding = None try: query_embedding = await self.embedding_store.get_embedding(query) except Exception as e: self._disable_embedding(f"search: {type(e).__name__}: {e}") + if query_embedding is None or not self._embedding_dim_matches(query_embedding): + if query_embedding is not None: + self._disable_embedding( + f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", + ) return [] - if query_embedding is None: - return [] - if not self._embedding_dim_matches(query_embedding): - self._disable_embedding( - f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", - ) + + # get_embedding above yielded control; a concurrent clear() drops the + # index to None once embedding is disabled, and a reindex may have swapped + # it out. Re-read after the await before dereferencing ``.ntotal`` so we + # never touch a None/emptied index (mirrors the entry guard). + index = self._faiss_index + if index is None or index.ntotal == 0: return [] q = self._prepare(query_embedding) - ntotal = self._faiss_index.ntotal + ntotal = index.ntotal + + # Small-index shortcut: when ntotal is below the brute-force threshold, + # the HNSW graph would visit nearly every node anyway. Bypass the graph + # and do an exact brute-force scan via index.storage (the IndexFlat that + # backs IndexHNSWFlat). This is faster (no graph traversal overhead) + # and gives 100% recall. Results are in the same (scores, rows) format + # so _collect_hits handles tombstones and filters unchanged. + # + # The threshold is max(1, sqrt(max(limit, ef_search))) * M, where ef_search + # = limit * 5. This ties brute-force coverage to M (graph degree): higher + # M means more neighbors per node and costlier graph traversal, so + # brute-force is preferred for larger indexes. sqrt scales sublinearly + # with the result count because HNSW's advantage grows sublinearly with + # limit — the graph need not be explored much more for larger result sets. + ef_search = limit * 5 + if ntotal < max(1, (max(limit, ef_search)) ** 0.5) * self.hnsw_m: + k = ntotal if search_filter else min(ntotal, limit + len(self._tombstones)) + scores, rows = index.storage.search(q, k) + return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter) if not search_filter: # No filter: simple over-fetch to cover tombstones. k = min(ntotal, limit + len(self._tombstones)) - scores, rows = self._faiss_index.search(q, k) + self._set_ef_search(index, limit) + scores, rows = index.search(q, k) return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter) # With filter: progressively increase k until we collect enough results # or exhaust the entire index. k = min(ntotal, 3 * limit) while True: - scores, rows = self._faiss_index.search(q, k) + self._set_ef_search(index, limit) + scores, rows = index.search(q, k) results = self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter) if len(results) >= limit or k >= ntotal: return results diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index ef3ed749..889a792a 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -450,6 +450,18 @@ class LocalFileStore(BaseFileStore): await self.keyword_index.dump() await self.file_graph.dump() + # -- maintenance ----------------------------------------------------------- + + async def optimize_index(self) -> None: + """Idle-time maintenance: compact the keyword index when present. + + The in-memory chunk map and file graph carry no deferred compaction + debt; the keyword index may hold lazy-deleted docs, so delegate to its + ``optimize_index`` for physical reclaim. + """ + if self.keyword_index: + await self.keyword_index.optimize_index() + # -- CRUD ----------------------------------------------------------------- async def upsert(self, files: list[tuple[FileNode, list[FileChunk]]]) -> None: @@ -539,6 +551,16 @@ class LocalFileStore(BaseFileStore): assert self.file_graph is not None paths = [path] if isinstance(path, str) else path nodes: list[FileNode] = await self.file_graph.get_nodes(paths) + await self._delete_nodes(nodes) + + async def _delete_nodes(self, nodes: list[FileNode]) -> None: + """Delete already-resolved nodes and their chunks. + + Split out so subclasses that need the node list before deletion (e.g. to + capture chunk ids for a vector index) can reuse it instead of querying the + graph a second time. + """ + assert self.file_graph is not None if not nodes: return deleted_chunk_ids = [cid for n in nodes for cid in n.chunk_ids] diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 18061b6b..98650c55 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -68,6 +68,12 @@ jobs: - backend: dream_finish_step file_catalog: dream + optimize_index_cron: + backend: cron + cron: "0 2 * * *" + steps: + - backend: optimize_index_step + auto_dream: backend: base description: "Auto-dream: scan today's day-index and daily notes, globally extract merged units/topics, integrate digest units, write interests.yaml, and persist the dream catalog." @@ -728,6 +734,7 @@ components: file_store: default: backend: local + # backend: faiss store_name: local # embedding_store: default embedding_store: "" diff --git a/reme/steps/index/__init__.py b/reme/steps/index/__init__.py index 79881d08..3e8a16bd 100644 --- a/reme/steps/index/__init__.py +++ b/reme/steps/index/__init__.py @@ -7,6 +7,7 @@ from .draft import AddDraftStep, ReadAllDraftStep from .log_changes import LogChangesStep from .node_search import NodeSearchStep from .init_changes import InitChangesStep +from .optimize_index import OptimizeIndexStep from .search import SearchStep from .search_v2 import SearchV2Step from .traverse import TraverseStep @@ -33,6 +34,7 @@ __all__ = [ "LogChangesStep", "NodeSearchStep", "ReadAllDraftStep", + "OptimizeIndexStep", "SearchStep", "SearchV2Step", "TraverseStep", diff --git a/reme/steps/index/optimize_index.py b/reme/steps/index/optimize_index.py new file mode 100644 index 00000000..df80d8c8 --- /dev/null +++ b/reme/steps/index/optimize_index.py @@ -0,0 +1,16 @@ +"""Idle-time file store maintenance, scheduled off the request path (e.g. cron).""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("optimize_index_step") +class OptimizeIndexStep(BaseStep): + """Call ``file_store.optimize_index()`` so backends can compact derived index state.""" + + async def execute(self): + assert self.context is not None + await self.file_store.optimize_index() + self.context.response.metadata["optimized_index"] = True + self.logger.info(f"[{self.name}] optimized file_store index") + return self.context.response diff --git a/tests/unit/test_faiss_index_maintenance.py b/tests/unit/test_faiss_index_maintenance.py new file mode 100644 index 00000000..15960f18 --- /dev/null +++ b/tests/unit/test_faiss_index_maintenance.py @@ -0,0 +1,356 @@ +"""Tests for FAISS index maintenance: backfill incremental sync and idle-time optimize_index.""" + +# pylint: disable=protected-access + +import asyncio +import os +import tempfile +import time + +import numpy as np +import pytest + +from reme.components.file_store import FaissLocalFileStore, LocalFileStore +from reme.schema import FileChunk, FileNode +from reme.steps.index import OptimizeIndexStep + + +class temp_chdir: + """Temporarily chdir into a test workspace.""" + + def __init__(self, path): + self.path = path + self.old = None + + def __enter__(self): + self.old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self.old) + + +class FakeEmbeddingStore: + """Small deterministic embedding provider used by file-store tests.""" + + dimensions = 2 + max_batch_size = 10 + + def _embed(self, text: str) -> np.ndarray: + if "beta" in text or "fresh" in text: + return np.array([0.0, 1.0], dtype=np.float16) + return np.array([1.0, 0.0], dtype=np.float16) + + async def health_check(self, _timeout: float = 2.0) -> bool: + """Report the fake embedding service as healthy.""" + return True + + async def get_embedding(self, input_text: str, **_kwargs) -> np.ndarray: + """Return a deterministic embedding for a single text.""" + return self._embed(input_text) + + async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]: + """Attach deterministic embeddings to file chunks.""" + for chunk_node in nodes: + chunk_node.embedding = self._embed(chunk_node.text) + return nodes + + +def run(coro): + """Run an async test body.""" + return asyncio.run(coro) + + +def node(path: str) -> FileNode: + """Build a minimal file node.""" + return FileNode(path=path, st_mtime=1.0) + + +def chunk(chunk_id: str, path: str, text: str, **metadata) -> FileChunk: + """Build a minimal file chunk.""" + return FileChunk(id=chunk_id, path=path, text=text, start_line=1, end_line=1, metadata=metadata) + + +def _new_faiss_store(name, **kwargs): + """Construct a FAISS store with embedding disabled at bind time.""" + try: + store = FaissLocalFileStore(name=name, embedding_store="", **kwargs) + except ImportError: + pytest.skip("faiss is not installed") + return store + + +async def _settle_reindex(store, timeout=5.0): + """Wait until no async reindex is pending or in flight.""" + deadline = time.monotonic() + timeout + while store._reindex_event.is_set() or store._reindex_busy: + if time.monotonic() > deadline: + raise AssertionError("async reindex did not settle in time") + await asyncio.sleep(0.005) + + +async def _seed_unembedded_chunk(store: LocalFileStore, chunk_id: str, path: str, text: str) -> None: + """Attach one chunk without a vector, keeping the graph invariant intact.""" + store.file_chunks[chunk_id] = chunk(chunk_id, path, text) + file_node = node(path) + file_node.chunk_ids = [chunk_id] + await store.file_graph.upsert_nodes([file_node]) + + +def test_faiss_backfill_adds_incrementally_without_rebuild(): + """Backfilled vectors are added to the live index; no full rebuild happens + while tombstones stay under the compaction threshold.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_backfill_incr") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.upsert([(node("b.md"), [chunk("b", "b.md", "beta text")])]) + # Same-id text change tombstones the old row (below the 128 floor). + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha updated")])]) + assert store._tombstones == {0} + + await _seed_unembedded_chunk(store, "c", "c.md", "alpha extra") + await store._backfill_missing_embeddings() + + # The new vector is live; the surviving tombstone proves the index + # was extended in place rather than rebuilt. + assert set(store._id_to_row) == {"a", "b", "c"} + assert store._tombstones == {0} + assert store._reindex_worker_task is None + assert {c.id for c in await store.vector_search("alpha", 10, {})} == {"a", "b", "c"} + await store.close() + + run(go()) + + +def test_faiss_backfill_compacts_when_tombstones_cross_threshold(): + """Tombstone pressure at backfill time still triggers a full rebuild.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_backfill_compact") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha updated")])]) + assert store._tombstones == {0} + + # Lower the threshold only now, so the upsert above did not compact. + store.max_tombstones = 1 + await _seed_unembedded_chunk(store, "c", "c.md", "alpha extra") + await store._backfill_missing_embeddings() + + assert store._tombstones == set() + assert set(store._id_to_row) == {"a", "c"} + assert {c.id for c in await store.vector_search("alpha", 10, {})} == {"a", "c"} + await store.close() + + run(go()) + + +def test_faiss_backfill_mass_delta_uses_async_reindex(): + """With async_reindex, a delta that rivals the live rows (initial mass + backfill into an empty index) goes through the background worker.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_backfill_async_mass", async_reindex=True) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + await _seed_unembedded_chunk(store, "a", "a.md", "alpha text") + await _seed_unembedded_chunk(store, "b", "b.md", "beta text") + await store._backfill_missing_embeddings() + + assert store._reindex_worker_task is not None # routed off-loop + await _settle_reindex(store) + assert set(store._id_to_row) == {"a", "b"} + assert [c.id for c in await store.vector_search("alpha", 5, {})][0] == "a" + await store.close() + + run(go()) + + +def test_faiss_backfill_small_delta_adds_inline_in_async_mode(): + """With async_reindex, a small delta is added inline without a worker.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_backfill_async_small", async_reindex=True) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.upsert([(node("b.md"), [chunk("b", "b.md", "beta text")])]) + + await _seed_unembedded_chunk(store, "c", "c.md", "alpha extra") + await store._backfill_missing_embeddings() + + # 1 new vector < 2 live rows -> inline add, no background worker. + assert store._reindex_worker_task is None + assert set(store._id_to_row) == {"a", "b", "c"} + assert {c.id for c in await store.vector_search("alpha", 10, {})} >= {"a", "c"} + await store.close() + + run(go()) + + +# -- optimize_index (idle-time maintenance) ------------------------------------- + + +def test_tombstone_threshold_scales_and_honors_override(): + """Threshold math: full vs half scale, 128 floor, and max_tombstones override.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_threshold") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = type("FakeIndex", (), {"ntotal": 1000})() + + assert store._tombstone_threshold() == 300 # 0.3 * 1000 + assert store._tombstone_threshold(scale=0.5) == 150 # half the write-path bar + store._faiss_index = type("FakeIndex", (), {"ntotal": 100})() + assert store._tombstone_threshold(scale=0.5) == 128 # floor dominates + + store.max_tombstones = 7 + assert store._tombstone_threshold() == 7 + assert store._tombstone_threshold(scale=0.5) == 7 # fixed override ignores scale + + store._faiss_index = None # avoid persisting the fake index on close + await store.close() + + run(go()) + + +def test_local_store_optimize_index_is_noop(): + """LocalFileStore.optimize_index() completes without touching store state.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = LocalFileStore(name="t_local_optimize", embedding_store="") + await store.start() + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + + await store.optimize_index() + + assert set(store.file_chunks) == {"a"} + assert [c.id for c in await store.keyword_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_faiss_optimize_index_noop_below_threshold(): + """optimize_index() keeps tombstones when they are under the half bar.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_optimize_noop") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha updated")])]) + assert store._tombstones == {0} + + await store.optimize_index() # 1 tombstone <= 128 floor -> untouched + + assert store._tombstones == {0} + assert store._reindex_worker_task is None + await store.close() + + run(go()) + + +def test_faiss_optimize_index_compacts_inline_when_sync(): + """optimize_index() rebuilds inline once tombstones exceed the (overridden) bar.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_optimize_sync") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + files = [(node(f"n{i}.md"), [chunk(f"c{i}", f"n{i}.md", "alpha text")]) for i in range(4)] + await store.upsert(files) + updated = [(node(f"n{i}.md"), [chunk(f"c{i}", f"n{i}.md", f"alpha v2 {i}")]) for i in range(3)] + await store.upsert(updated) + assert len(store._tombstones) == 3 + + store.max_tombstones = 2 # lower the bar only for optimize_index + await store.optimize_index() + + assert store._tombstones == set() + assert set(store._id_to_row) == {"c0", "c1", "c2", "c3"} + assert store._reindex_worker_task is None # inline path + assert len(await store.vector_search("alpha", 10, {})) == 4 + await store.close() + + run(go()) + + +def test_faiss_optimize_index_uses_worker_when_async(): + """optimize_index() submits the rebuild to the background worker under async_reindex.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_optimize_async", async_reindex=True) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + files = [(node(f"n{i}.md"), [chunk(f"c{i}", f"n{i}.md", "alpha text")]) for i in range(4)] + await store.upsert(files) + updated = [(node(f"n{i}.md"), [chunk(f"c{i}", f"n{i}.md", f"alpha v2 {i}")]) for i in range(3)] + await store.upsert(updated) + assert len(store._tombstones) == 3 + + store.max_tombstones = 2 + await store.optimize_index() + + assert store._reindex_worker_task is not None # routed off-loop + await _settle_reindex(store) + assert store._tombstones == set() + assert set(store._id_to_row) == {"c0", "c1", "c2", "c3"} + await store.close() + + run(go()) + + +def test_optimize_index_step_calls_file_store_optimize_index(): + """The cron-facing step delegates to file_store.optimize_index() and reports success.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = LocalFileStore(name="t_optimize_step", embedding_store="") + await store.start() + + calls = [] + original_optimize = store.optimize_index + + async def counting_optimize(): + calls.append(True) + await original_optimize() + + store.optimize_index = counting_optimize + step = OptimizeIndexStep(file_store=store) + await step() + + assert calls == [True] + assert step.context.response.metadata["optimized_index"] is True + await store.close() + + run(go()) diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index f3c48b9e..e8ab7cdb 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -1,12 +1,15 @@ """Regression tests for LocalFileStore / FaissLocalFileStore consistency.""" -# pylint: disable=protected-access +# pylint: disable=protected-access,too-many-lines import asyncio import base64 +import datetime import json import os import tempfile +import threading +import time import numpy as np import pytest @@ -925,3 +928,614 @@ def test_faiss_date_filter_progressive_recall(): await store.close() run(go()) + + +def test_faiss_vector_search_survives_concurrent_index_drop(): + """vector_search returns [] (not AttributeError) if a concurrent clear() + drops _faiss_index while the query embedding is being computed (TOCTOU).""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_search_drop") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + + real_get_embedding = store.embedding_store.get_embedding + + async def dropping_get_embedding(text, **kwargs): + emb = await real_get_embedding(text, **kwargs) + # Simulate a concurrent clear() landing during the await. + store._faiss_index = None + return emb + + store.embedding_store.get_embedding = dropping_get_embedding + + # Must return [] rather than raising AttributeError on None.ntotal. + assert await store.vector_search("alpha", 5, {}) == [] + await store.close() + + run(go()) + + +# -- HNSW construction parameter persistence tests --------------------------- + + +def test_faiss_rebuilds_on_hnsw_m_mismatch(): + """Reopening with a different hnsw_m must rebuild: M is structural and the + persisted graph topology cannot serve the new configuration.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + # Phase 1: build and persist with hnsw_m=8. + store_a = _new_faiss_store("t_faiss_m_mismatch", hnsw_m=8) + await store_a.start() + store_a.embedding_store = FakeEmbeddingStore() + store_a._faiss_index = store_a._new_index() + await store_a.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + assert store_a._faiss_index.hnsw.nb_neighbors(0) // 2 == 8 + await store_a.close() # _close() → dump() writes sidecar + + # Phase 2: reopen the same workspace with hnsw_m=32. + store_b = _new_faiss_store("t_faiss_m_mismatch", hnsw_m=32) + await store_b.start() # loads chunks; FAISS skipped (embedding_store is None) + store_b.embedding_store = FakeEmbeddingStore() + + # The sidecar was built with M=8; config says 32 → reject and rebuild. + assert await store_b._try_load_sidecar() is False + assert not store_b.faiss_path.exists() # sidecar wiped on rejection + store_b._rebuild_index() + + # Rebuilt index honors the new M. + assert store_b._faiss_index.hnsw.nb_neighbors(0) // 2 == 32 + assert [c.id for c in await store_b.vector_search("alpha", 5, {})] == ["a"] + await store_b.close() + + run(go()) + + +def test_faiss_rebuilds_on_hnsw_m_mismatch_empty_index(): + """An empty sidecar must not bypass the M check: M is baked into the + serialized graph, so later insertions would use stale connectivity.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + # Phase 1: persist an *empty* index built with hnsw_m=8. + store_a = _new_faiss_store("t_faiss_m_mismatch_empty", hnsw_m=8) + await store_a.start() + store_a.embedding_store = FakeEmbeddingStore() + store_a._faiss_index = store_a._new_index() + assert store_a._faiss_index.ntotal == 0 + await store_a.close() # _close() → dump() writes sidecar + + # Phase 2: reopen with hnsw_m=32; the empty M=8 sidecar must be rejected. + store_b = _new_faiss_store("t_faiss_m_mismatch_empty", hnsw_m=32) + await store_b.start() + store_b.embedding_store = FakeEmbeddingStore() + + assert await store_b._try_load_sidecar() is False + assert not store_b.faiss_path.exists() # sidecar wiped on rejection + store_b._rebuild_index() + + # New insertions use the configured M, not the stale persisted one. + await store_b.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + assert store_b._faiss_index.hnsw.nb_neighbors(0) // 2 == 32 + await store_b.close() + + run(go()) + + +def test_faiss_rejects_stale_sidecar_after_partial_dump(): + """A sidecar whose vectors belong to an older chunk generation must be + rejected even when the live ID set is unchanged (same-ID in-place update). + + Reproduces the crash window inside dump(): the authoritative chunk JSONL + is written, then the process dies before _write_sidecar(). On restart the + stale sidecar passes every shape check (type/dim/M/rows/live-ID set); only + the embedding content digest can tell its vectors are a generation behind. + """ + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + # t0: consistent state on disk (c1 = alpha -> [1,0]). + store_a = _new_faiss_store("t_faiss_stale_sidecar") + await store_a.start() + store_a.embedding_store = FakeEmbeddingStore() + store_a._faiss_index = store_a._new_index() + await store_a.upsert([(node("a.md"), [chunk("c1", "a.md", "alpha topic")])]) + await store_a.dump() + + # t1: same-ID in-place update (c1 = beta -> [0,1]). + await store_a.upsert([(node("a.md"), [chunk("c1", "a.md", "beta topic")])]) + + # t2: simulate a crash between the two writes in dump(): only the + # parent's JSONL write lands; the sidecar stays at the alpha + # generation. (No close() -- the process is presumed dead.) + await LocalFileStore.dump(store_a) + + # t3: restart. The stale sidecar must be rejected by the digest. + store_b = _new_faiss_store("t_faiss_stale_sidecar") + await store_b.start() + store_b.embedding_store = FakeEmbeddingStore() + assert store_b.file_chunks["c1"].text == "beta topic" + + assert await store_b._try_load_sidecar() is False + assert not store_b.faiss_path.exists() # sidecar wiped on rejection + store_b._rebuild_index() + + # t4: the rebuilt index serves the current generation: a beta query + # scores ~1.0 instead of 0.0 against the stale alpha vector. + results = await store_b.vector_search("beta", 5, {}) + assert [c.id for c in results] == ["c1"] + assert results[0].scores["vector"] > 0.5 + await store_b.close() + + run(go()) + + +def test_faiss_ef_construction_hot_update_without_rebuild(): + """Reopening with a different hnsw_ef_construction must NOT rebuild: it only + affects future add() edge formation; the live efConstruction is hot-updated.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + # Phase 1: build and persist with efConstruction=40. + store_a = _new_faiss_store("t_faiss_ef_hot", hnsw_ef_construction=40) + await store_a.start() + store_a.embedding_store = FakeEmbeddingStore() + store_a._faiss_index = store_a._new_index() + await store_a.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + assert store_a._faiss_index.hnsw.efConstruction == 40 + await store_a.close() + + # Phase 2: reopen with efConstruction=128 (M unchanged). + store_b = _new_faiss_store("t_faiss_ef_hot", hnsw_ef_construction=128) + await store_b.start() + store_b.embedding_store = FakeEmbeddingStore() + + # Sidecar loads successfully — M matches, only efConstruction differs. + assert await store_b._try_load_sidecar() is True + assert store_b.faiss_path.exists() # sidecar retained (no rebuild) + + # efConstruction was hot-updated to the new config value. + assert store_b._faiss_index.hnsw.efConstruction == 128 + + # Search still works on the loaded (not rebuilt) index. + assert [c.id for c in await store_b.vector_search("alpha", 5, {})] == ["a"] + await store_b.close() + + run(go()) + + +def test_faiss_small_index_uses_brute_force_scan(): + """Below the brute-force threshold, vector_search does an exact scan via + index.storage; the HNSW path (_set_ef_search) is used only when large enough.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_brute_force") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + # 80 chunks with valid date paths so we can test filtered search too. + # With M=32 (default): limit=5 -> threshold=sqrt(25)*32=160 > 80 + # -> brute-force. limit=1 -> threshold=sqrt(5)*32~=72 < 80 -> HNSW. + base = datetime.date(2026, 1, 1) + files = [] + for i in range(80): + d = base + datetime.timedelta(days=i) + path = f"daily/{d.isoformat()}/note.md" + files.append((node(path), [chunk(f"c{i}", path, "alpha text")])) + await store.upsert(files) + + # _set_ef_search is only called on the HNSW path. Spying on it + # tells us which branch vector_search took. + ef_calls: list[int] = [] + original_set_ef = store._set_ef_search + + def spy_set_ef(idx, lim): + ef_calls.append(lim) + original_set_ef(idx, lim) + + store._set_ef_search = spy_set_ef + + # ntotal=80, limit=5 -> threshold=sqrt(25)*32=160 -> 80 < 160 -> brute-force. + results = await store.vector_search("alpha", 5, {}) + assert len(results) == 5 + assert not ef_calls # brute-force path taken + + # Brute-force + filter: scans all vectors, _collect_hits filters. + filt = {"start_date": "2026-01-15", "end_date": "2026-01-17"} + results = await store.vector_search("alpha", 5, filt) + assert {r.id for r in results} == {"c14", "c15", "c16"} + assert not ef_calls # still brute-force + + # ntotal=80, limit=1 -> threshold=sqrt(5)*32~=72 -> 80 >= 72 -> HNSW graph search. + results = await store.vector_search("alpha", 1, {}) + assert len(results) == 1 + assert len(ef_calls) >= 1 # HNSW path taken, efSearch was set + + store._set_ef_search = original_set_ef + await store.close() + + run(go()) + + +# -- Async reindex tests ----------------------------------------------------- + + +def _new_faiss_store(name, **kwargs): + """Construct a started FAISS store with a fake embedding backend and empty index.""" + try: + store = FaissLocalFileStore(name=name, embedding_store="", **kwargs) + except ImportError: + pytest.skip("faiss is not installed") + return store + + +async def _settle_reindex(store, timeout=5.0): + """Wait until no async reindex is pending or in flight.""" + deadline = time.monotonic() + timeout + while store._reindex_event.is_set() or store._reindex_busy: + if time.monotonic() > deadline: + raise AssertionError("async reindex did not settle in time") + await asyncio.sleep(0.005) + + +def test_faiss_async_reindex_disabled_by_default(): + """Default store keeps the synchronous compaction path: no background worker.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_sync_default", max_tombstones=2) + assert store.async_reindex is False + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + files = [(node(f"n{i}.md"), [chunk(f"c{i}", f"n{i}.md", "alpha text")]) for i in range(4)] + await store.upsert(files) + await store.delete([f"n{i}.md" for i in range(3)]) + + # Synchronous rebuild ran inline; no background worker was created. + assert store._reindex_worker_task is None + assert store._tombstones == set() + assert set(store._id_to_row) == {"c3"} + assert [c.id for c in await store.vector_search("alpha", 10, {})] == ["c3"] + await store.close() + + run(go()) + + +def test_faiss_async_reindex_triggered_by_compaction(): + """Crossing the tombstone threshold submits a background rebuild whose result + matches a synchronous rebuild (deleted ids gone, tombstones cleared).""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_async_compact", async_reindex=True, max_tombstones=2) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + files = [(node(f"n{i}.md"), [chunk(f"c{i}", f"n{i}.md", "alpha text")]) for i in range(4)] + await store.upsert(files) + assert store._reindex_worker_task is None # below threshold, nothing submitted yet + + await store.delete([f"n{i}.md" for i in range(3)]) # 3 tombstones >= 2 -> submit + assert store._reindex_worker_task is not None + await _settle_reindex(store) + + assert set(store._id_to_row) == {"c3"} + assert store._tombstones == set() + assert [c.id for c in await store.vector_search("alpha", 10, {})] == ["c3"] + await store.close() + + run(go()) + + +def test_faiss_async_reindex_no_lost_writes_during_build(): + """Writes that land while an async rebuild is in flight are not lost: a + follow-up rebuild folds them in (eventual consistency).""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_async_nolost", async_reindex=True) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + + started = threading.Event() + release = threading.Event() + real_build = store._build_index_blocking + + def gated_build(dim, vectors): + started.set() + release.wait() + return real_build(dim, vectors) + + store._build_index_blocking = gated_build + + store._submit_reindex() # snapshot == {a} + while not started.is_set(): + await asyncio.sleep(0.005) + + # Concurrent writes on the live index while the build is blocked: + await store.upsert([(node("b.md"), [chunk("b", "b.md", "beta text")])]) # brand new + await store.upsert([(node("a.md"), [chunk("a", "a.md", "beta text")])]) # changed text + + release.set() + await _settle_reindex(store) + + # After the follow-up rebuild the index reflects both concurrent writes; + # the changed chunk now embeds as "beta". + assert set(store._id_to_row) == {"a", "b"} + assert {c.id for c in await store.vector_search("beta", 5, {})} == {"a", "b"} + await store.close() + + run(go()) + + +def test_faiss_async_reindex_single_worker_coalesces(): + """Only one reindex runs at a time; repeated submissions coalesce and the + worker stays a single long-lived task.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_async_single", async_reindex=True) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + + started = threading.Event() + release = threading.Event() + active = {"n": 0} + peak = {"n": 0} + real_build = store._build_index_blocking + + def gated_build(dim, vectors): + active["n"] += 1 + peak["n"] = max(peak["n"], active["n"]) + started.set() + release.wait() + try: + return real_build(dim, vectors) + finally: + active["n"] -= 1 + + store._build_index_blocking = gated_build + + store._submit_reindex() + while not started.is_set(): + await asyncio.sleep(0.005) + worker = store._reindex_worker_task + + # Several more submissions while the first build is blocked collapse into + # the flag rather than spawning parallel builds or a second worker. + for _ in range(5): + store._submit_reindex() + assert store._reindex_worker_task is worker + + release.set() + await _settle_reindex(store) + + assert peak["n"] == 1 # never two builds at once + assert set(store._id_to_row) == {"a"} + await store.close() + + run(go()) + + +def test_faiss_async_reindex_cancelled_on_close(): + """close() stops an in-flight reindex without hanging on the worker thread.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_async_close", async_reindex=True) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + + started = threading.Event() + + def slow_build(_dim, _vectors): + started.set() + while not store._closing: + time.sleep(0.005) + + store._build_index_blocking = slow_build + store._submit_reindex() + while not started.is_set(): + await asyncio.sleep(0.005) + + await store.close() # sets _closing, cancels the worker; the build thread exits + assert store._reindex_worker_task is None + + run(go()) + + +def test_faiss_close_does_not_leave_orphan_reindex(): + """The final dump in _close() must not submit a background reindex.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_close_orphan", async_reindex=True, max_tombstones=2) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + files = [(node(f"n{i}.md"), [chunk(f"c{i}", f"n{i}.md", "alpha text")]) for i in range(4)] + await store.upsert(files) + + # Block the build so the reindex is still in flight (tombstones uncompacted) + # when close runs. + started = threading.Event() + + def slow_build(_dim, _vectors): + started.set() + while not store._closing: + time.sleep(0.005) + + store._build_index_blocking = slow_build + await store.delete([f"n{i}.md" for i in range(3)]) # 3 tombstones >= 2 -> submit + while not started.is_set(): + await asyncio.sleep(0.005) + assert len(store._tombstones) >= store.max_tombstones # not compacted yet + + # Closing cancels the in-flight worker; the _closing guard stops the + # final dump from submitting an orphan reindex. + await store.close() + assert store._reindex_worker_task is None + + run(go()) + + +def test_faiss_async_close_persists_stale_snapshot_on_same_id_update(): + """Regression: close() during a follow-up rebuild can persist a stale snapshot. + + Bug: async rebuild snapshots alpha; "a" is updated to beta; the first build + swaps stale alpha back; close() cancels the follow-up rebuild before it can + swap beta; the stale alpha index is persisted and accepted on reopen. + + Expected after fix: searching with the beta vector scores ~1.0. + """ + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_stale_close", async_reindex=True) + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + + # Chunk "a" with alpha text -> embedding [1.0, 0.0]. + await store.upsert([(node("note.md"), [chunk("a", "note.md", "alpha text")])]) + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + + first_started = threading.Event() + release_first = threading.Event() + second_started = threading.Event() + real_build = store._build_index_blocking + build_count = {"n": 0} + + def gated_build(dim, vectors): + build_count["n"] += 1 + if build_count["n"] == 1: + # First build: hold the alpha snapshot. + first_started.set() + release_first.wait() + else: + # Follow-up build: signal started, then hold until close(). + second_started.set() + while not store._closing: + time.sleep(0.005) + return real_build(dim, vectors) + + store._build_index_blocking = gated_build + + # Step 1: submit the first rebuild (snapshot == {a: alpha}). + store._submit_reindex() + while not first_started.is_set(): + await asyncio.sleep(0.005) + + await store.upsert([(node("note.md"), [chunk("a", "note.md", "beta text")])]) + release_first.set() + while not second_started.is_set(): + await asyncio.sleep(0.005) + await store.close() + assert store._reindex_worker_task is None + + # Step 6: reopen and verify the persisted state. + reopened = _new_faiss_store("t_faiss_stale_close", async_reindex=True) + await reopened.start() + reopened.embedding_store = FakeEmbeddingStore() + + # The authoritative chunk JSONL has "beta text". + assert reopened.file_chunks["a"].text == "beta text" + + # The stale sidecar is accepted because the ID set is unchanged. + assert await reopened._try_load_sidecar() is True + + results = await reopened.vector_search("beta", 5, {}) + assert len(results) == 1 + assert results[0].id == "a" + score = results[0].scores["vector"] + assert score > 0.5, ( + f"BUG: persisted FAISS index has stale alpha vector; " f"beta query scored {score:.4f} instead of ~1.0" + ) + await reopened.close() + + run(go()) + + +def test_faiss_clear_waits_for_in_flight_dump(): + """clear() serializes with dump() through _faiss_dump_lock.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_clear_lock") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + + first_started = asyncio.Event() + release = asyncio.Event() + original_write_sidecar = store._write_sidecar + + async def blocking_write_sidecar(): + first_started.set() + await release.wait() + + store._write_sidecar = blocking_write_sidecar + dump_task = asyncio.create_task(store.dump()) + await first_started.wait() + + clear_task = asyncio.create_task(store.clear()) + await asyncio.sleep(0.02) + assert not clear_task.done() # blocked on _faiss_dump_lock held by dump + + release.set() + await asyncio.gather(dump_task, clear_task) + assert store._id_to_row == {} + + store._write_sidecar = original_write_sidecar + await store.close() + + run(go()) + + +def test_faiss_delete_queries_graph_once(): + """delete() resolves nodes once and reuses them (no redundant get_nodes).""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_faiss_store("t_faiss_delete_once") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store._faiss_index = store._new_index() + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + + calls: list = [] + original_get_nodes = store.file_graph.get_nodes + + async def counting_get_nodes(paths=None): + calls.append(paths) + return await original_get_nodes(paths) + + store.file_graph.get_nodes = counting_get_nodes + await store.delete("a.md") + + assert len(calls) == 1 + assert "a" not in store._id_to_row + + store.file_graph.get_nodes = original_get_nodes + await store.close() + + run(go())