feat(faiss_file_store): upgrade FAISS to HNSW index with async reindex (#390)

* feat(file_store): upgrade FAISS to HNSW index with async reindex and path constraint

- Replace IndexFlatIP with IndexHNSWFlat for better recall/speed tradeoff
- Add dynamic efSearch (limit * 5) scaled to query request size
- Add async_reindex option: background rebuild with generation-based invalidation
- Extract _delete_nodes() in LocalFileStore for subclass reuse
- Add unit tests for file store consistency

* fix: resolve pylint warnings in faiss store and test file

* refactor(file_store): replace generation-based reindex with event-flag worker

- Replace _reindex_generation/lock/task with a single long-lived worker
  coroutine consuming an asyncio.Event flag; repeated submissions coalesce
- Use local index reference in vector_search to avoid TOCTOU on self._faiss_index
- Pass index explicitly to _set_ef_search for consistency
- Track _index_writes to re-arm reindex after concurrent writes
- Update tests to match new internal API

* fix: resolve pylint too-many-return-statements and implicit-booleaness warnings

* feat(file_store): add refine maintenance hook and incremental embedding backfill

- Add refine() idle-time maintenance hook to BaseFileStore/LocalFileStore
- FaissLocalFileStore: incremental vector add on backfill instead of full rebuild
- Dynamic tombstone compaction threshold scaled by index size
- Add RefineStoreStep with daily cron job (refine_store_cron)
- Enable faiss backend and embedding_store by default in default.yaml
- Add unit tests for faiss index maintenance

* chore(deps): promote faiss-cpu to core dependencies

faiss backend is now the default file_store, so faiss-cpu moves from
the optional [core] extra to the base dependencies list.

* feat: rename refine_store to optimize_index and add vecdb_path_constraint

- Rename refine_store step to optimize_index with cron job scheduling
- Add vecdb_path_constraint to file_store components
- Update default.yaml with optimize_index_cron and faiss backend comment
- Update memory_search docs (en/zh) for FAISS vector management
- Update unit tests for index maintenance

* feat(faiss): add embedding digest to reject stale sidecar after partial dump

Add _chunks_embedding_digest() that computes an order-independent SHA-256
over (chunk_id, float16 embedding) pairs. The digest is written into the
idmap sidecar at dump time and verified at load time. A mismatch means the
sidecar vectors belong to a different chunk generation than the authoritative
JSONL — detectable even when the live-ID set is unchanged (same-ID in-place
update crash window).

Add test_faiss_rejects_stale_sidecar_after_partial_dump reproducing the
crash-between-writes scenario and asserting digest-based rejection.

Compress verbose docstrings/comments in existing tests for pylint line
budget.

---------

Co-authored-by: sa-buc <jiangniurou.xyf@dail-algo011164204033.ET135>
This commit is contained in:
xyf2020 2026-07-27 19:54:38 +08:00 committed by GitHub
parent f34dcdb09b
commit 4eb2adf961
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1486 additions and 32 deletions

1
.gitignore vendored
View file

@ -75,3 +75,4 @@ benchmark/results/
# integration tests outputs
tests/integration/logs/
daily/

View file

@ -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

View file

@ -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 场景下召回效率更高。
## 怎么搜索

View file

@ -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.
"""

View file

@ -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

View file

@ -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]

View file

@ -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: ""

View file

@ -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",

View file

@ -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

View file

@ -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())

View file

@ -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())