From d9448fca120c6a426f0145d8e3bec997ef9e6c81 Mon Sep 17 00:00:00 2001 From: lc285652 Date: Wed, 5 Aug 2026 15:20:22 +0800 Subject: [PATCH] fix(file_store): fix zvec collection sync and content validation, declare zvec dependency --- pyproject.toml | 1 + .../file_store/zvec_local_file_store.py | 87 +++++++-- tests/unit/test_zvec_file_store.py | 179 ++++++++++++++++++ 3 files changed, 250 insertions(+), 17 deletions(-) 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/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py index 8d7209c1..80129eed 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -3,6 +3,7 @@ import hashlib import json import shutil +import time from uuid import uuid4 import aiofiles @@ -14,6 +15,8 @@ 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") @@ -33,7 +36,9 @@ class ZvecLocalFileStore(LocalFileStore): 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. + 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 @@ -43,7 +48,8 @@ class ZvecLocalFileStore(LocalFileStore): idle-time index compaction hook. zvec is imported lazily inside ``__init__`` so that merely importing this - module does not require the optional dependency. + module does not require the optional dependency; the backend targets the + zvec version declared in ``pyproject.toml`` (``zvec>=0.6.0``). """ def __init__( @@ -67,7 +73,7 @@ class ZvecLocalFileStore(LocalFileStore): import zvec except ImportError as e: raise ImportError( - "zvec is required for ZvecLocalFileStore. Install with `pip install zvec`.", + "zvec is required for ZvecLocalFileStore. Install with `pip install 'zvec>=0.6.0'`.", ) from e return zvec @@ -113,6 +119,19 @@ class ZvecLocalFileStore(LocalFileStore): 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 @@ -226,7 +245,8 @@ class ZvecLocalFileStore(LocalFileStore): - 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 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 @@ -250,6 +270,7 @@ class ZvecLocalFileStore(LocalFileStore): 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 @@ -263,6 +284,37 @@ class ZvecLocalFileStore(LocalFileStore): 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() @@ -298,28 +350,30 @@ class ZvecLocalFileStore(LocalFileStore): return assert self.file_graph is not None - # Snapshot pre-upsert chunk state so we can diff after the parent ran. + # 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} - old_text_by_id = { - cid: chunk.text - for n in old_nodes - for cid in n.chunk_ids - if (chunk := self.file_chunks.get(cid)) is not None - } 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, old_text_by_id) + 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]], - old_text_by_id: dict[str, str], ) -> None: - """Apply add / delete deltas to the zvec collection natively.""" + """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: @@ -328,9 +382,8 @@ class ZvecLocalFileStore(LocalFileStore): 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 - if cid in self._indexed_ids and old_text_by_id.get(cid) == chunk.text: - continue # unchanged text -> vector already in the collection to_upsert.append(chunk) self._delete_docs(to_delete) self._upsert_docs(to_upsert) @@ -400,7 +453,7 @@ class ZvecLocalFileStore(LocalFileStore): 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.VectorQuery(field_name="embedding", vector=vector), + self._zvec.Query(field_name="embedding", vector=vector), topk=max(1, topk), ) return [(doc.id, 1.0 - float(doc.score)) for doc in docs] diff --git a/tests/unit/test_zvec_file_store.py b/tests/unit/test_zvec_file_store.py index de85cfed..ba1566f4 100644 --- a/tests/unit/test_zvec_file_store.py +++ b/tests/unit/test_zvec_file_store.py @@ -6,6 +6,9 @@ import asyncio import json import os import tempfile +import tomllib +import warnings +from pathlib import Path import numpy as np import pytest @@ -71,6 +74,35 @@ 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 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: @@ -142,6 +174,37 @@ def test_zvec_same_id_text_change_updates_vector_in_place(): 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.""" @@ -203,6 +266,22 @@ def test_zvec_vector_search_applies_post_filter(): 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.""" @@ -384,6 +463,94 @@ def test_zvec_hnsw_m_mismatch_triggers_rebuild(): 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.""" @@ -416,6 +583,18 @@ def test_zvec_stale_embedding_dimension_rebuilds_from_backfill(): 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 ------------------------------------------------------------