From ad4f23e4dc00bc77418ae5d5c5e4e91381c17743 Mon Sep 17 00:00:00 2001 From: lichen2015 Date: Wed, 5 Aug 2026 22:08:30 +0800 Subject: [PATCH] feat(file_store): add ZvecLocalFileStore backend (#410) * feat(file_store): add ZvecLocalFileStore backend - Implement ZvecLocalFileStore with native zvec collection for ANN search. - Keep JSONL chunks as the source of truth; rebuild collection from chunks when sidecar digest/dimension/HNSW M mismatch is detected. - Add dedicated unit tests in tests/unit/test_zvec_file_store.py. - Parametrize existing file_store consistency tests to cover both LocalFileStore and ZvecLocalFileStore. - Register the new backend in reme/components/file_store/__init__.py. * fix(file_store): fix zvec collection sync and content validation, declare zvec dependency --- pyproject.toml | 1 + reme/components/file_store/__init__.py | 2 + .../file_store/zvec_local_file_store.py | 476 ++++++++++++++ tests/unit/test_file_store_consistency.py | 62 +- tests/unit/test_zvec_file_store.py | 614 ++++++++++++++++++ 5 files changed, 1141 insertions(+), 14 deletions(-) create mode 100644 reme/components/file_store/zvec_local_file_store.py create mode 100644 tests/unit/test_zvec_file_store.py diff --git a/pyproject.toml b/pyproject.toml index 2e1efc6f..cba14f77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ core = [ "openai-codex>=0.144.4", "pproxy>=2.7.9,<3", "faiss-cpu>=1.13.2", + "zvec>=0.6.0", "jieba>=0.42.1", "rjieba>=0.2.1", "neo4j>=6.2.0", diff --git a/reme/components/file_store/__init__.py b/reme/components/file_store/__init__.py index 61e9b14f..116330f8 100644 --- a/reme/components/file_store/__init__.py +++ b/reme/components/file_store/__init__.py @@ -8,9 +8,11 @@ shipped today. from .base_file_store import BaseFileStore from .faiss_local_file_store import FaissLocalFileStore from .local_file_store import LocalFileStore +from .zvec_local_file_store import ZvecLocalFileStore __all__ = [ "BaseFileStore", "FaissLocalFileStore", "LocalFileStore", + "ZvecLocalFileStore", ] diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py new file mode 100644 index 00000000..80129eed --- /dev/null +++ b/reme/components/file_store/zvec_local_file_store.py @@ -0,0 +1,476 @@ +"""Zvec-backed file store: chunk JSONL stays authoritative; a zvec collection replaces the linear vector scan.""" + +import hashlib +import json +import shutil +import time +from uuid import uuid4 + +import aiofiles +import numpy as np + +from .local_file_store import LocalFileStore +from ..component_registry import R +from ...schema import FileChunk, FileNode + +# Batch size for bulk inserts during a rebuild. +_ZVEC_INSERT_BATCH_SIZE = 1024 +# Batch size for the startup fetch that verifies persisted collection contents. +_ZVEC_VERIFY_BATCH_SIZE = 1024 + + +@R.register("zvec") +class ZvecLocalFileStore(LocalFileStore): + """LocalFileStore variant whose vector_search is backed by a zvec collection. + + Chunk persistence is unchanged (JSONL, owned by the parent); the zvec + collection only stores ``(chunk_id, embedding)`` pairs and serves ANN + queries. ``self.file_chunks`` remains the source of truth: if the + collection directory or its digest sidecar is missing or stale, the + collection is rebuilt from the chunks. + + zvec (https://zvec.org) is an in-process vector database, so unlike the + FAISS backend there is no tombstone bookkeeping: documents are updated and + removed natively via ``Collection.upsert`` / ``Collection.delete``, and + the collection persists itself inside its own directory. A small JSON + sidecar records an order-independent digest of the live + ``(chunk_id, embedding)`` set at dump time; on load a digest mismatch + (crash between the chunk dump and the collection flush, externally mixed + files, model change) triggers a clean rebuild, and the opened collection is + additionally verified against the ids and vectors it should hold so a + damaged or externally modified collection is rebuilt instead of served. + + HNSW parameters (``hnsw_m``, ``hnsw_ef_construction``) are baked into the + collection schema at creation time; changing ``hnsw_m`` on an existing + store triggers a rebuild so the graph topology matches the configuration. + + ``optimize_index`` maps directly onto ``Collection.optimize()``, zvec's + idle-time index compaction hook. + + zvec is imported lazily inside ``__init__`` so that merely importing this + module does not require the optional dependency; the backend targets the + zvec version declared in ``pyproject.toml`` (``zvec>=0.6.0``). + """ + + def __init__( + self, + hnsw_m: int = 32, + hnsw_ef_construction: int = 64, + **kwargs, + ): + super().__init__(**kwargs) + self._zvec = self._import_zvec() + self.hnsw_m = hnsw_m + self.hnsw_ef_construction = hnsw_ef_construction + self.zvec_path = self.component_metadata_path / f"zvec_index_{self.name}_{self.store_version}" + self.zvec_sidecar_path = self.component_metadata_path / f"zvec_sidecar_{self.name}_{self.store_version}.json" + self._collection = None # zvec.Collection | None + self._indexed_ids: set[str] = set() # chunk ids currently present in the collection + + @staticmethod + def _import_zvec(): + try: + import zvec + except ImportError as e: + raise ImportError( + "zvec is required for ZvecLocalFileStore. Install with `pip install 'zvec>=0.6.0'`.", + ) from e + return zvec + + # -- helpers ---------------------------------------------------------- + + @property + def _dim(self) -> int: + return self.embedding_store.dimensions if self.embedding_store is not None else 0 + + def _collection_schema(self): + """Vector-only schema: chunk data lives in the JSONL, zvec is pure ANN.""" + zvec = self._zvec + return zvec.CollectionSchema( + name=self.name, + fields=[], + vectors=[ + zvec.VectorSchema( + name="embedding", + data_type=zvec.DataType.VECTOR_FP32, + dimension=self._dim, + index_param=zvec.HnswIndexParam( + metric_type=zvec.MetricType.COSINE, + m=self.hnsw_m, + ef_construction=self.hnsw_ef_construction, + ), + ), + ], + ) + + def _create_collection(self): + """Create a fresh (empty) collection, replacing any directory on disk.""" + # Release any open handle first: zvec holds an in-process lock on the + # collection directory, so the old object must be dropped before the + # directory is wiped and re-created. + self._collection = None + if self.zvec_path.exists(): + shutil.rmtree(self.zvec_path, ignore_errors=True) + self._indexed_ids = set() + return self._zvec.create_and_open(path=str(self.zvec_path), schema=self._collection_schema()) + + def _to_doc(self, chunk: FileChunk): + """Build a zvec Doc carrying only the id and the float32 vector.""" + vector = np.asarray(chunk.embedding, dtype=np.float32).tolist() + return self._zvec.Doc(id=chunk.id, vectors={"embedding": vector}) + + @staticmethod + def _vector_fingerprint(vector) -> bytes: + """Canonical bytes of a vector, comparable across an insert/fetch round trip. + + Vectors are inserted as float32 and hashed as float16 (the JSONL + serialization dtype), so normalizing through float32 then float16 makes + an in-memory chunk embedding and the vector read back from the + collection byte-identical whenever they represent the same value. + """ + if vector is None: + return b"" + return np.asarray(np.asarray(vector, dtype=np.float32), dtype=np.float16).tobytes() + + def _upsert_docs(self, chunks: list[FileChunk]) -> None: + if not chunks or self._collection is None: + return + self._collection.upsert([self._to_doc(c) for c in chunks]) + self._indexed_ids.update(c.id for c in chunks) + + def _delete_docs(self, chunk_ids: list[str]) -> None: + if self._collection is None: + return + stale = [cid for cid in chunk_ids if cid in self._indexed_ids] + if not stale: + return + self._collection.delete(ids=stale) + self._indexed_ids.difference_update(stale) + + def _rebuild_collection(self) -> None: + """Rebuild the zvec collection from self.file_chunks (the source of truth).""" + self._collection = self._create_collection() + chunks = [c for c in self.file_chunks.values() if self._embedding_dim_matches(c.embedding)] + for start in range(0, len(chunks), _ZVEC_INSERT_BATCH_SIZE): + batch = chunks[start : start + _ZVEC_INSERT_BATCH_SIZE] + self._collection.insert([self._to_doc(c) for c in batch]) + self._indexed_ids.update(c.id for c in batch) + self.logger.info(f"{self.name}: rebuilt zvec collection with {len(chunks)} vectors at {self.zvec_path}") + + async def _after_embedding_backfill(self) -> None: + """Make newly backfilled vectors visible to zvec. + + zvec upserts are incremental by nature, so backfilled vectors are added + directly to the live collection; no rebuild or tombstone accounting is + needed. + """ + if self.embedding_store is None or self._dim == 0: + return + if self._collection is None: + self._rebuild_collection() + return + to_add = [ + chunk + for cid, chunk in self.file_chunks.items() + if cid not in self._indexed_ids and self._embedding_dim_matches(chunk.embedding) + ] + self._upsert_docs(to_add) + + # -- maintenance ------------------------------------------------------ + + async def optimize_index(self) -> None: + """Idle-time maintenance: delegate to zvec's native index optimization. + + ``Collection.optimize()`` merges staged segments into the persistent + HNSW index, the direct analogue of a tombstone compaction pass. + """ + await super().optimize_index() + if self._collection is None: + return + try: + self._collection.optimize() + self.logger.info(f"{self.name}: zvec collection optimized") + except Exception as e: + self.logger.exception(f"{self.name}: zvec optimize failed: {e}") + + # -- lifecycle ---------------------------------------------------------- + + async def _close(self) -> None: + """Persist via the parent, then release the collection handle. + + zvec holds an in-process lock on the collection directory for the + lifetime of the ``Collection`` object; dropping the reference releases + it so another store instance can reopen the same directory. + """ + await super()._close() # parent dump() flushes the live collection + self._collection = None + self._indexed_ids = set() + + # -- 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 or has + round-tripped through the JSONL. Written into the sidecar at dump time + and recomputed from ``self.file_chunks`` at load time: a mismatch means + the collection belongs to a different chunk generation than the + authoritative JSONL. + """ + 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 the zvec collection (open or rebuild).""" + await super().load() + if self.embedding_store is None or self._dim == 0: + self._collection = None + return + if not await self._try_open_collection(): + self._rebuild_collection() + + async def _try_open_collection(self) -> bool: + """Open the persisted collection and validate it against the chunks. + + On any mismatch or open error the collection directory and sidecar are + wiped so the caller can rebuild from chunks cleanly. Validated: + + - vector dimension against the active embedding model; + - HNSW ``M`` against the active config (baked into the graph topology); + - the sidecar embedding digest against the authoritative JSONL; + - the collection contents against the ids and vectors it should hold. + """ + if not (self.zvec_path.exists() and self.zvec_sidecar_path.exists()): + return False + collection = None + try: + async with aiofiles.open(self.zvec_sidecar_path, encoding=self.encoding) as f: + sidecar = json.loads(await f.read()) + if sidecar.get("digest") != self._chunks_embedding_digest(): + raise ValueError("zvec sidecar embedding digest does not match persisted chunks") + indexed_ids = set(sidecar.get("ids", [])) + expected_ids = { + cid for cid, chunk in self.file_chunks.items() if self._embedding_dim_matches(chunk.embedding) + } + if indexed_ids != expected_ids: + raise ValueError("zvec sidecar ids do not match persisted chunks") + + collection = self._zvec.open(path=str(self.zvec_path)) + vector_schema = collection.schema.vectors[0] + if vector_schema.dimension != self._dim: + raise ValueError(f"zvec dim {vector_schema.dimension} != embedding dim {self._dim}") + persisted_m = vector_schema.index_param.m + if persisted_m != self.hnsw_m: + raise ValueError(f"zvec HNSW M mismatch: persisted={persisted_m}, configured={self.hnsw_m}") + self._verify_collection_contents(collection, indexed_ids) + + self._collection = collection + self._indexed_ids = indexed_ids + self.logger.info(f"Opened zvec collection: {len(indexed_ids)} vectors from {self.zvec_path}") + return True + except Exception as e: + self.logger.warning(f"Failed to open zvec collection, will rebuild: {e}") + collection = None # release the handle before wiping the directory + if self.zvec_path.exists(): + shutil.rmtree(self.zvec_path, ignore_errors=True) + self.zvec_sidecar_path.unlink(missing_ok=True) + return False + + def _verify_collection_contents(self, collection, expected_ids: set[str]) -> None: + """Check that the opened collection really holds the expected (id, vector) pairs. + + The sidecar only binds the collection to a chunk generation; it cannot + show that the collection itself lost, gained, or corrupted documents + (a crash between flushes, an external write, a partial copy). The + document count catches missing or extra ids cheaply, and the fetched + vectors catch same-id corruption. Any failure raises so the caller + rebuilds from the authoritative chunks. + """ + started_at = time.monotonic() + doc_count = collection.stats.doc_count + if doc_count != len(expected_ids): + raise ValueError(f"zvec collection holds {doc_count} documents, expected {len(expected_ids)}") + + ids = sorted(expected_ids) + for start in range(0, len(ids), _ZVEC_VERIFY_BATCH_SIZE): + batch = ids[start : start + _ZVEC_VERIFY_BATCH_SIZE] + docs = collection.fetch(batch, include_vector=True) + for cid in batch: + doc = docs.get(cid) + if doc is None: + raise ValueError(f"zvec collection is missing chunk {cid}") + indexed = self._vector_fingerprint(doc.vectors.get("embedding")) + if indexed != self._vector_fingerprint(self.file_chunks[cid].embedding): + raise ValueError(f"zvec vector for chunk {cid} does not match the persisted chunk") + self.logger.info( + f"{self.name}: verified zvec collection contents: docs={doc_count}, " + f"elapsed={time.monotonic() - started_at:.3f}s", + ) + + async def dump(self) -> None: + """Persist chunks JSONL via the parent, then flush zvec and write the sidecar.""" + await super().dump() + if self._collection is None or self.embedding_store is None: + return + try: + self._collection.flush() + await self._write_sidecar() + self.logger.info(f"Saved zvec collection: {len(self._indexed_ids)} vectors to {self.zvec_path}") + except Exception as e: + self.logger.exception(f"Failed to persist zvec collection: {e}") + + async def _write_sidecar(self) -> None: + """Atomically write the digest sidecar binding the collection to the chunk generation.""" + tmp = self.zvec_sidecar_path.with_name(f".{self.zvec_sidecar_path.name}.{uuid4().hex}.tmp") + payload = json.dumps( + { + "ids": sorted(self._indexed_ids), + "digest": self._chunks_embedding_digest(), + }, + ) + try: + async with aiofiles.open(tmp, "w", encoding=self.encoding) as f: + await f.write(payload) + tmp.replace(self.zvec_sidecar_path) + finally: + tmp.unlink(missing_ok=True) + + # -- CRUD overrides --------------------------------------------------- + + async def upsert(self, files: list[tuple[FileNode, list[FileChunk]]]) -> None: + if not files: + return + assert self.file_graph is not None + + # Snapshot pre-upsert chunk ids so we can drop chunks the new revision removed. + old_nodes = await self.file_graph.get_nodes([node.path for node, _ in files]) + old_ids_by_path = {n.path: set(n.chunk_ids) for n in old_nodes} + await super().upsert(files) + + if self._collection is None or self.embedding_store is None: + return + self._sync_collection_after_upsert(files, old_ids_by_path) + + def _sync_collection_after_upsert( + self, + files: list[tuple[FileNode, list[FileChunk]]], + old_ids_by_path: dict[str, set[str]], + ) -> None: + """Apply add / delete deltas to the zvec collection natively. + + Every eligible chunk of the request is re-upserted instead of being + diffed against its previous text: the parent accepts a caller-provided + embedding, so unchanged text does not imply an unchanged vector. zvec + upsert is idempotent, so re-sending an identical vector is cheap and + keeps the collection consistent with ``self.file_chunks``. Chunks that + no longer carry a usable vector are removed so the indexed id set stays + exactly the set of embeddable chunks. + """ + to_delete: list[str] = [] + to_upsert: list[FileChunk] = [] + for node, _ in files: + new_ids = set(node.chunk_ids) + to_delete.extend(old_ids_by_path.get(node.path, set()) - new_ids) + for cid in new_ids: + chunk = self.file_chunks.get(cid) + if chunk is None or not self._embedding_dim_matches(chunk.embedding): + to_delete.append(cid) # _delete_docs ignores ids that were never indexed + continue + to_upsert.append(chunk) + self._delete_docs(to_delete) + self._upsert_docs(to_upsert) + + async def delete(self, path: str | list[str]) -> None: + assert self.file_graph is not None + 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 self._delete_nodes(nodes) # reuse resolved nodes; avoids a second get_nodes + self._delete_docs(deleted_ids) + + async def clear(self) -> None: + await super().clear() + if self._collection is not None: + try: + self._collection.destroy() # drops the collection directory + except Exception: # pragma: no cover - defensive + shutil.rmtree(self.zvec_path, ignore_errors=True) + elif self.zvec_path.exists(): + shutil.rmtree(self.zvec_path, ignore_errors=True) + self._indexed_ids = set() + self.zvec_sidecar_path.unlink(missing_ok=True) + self._collection = self._create_collection() if self.embedding_store is not None else None + + # -- 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 limit <= 0 or self._collection is None or not self._indexed_ids: + 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 [] + + # get_embedding above yielded control; a concurrent clear() may have + # swapped or dropped the collection. Re-read before dereferencing. + collection = self._collection + if collection is None or not self._indexed_ids: + return [] + + vector = np.asarray(query_embedding, dtype=np.float32).tolist() + ntotal = len(self._indexed_ids) + + if not search_filter: + hits = self._query_collection(collection, vector, min(limit, ntotal)) + return self._collect_hits(hits, limit, search_filter) + + # With a post-filter: progressively increase k until we collect enough + # results or exhaust the reachable index. + k = min(ntotal, 3 * limit) + while True: + hits = self._query_collection(collection, vector, k) + results = self._collect_hits(hits, limit, search_filter) + if len(results) >= limit or k >= ntotal: + return results + k = min(ntotal, k * 2) + + def _query_collection(self, collection, vector: list[float], topk: int) -> list[tuple[str, float]]: + """Run an ANN query and convert cosine distance to similarity (higher = closer).""" + docs = collection.query( + self._zvec.Query(field_name="embedding", vector=vector), + topk=max(1, topk), + ) + return [(doc.id, 1.0 - float(doc.score)) for doc in docs] + + def _collect_hits( + self, + hits: list[tuple[str, float]], + limit: int, + search_filter: dict | None = None, + ) -> list[FileChunk]: + """Map zvec hits back to chunks, skipping stale ids and filtered-out chunks.""" + results: list[FileChunk] = [] + for chunk_id, score in hits: + chunk = self.file_chunks.get(chunk_id) + if chunk is None or not self._matches_search_filter(chunk, search_filter): + continue + results.append(chunk.model_copy(update={"scores": {"vector": score, "score": score}})) + if len(results) >= limit: + break + return results diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index e8ab7cdb..cd61cbdd 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -14,7 +14,7 @@ import time import numpy as np import pytest -from reme.components.file_store import FaissLocalFileStore, LocalFileStore +from reme.components.file_store import FaissLocalFileStore, LocalFileStore, ZvecLocalFileStore from reme.components.file_store import local_file_store as local_file_store_module from reme.schema import FileChunk, FileNode from reme.utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst @@ -148,6 +148,32 @@ def chunk(chunk_id: str, path: str, text: str, **metadata) -> FileChunk: return FileChunk(id=chunk_id, path=path, text=text, start_line=1, end_line=1, metadata=metadata) +def _new_local_store(name, **kwargs): + """Construct a LocalFileStore with embedding disabled at bind time.""" + return LocalFileStore(name=name, embedding_store="", **kwargs) + + +def _new_zvec_store(name, **kwargs): + """Construct a zvec store with embedding disabled at bind time.""" + try: + store = ZvecLocalFileStore(name=name, embedding_store="", **kwargs) + except ImportError: + pytest.skip("zvec is not installed") + return store + + +def _ensure_zvec_collection(store): + """Materialize the zvec collection once an embedding backend is attached. + + Fresh zvec stores start with no collection because ``embedding_store=""``. + Tests that attach a fake provider after ``start()`` must explicitly create + the collection before the first upsert, otherwise vectors are accepted by + the parent but never synced into zvec. + """ + if isinstance(store, ZvecLocalFileStore) and store._collection is None and store.embedding_store is not None: + store._collection = store._create_collection() + + async def set_chunks_with_graph(store: LocalFileStore, chunks: dict[str, FileChunk]) -> None: """Seed a graph/chunk snapshot that satisfies the persistence invariant.""" store.file_chunks = chunks @@ -454,14 +480,16 @@ def test_chunk_persistence_loads_legacy_json_embedding_list(): run(go()) -def test_same_chunk_id_with_changed_text_gets_new_embedding(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_same_chunk_id_with_changed_text_gets_new_embedding(store_factory): """Changing a chunk text refreshes its embedding.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_embedding_reuse", embedding_store="") + store = store_factory(name="t_embedding_reuse") await store.start() store.embedding_store = FakeEmbeddingStore() + _ensure_zvec_collection(store) await store.upsert([(node("note.md"), [chunk("same", "note.md", "alpha text")])]) assert store.file_chunks["same"].embedding.tolist() == [1.0, 0.0] @@ -474,12 +502,13 @@ def test_same_chunk_id_with_changed_text_gets_new_embedding(): run(go()) -def test_load_backfills_missing_embeddings_from_persisted_chunks(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_load_backfills_missing_embeddings_from_persisted_chunks(store_factory): """Startup backfills old chunks in the background and persists vectors.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + store = store_factory(name="t_embedding_backfill") await store.start() await store.upsert( [ @@ -489,7 +518,7 @@ def test_load_backfills_missing_embeddings_from_persisted_chunks(): ) await store.close() - store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + store = store_factory(name="t_embedding_backfill") store.embedding_store = FakeEmbeddingStore() await store.start() await store._embedding_backfill_task @@ -498,7 +527,7 @@ def test_load_backfills_missing_embeddings_from_persisted_chunks(): assert store.file_chunks["b"].embedding.tolist() == [0.0, 1.0] await store.close() - store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + store = store_factory(name="t_embedding_backfill") await store.start() assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] assert store.file_chunks["b"].embedding.tolist() == [0.0, 1.0] @@ -588,12 +617,13 @@ def test_load_skips_backfill_when_embedding_health_check_fails(): run(go()) -def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions(store_factory): """Loading persisted chunks re-embeds vectors that do not match current dimensions.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_embedding_stale_dim", embedding_store="") + store = store_factory(name="t_embedding_stale_dim") await store.start() stale = chunk("a", "a.md", "alpha text") stale.embedding = np.array([1.0], dtype=np.float16) @@ -601,7 +631,7 @@ def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions(): await store.dump() await store.close() - store = LocalFileStore(name="t_embedding_stale_dim", embedding_store="") + store = store_factory(name="t_embedding_stale_dim") fake = CountingFakeEmbeddingStore() store.embedding_store = fake await store.start() @@ -689,14 +719,16 @@ def test_upsert_reembeds_prefilled_chunk_with_stale_dimension(): run(go()) -def test_search_filter_applies_to_vector_and_keyword_results(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_search_filter_applies_to_vector_and_keyword_results(store_factory): """Search filters apply consistently to vector and keyword results.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_filter", embedding_store="") + store = store_factory(name="t_filter") await store.start() store.embedding_store = FakeEmbeddingStore() + _ensure_zvec_collection(store) await store.upsert( [ @@ -860,14 +892,16 @@ def test_date_filter_extract_and_match(): assert LocalFileStore._matches_search_filter(c, {"start_date": "2026-01-01", "end_date": "2026-12-31"}) is True -def test_date_filter_with_vector_and_keyword_search(): +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_zvec_store]) +def test_date_filter_with_vector_and_keyword_search(store_factory): """vector_search and keyword_search respect start_date/end_date filters.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): - store = LocalFileStore(name="t_date_search", embedding_store="") + store = store_factory(name="t_date_search") await store.start() store.embedding_store = FakeEmbeddingStore() + _ensure_zvec_collection(store) await store.upsert( [ diff --git a/tests/unit/test_zvec_file_store.py b/tests/unit/test_zvec_file_store.py new file mode 100644 index 00000000..ba1566f4 --- /dev/null +++ b/tests/unit/test_zvec_file_store.py @@ -0,0 +1,614 @@ +"""Tests for ZvecLocalFileStore: native upsert/delete sync, persistence, and rebuild triggers.""" + +# pylint: disable=protected-access + +import asyncio +import json +import os +import tempfile +import tomllib +import warnings +from pathlib import Path + +import numpy as np +import pytest + +from reme.components.file_store import ZvecLocalFileStore +from reme.schema import FileChunk, FileNode + + +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 embedded_chunk(chunk_id: str, path: str, text: str, embedding: list[float]) -> FileChunk: + """Build a chunk carrying an explicit, caller-provided embedding.""" + file_chunk = chunk(chunk_id, path, text) + file_chunk.embedding = np.asarray(embedding, dtype=np.float16) + return file_chunk + + +def _track_rebuilds(store: ZvecLocalFileStore) -> list[bool]: + """Record calls to the collection rebuild path.""" + rebuilds: list[bool] = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + return rebuilds + + +def _tamper_with_collection(store: ZvecLocalFileStore, mutate) -> None: + """Mutate the persisted collection behind the store's back, then release the handle. + + zvec holds an in-process lock on the collection directory, so the store must + already be closed and the local handle must be dropped before another store + reopens the same path. + """ + zvec = pytest.importorskip("zvec") + collection = zvec.open(path=str(store.zvec_path)) + mutate(zvec, collection) + collection.flush() + del collection + + +def _new_zvec_store(name, **kwargs): + """Construct a zvec store with embedding disabled at bind time.""" + try: + store = ZvecLocalFileStore(name=name, embedding_store="", **kwargs) + except ImportError: + pytest.skip("zvec is not installed") + return store + + +async def _started_store(name, **kwargs) -> ZvecLocalFileStore: + """Start a store with a fake embedding provider and a live collection.""" + store = _new_zvec_store(name, **kwargs) + store.embedding_store = FakeEmbeddingStore() + await store.start() + if store._collection is None: + store._collection = store._create_collection() + return store + + +async def _seed_unembedded_chunk(store: ZvecLocalFileStore, 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]) + + +# -- CRUD / search ------------------------------------------------------------ + + +def test_zvec_upsert_and_vector_search_ranks_by_similarity(): + """Upserted chunks are searchable; results are ranked by cosine similarity.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_basic") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.upsert([(node("b.md"), [chunk("b", "b.md", "beta text")])]) + + results = await store.vector_search("alpha", 10, {}) + assert [c.id for c in results] == ["a", "b"] + # Identical vector -> cosine distance 0 -> similarity score 1. + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + assert results[0].scores["score"] == results[0].scores["vector"] + await store.close() + + run(go()) + + +def test_zvec_same_id_text_change_updates_vector_in_place(): + """A same-id text change replaces the vector via native zvec upsert.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_update") + + 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"] + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "beta text")])]) + + assert store._indexed_ids == {"a"} + results = await store.vector_search("beta", 5, {}) + assert [c.id for c in results] == ["a"] + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + await store.close() + + run(go()) + + +def test_zvec_explicit_embedding_change_updates_collection(): + """A same-id, same-text update with a new explicit embedding replaces the vector.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_explicit_embedding") + + await store.upsert([(node("a.md"), [embedded_chunk("same", "a.md", "same text", [1.0, 0.0])])]) + await store.upsert([(node("a.md"), [embedded_chunk("same", "a.md", "same text", [0.0, 1.0])])]) + + # "beta" embeds to [0, 1]; the collection must hold the new vector. + results = await store.vector_search("beta", 5, {}) + assert [c.id for c in results] == ["same"] + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + await store.close() + + # The stale vector must not survive a restart either. + reopened = _new_zvec_store("t_zvec_explicit_embedding") + reopened.embedding_store = FakeEmbeddingStore() + rebuilds = _track_rebuilds(reopened) + await reopened.start() + + assert not rebuilds # the collection was already in sync, no repair needed + results = await reopened.vector_search("beta", 5, {}) + assert [c.id for c in results] == ["same"] + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + await reopened.close() + + run(go()) + + +def test_zvec_upsert_removes_stale_chunks_from_collection(): + """Re-upserting a path deletes vectors of chunks that no longer exist.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_stale") + + await store.upsert([(node("a.md"), [chunk("old", "a.md", "alpha old")])]) + await store.upsert([(node("a.md"), [chunk("new", "a.md", "alpha new")])]) + + assert store._indexed_ids == {"new"} + assert [c.id for c in await store.vector_search("alpha", 10, {})] == ["new"] + await store.close() + + run(go()) + + +def test_zvec_delete_removes_vectors(): + """Deleting a path removes its chunk vectors from the collection.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_delete") + + await store.upsert( + [ + (node("a.md"), [chunk("a", "a.md", "alpha text")]), + (node("b.md"), [chunk("b", "b.md", "beta text")]), + ], + ) + await store.delete("a.md") + + assert store._indexed_ids == {"b"} + assert [c.id for c in await store.vector_search("alpha", 10, {})] == ["b"] + await store.close() + + run(go()) + + +def test_zvec_vector_search_applies_post_filter(): + """The shared search_filter semantics apply on top of ANN results.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_filter") + + await store.upsert( + [ + (node("a.md"), [chunk("a", "a.md", "alpha one", kind="x")]), + (node("b.md"), [chunk("b", "b.md", "alpha two", kind="y")]), + ], + ) + + assert [c.id for c in await store.vector_search("alpha", 10, {"path": "b.md"})] == ["b"] + assert [c.id for c in await store.vector_search("alpha", 10, {"kind": "x"})] == ["a"] + assert await store.vector_search("alpha", 10, {"path": "c.md"}) == [] + await store.close() + + run(go()) + + +def test_zvec_vector_search_uses_current_query_api(): + """The ANN query path must not rely on zvec's deprecated VectorQuery alias.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_query_api") + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_clear_resets_collection(): + """clear() drops the collection, sidecar, and indexed-id state.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_clear") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.dump() + assert store.zvec_sidecar_path.exists() + + await store.clear() + + assert store._indexed_ids == set() + assert not store.zvec_sidecar_path.exists() + assert await store.vector_search("alpha", 10, {}) == [] + # The store stays usable after clear. + await store.upsert([(node("b.md"), [chunk("b", "b.md", "beta text")])]) + assert [c.id for c in await store.vector_search("beta", 10, {})] == ["b"] + await store.close() + + run(go()) + + +def test_zvec_keyword_only_mode_keeps_keyword_search_working(): + """Without an embedding store, vector search is empty but keyword search works.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_zvec_store("t_zvec_keyword_only") + await store.start() + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "uniquezvecword only")])]) + + assert store._collection is None + assert await store.vector_search("uniquezvecword", 5, {}) == [] + assert [c.id for c in await store.keyword_search("uniquezvecword", 5, {})] == ["a"] + await store.close() + + run(go()) + + +# -- embedding backfill --------------------------------------------------------- + + +def test_zvec_backfill_adds_incrementally(): + """Backfilled vectors are upserted into the live collection without a rebuild.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_backfill") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await _seed_unembedded_chunk(store, "c", "c.md", "alpha extra") + + rebuilds = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + + await store._backfill_missing_embeddings() + + assert not rebuilds + assert store._indexed_ids == {"a", "c"} + assert {c.id for c in await store.vector_search("alpha", 10, {})} == {"a", "c"} + await store.close() + + run(go()) + + +# -- persistence ---------------------------------------------------------------- + + +def test_zvec_persistence_round_trip_reopens_without_rebuild(): + """dump() + a fresh store reattaches the persisted collection via the sidecar.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_persist") + await seed.upsert( + [ + (node("a.md"), [chunk("a", "a.md", "alpha text")]), + (node("b.md"), [chunk("b", "b.md", "beta text")]), + ], + ) + await seed.close() + + store = _new_zvec_store("t_zvec_persist") + store.embedding_store = FakeEmbeddingStore() + + rebuilds = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + + await store.start() + + assert not rebuilds + assert store._indexed_ids == {"a", "b"} + assert [c.id for c in await store.vector_search("alpha", 5, {})][0] == "a" + assert [c.id for c in await store.vector_search("beta", 5, {})][0] == "b" + await store.close() + + run(go()) + + +def test_zvec_digest_mismatch_triggers_rebuild(): + """A sidecar whose digest fell behind the chunk JSONL forces a clean rebuild.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_digest") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + sidecar = json.loads(seed.zvec_sidecar_path.read_text()) + sidecar["digest"] = "0" * 64 + seed.zvec_sidecar_path.write_text(json.dumps(sidecar)) + + store = _new_zvec_store("t_zvec_digest") + store.embedding_store = FakeEmbeddingStore() + + rebuilds = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + + await store.start() + + assert rebuilds == [True] + assert store._indexed_ids == {"a"} + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_missing_sidecar_triggers_rebuild(): + """A collection directory without its sidecar cannot be trusted and is rebuilt.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_no_sidecar") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + seed.zvec_sidecar_path.unlink() + + store = _new_zvec_store("t_zvec_no_sidecar") + store.embedding_store = FakeEmbeddingStore() + await store.start() + + assert store._indexed_ids == {"a"} + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_hnsw_m_mismatch_triggers_rebuild(): + """A persisted collection built with a different HNSW M is rebuilt.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_m_mismatch", hnsw_m=16) + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + store = _new_zvec_store("t_zvec_m_mismatch", hnsw_m=48) + store.embedding_store = FakeEmbeddingStore() + + rebuilds = [] + original_rebuild = store._rebuild_collection + store._rebuild_collection = lambda: rebuilds.append(True) or original_rebuild() + + await store.start() + + assert rebuilds == [True] + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_missing_collection_document_triggers_rebuild(): + """A collection that lost a document is rebuilt even though the sidecar matches.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_missing_doc") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + _tamper_with_collection(seed, lambda _zvec, collection: collection.delete(ids=["a"])) + # The sidecar still claims the chunk is indexed. + assert json.loads(seed.zvec_sidecar_path.read_text())["ids"] == ["a"] + + store = _new_zvec_store("t_zvec_missing_doc") + store.embedding_store = FakeEmbeddingStore() + rebuilds = _track_rebuilds(store) + await store.start() + + assert rebuilds == [True] + assert store._indexed_ids == {"a"} + assert store._collection.stats.doc_count == 1 + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_unexpected_collection_document_triggers_rebuild(): + """A collection holding a document no chunk owns is rebuilt.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_extra_doc") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + _tamper_with_collection( + seed, + lambda zvec, collection: collection.upsert( + [zvec.Doc(id="ghost", vectors={"embedding": [0.0, 1.0]})], + ), + ) + + store = _new_zvec_store("t_zvec_extra_doc") + store.embedding_store = FakeEmbeddingStore() + rebuilds = _track_rebuilds(store) + await store.start() + + assert rebuilds == [True] + assert store._indexed_ids == {"a"} + assert store._collection.stats.doc_count == 1 + assert [c.id for c in await store.vector_search("beta", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_zvec_corrupted_collection_vector_triggers_rebuild(): + """An expected id whose stored vector was replaced is rebuilt from the chunks.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_bad_vector") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + _tamper_with_collection( + seed, + lambda zvec, collection: collection.upsert( + [zvec.Doc(id="a", vectors={"embedding": [0.0, 1.0]})], + ), + ) + + store = _new_zvec_store("t_zvec_bad_vector") + store.embedding_store = FakeEmbeddingStore() + rebuilds = _track_rebuilds(store) + await store.start() + + assert rebuilds == [True] + # The rebuilt collection carries the vector of the authoritative chunk. + results = await store.vector_search("alpha", 5, {}) + assert [c.id for c in results] == ["a"] + assert results[0].scores["vector"] == pytest.approx(1.0, abs=1e-5) + await store.close() + + run(go()) + + +def test_zvec_stale_embedding_dimension_rebuilds_from_backfill(): + """Persisted vectors with a stale dimension are dropped and re-embedded.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await _started_store("t_zvec_stale_dim") + await seed.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await seed.close() + + class WideEmbeddingStore(FakeEmbeddingStore): + """Same provider with a different dimension.""" + + dimensions = 4 + + def _embed(self, text: str) -> np.ndarray: + base = super()._embed(text) + return np.concatenate([base, base]).astype(np.float16) + + store = _new_zvec_store("t_zvec_stale_dim") + store.embedding_store = WideEmbeddingStore() + await store.start() + # Startup backfill re-embeds the stale chunk in the background. + await store._embedding_backfill_task + + assert store._collection.schema.vectors[0].dimension == 4 + assert store._indexed_ids == {"a"} + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +# -- packaging ------------------------------------------------------------ + + +def test_zvec_declared_as_installable_dependency(): + """The registered backend needs a supported installation path in pyproject.""" + pyproject = Path(__file__).resolve().parents[2] / "pyproject.toml" + if not pyproject.exists(): # running against an installed package, not the repo + pytest.skip("pyproject.toml is not available") + optional = tomllib.loads(pyproject.read_text(encoding="utf-8"))["project"]["optional-dependencies"] + assert any(dep.replace(" ", "").startswith("zvec") for dep in optional["core"]) + + +# -- maintenance ------------------------------------------------------------ + + +def test_zvec_optimize_index_runs_native_optimize(): + """optimize_index() delegates to Collection.optimize() and keeps search intact.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = await _started_store("t_zvec_optimize") + + await store.upsert([(node("a.md"), [chunk("a", "a.md", "alpha text")])]) + await store.optimize_index() + + assert [c.id for c in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go())