diff --git a/reme/components/embedding_store/base_embedding_store.py b/reme/components/embedding_store/base_embedding_store.py index 7f7d46e2..500ece5b 100644 --- a/reme/components/embedding_store/base_embedding_store.py +++ b/reme/components/embedding_store/base_embedding_store.py @@ -44,11 +44,21 @@ class BaseEmbeddingStore(BaseComponent): async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]: """Get embeddings for texts.""" + def _embedding_dim_matches(self, embedding: np.ndarray | None) -> bool: + """Return whether an embedding matches the configured model dimension.""" + if embedding is None: + return False + dimensions = getattr(self, "dimensions", None) + # Base stores may not expose dimensions; only enforce the check when they do. + if dimensions is None: + return True + return len(embedding) == dimensions + async def get_node_embeddings(self, nodes: list[EmbNode], **kwargs) -> list[EmbNode]: """Embed each node's text in-place and return the same list.""" embeddings = await self.get_embeddings([n.text for n in nodes], **kwargs) if len(embeddings) == len(nodes): for node, vec in zip(nodes, embeddings): - if vec is not None: + if vec is not None and self._embedding_dim_matches(vec): node.embedding = vec return nodes diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index ea15affd..dc023a13 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -61,6 +61,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore): result = await asyncio.wait_for(self.as_embedding(["ping"]), timeout=timeout) if not result or result[0] is None: raise RuntimeError("empty embedding") + if len(result[0]) != self.dimensions: + raise RuntimeError(f"embedding dimension mismatch: {len(result[0])} != {self.dimensions}") self.is_healthy = True self.logger.info(f"{tag} -> OK") except asyncio.TimeoutError: @@ -111,11 +113,18 @@ class LocalEmbeddingStore(BaseEmbeddingStore): if not embeddings or len(embeddings) != len(texts): return [] out: list[tuple[int, str, np.ndarray]] = [] + bad_dims: dict[int, int] = {} for (idx, _text, key), raw in zip(batch, embeddings): if raw is None: continue - emb = self._normalize_dim(np.asarray(raw, dtype=np.float16)) + emb = np.asarray(raw, dtype=np.float16) + if not self._validate_dim(emb): + bad_dims[len(emb)] = bad_dims.get(len(emb), 0) + 1 + continue out.append((idx, key, emb)) + if bad_dims: + details = ", ".join(f"{count} with dim {dim}" for dim, count in sorted(bad_dims.items())) + self.logger.error(f"Embedding dimension mismatch in batch: expected {self.dimensions}; rejected {details}") return out async def _call_with_retry(self, texts: list[str], **kwargs) -> list[list[float] | None] | None: @@ -132,12 +141,9 @@ class LocalEmbeddingStore(BaseEmbeddingStore): return None return None - def _normalize_dim(self, emb: np.ndarray) -> np.ndarray: - if len(emb) == self.dimensions: - return emb - if len(emb) < self.dimensions: - return np.pad(emb, (0, self.dimensions - len(emb))) - return emb[: self.dimensions] + def _validate_dim(self, emb: np.ndarray) -> bool: + """Return whether an embedding exactly matches the configured dimension.""" + return len(emb) == self.dimensions # -- Cache -- diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 80668bf9..e39078f5 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -96,7 +96,7 @@ class FaissLocalFileStore(LocalFileStore): self._id_map = [] self._id_to_row = {} self._tombstones.clear() - chunks = [c for c in self.file_chunks.values() if c.embedding is not None] + chunks = [c for c in self.file_chunks.values() if self._embedding_dim_matches(c.embedding)] if not chunks: return vectors = np.stack([c.embedding for c in chunks]) @@ -138,7 +138,9 @@ class FaissLocalFileStore(LocalFileStore): live_ids = [cid for i, cid in enumerate(id_map) if i not in tombstones] if len(live_ids) != len(set(live_ids)): raise ValueError("FAISS id_map contains duplicate live chunk ids") - expected_ids = {cid for cid, chunk in self.file_chunks.items() if chunk.embedding is not None} + expected_ids = { + cid for cid, chunk in self.file_chunks.items() if self._embedding_dim_matches(chunk.embedding) + } if set(live_ids) != expected_ids: raise ValueError("FAISS sidecar live ids do not match persisted chunks") self._faiss_index = index @@ -213,7 +215,7 @@ class FaissLocalFileStore(LocalFileStore): self._tombstone(cid) for cid in new_ids: chunk = self.file_chunks.get(cid) - if chunk is None or chunk.embedding is None: + if chunk is None or not self._embedding_dim_matches(chunk.embedding): continue if cid in existing and old_text_by_id.get(cid) == chunk.text: continue @@ -251,9 +253,7 @@ class FaissLocalFileStore(LocalFileStore): # -- 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: - return [] - if self._faiss_index.ntotal == 0: + if self.embedding_store is None or not query or self._faiss_index is None or self._faiss_index.ntotal == 0: return [] try: @@ -263,6 +263,11 @@ class FaissLocalFileStore(LocalFileStore): 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}", + ) + return [] q = self._prepare(query_embedding) ntotal = self._faiss_index.ntotal diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 1b58bcb7..c7b17aed 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -78,6 +78,30 @@ class LocalFileStore(BaseFileStore): self.logger.error(f"{self.name}: embedding disabled, {reason}") self.embedding_store = None + def _embedding_dim_matches(self, embedding: np.ndarray | None) -> bool: + """Return whether an index embedding matches the active embedding model.""" + # With no active embedding store, no persisted/index vector is trustworthy. + if self.embedding_store is None or embedding is None: + return False + return len(embedding) == self.embedding_store.dimensions + + def _drop_stale_embedding(self, chunk: FileChunk, context: str) -> bool: + """Drop a chunk embedding when it does not match the active model dimension.""" + if self.embedding_store is None: + return False + if chunk.embedding is None or self._embedding_dim_matches(chunk.embedding): + return False + self.logger.warning( + f"{self.name}: stale embedding for chunk {chunk.id} during {context}: " + f"{len(chunk.embedding)} != {self.embedding_store.dimensions}; re-embedding", + ) + chunk.embedding = None + return True + + def _drop_stale_embeddings(self, chunks: list[FileChunk], context: str) -> None: + for chunk in chunks: + self._drop_stale_embedding(chunk, context) + # -- persistence ---------------------------------------------------------- async def load(self) -> None: @@ -91,11 +115,18 @@ class LocalFileStore(BaseFileStore): chunk = FileChunk.model_validate_json(line) self.file_chunks[chunk.id] = chunk self.logger.info(f"Loaded {len(self.file_chunks)} chunks from {self.chunks_path}") + self._invalidate_stale_embeddings() await self._sync_keyword_index_from_chunks() await self._backfill_missing_embeddings() except Exception as e: self.logger.exception(f"Failed to load {self.chunks_path}: {e}") + def _invalidate_stale_embeddings(self) -> None: + """Drop persisted embeddings whose dimension no longer matches the active model.""" + if self.embedding_store is None: + return + self._drop_stale_embeddings(list(self.file_chunks.values()), "load") + async def _backfill_missing_embeddings(self) -> None: """Embed persisted chunks that predate embedding being enabled.""" if not self.embedding_store or not self.file_chunks: @@ -106,12 +137,17 @@ class LocalFileStore(BaseFileStore): return self.logger.info(f"{self.name}: backfilling embeddings for {len(missing)} chunks") + if not await self.embedding_store.health_check(): + self._disable_embedding("backfill health check failed") + return + try: await self.embedding_store.get_node_embeddings(missing) except Exception as e: self._disable_embedding(f"backfill: {type(e).__name__}: {e}") return + self._drop_stale_embeddings(missing, "backfill") filled = sum(1 for chunk in missing if chunk.embedding is not None) if filled: self.logger.info(f"{self.name}: backfilled embeddings for {filled}/{len(missing)} chunks") @@ -213,9 +249,17 @@ class LocalFileStore(BaseFileStore): cached: dict[str, CachedEmbedding], needs_embed: list[FileChunk], ) -> None: - if not self.embedding_store or chunk.embedding is not None: + if not self.embedding_store: return - if chunk.id in cached and cached[chunk.id][0] == chunk.text: + if chunk.embedding is not None: + if self._embedding_dim_matches(chunk.embedding): + return + self._drop_stale_embedding(chunk, "upsert") + if ( + chunk.id in cached + and cached[chunk.id][0] == chunk.text + and self._embedding_dim_matches(cached[chunk.id][1]) + ): chunk.embedding = cached[chunk.id][1] elif chunk.text: needs_embed.append(chunk) @@ -227,6 +271,8 @@ class LocalFileStore(BaseFileStore): await self.embedding_store.get_node_embeddings(chunks) except Exception as e: self._disable_embedding(f"upsert: {type(e).__name__}: {e}") + return + self._drop_stale_embeddings(chunks, "upsert") async def delete(self, path: str | list[str]) -> None: assert self.file_graph is not None @@ -282,11 +328,16 @@ class LocalFileStore(BaseFileStore): 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}", + ) + return [] candidates = [ c for c in self.file_chunks.values() - if c.embedding is not None and self._matches_search_filter(c, search_filter) + if self._embedding_dim_matches(c.embedding) and self._matches_search_filter(c, search_filter) ] if not candidates: return [] diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index 288fa29a..7954d8c6 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -54,6 +54,36 @@ class FakeEmbeddingStore: return nodes +class CountingFakeEmbeddingStore(FakeEmbeddingStore): + """Fake embedding store that records node backfill requests.""" + + def __init__(self): + self.node_embedding_calls: list[list[str]] = [] + + async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]: + self.node_embedding_calls.append([node.id for node in nodes]) + return await super().get_node_embeddings(nodes, **_kwargs) + + +class UnhealthyCountingEmbeddingStore(CountingFakeEmbeddingStore): + """Fake embedding store that fails the backfill health gate.""" + + async def health_check(self, _timeout: float = 2.0) -> bool: + return False + + +class WrongDimEmbeddingStore(FakeEmbeddingStore): + """Fake embedding store that returns vectors with the wrong dimension.""" + + async def get_embedding(self, input_text: str, **_kwargs) -> np.ndarray: + return np.array([1.0], dtype=np.float16) + + async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]: + for chunk_node in nodes: + chunk_node.embedding = np.array([1.0], dtype=np.float16) + return nodes + + def run(coro): """Run an async test body.""" return asyncio.run(coro) @@ -173,6 +203,132 @@ def test_load_backfills_missing_embeddings_from_persisted_chunks(): run(go()) +def test_load_skips_backfill_when_embedding_health_check_fails(): + """Backfill disables embeddings before batching when the provider is unhealthy.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = LocalFileStore(name="t_embedding_backfill_unhealthy", embedding_store="") + await store.start() + store.file_chunks = {"a": chunk("a", "a.md", "alpha text")} + await store.dump() + await store.close() + + store = LocalFileStore(name="t_embedding_backfill_unhealthy", embedding_store="") + await store.start() + fake = UnhealthyCountingEmbeddingStore() + store.embedding_store = fake + await store.load() + + assert not fake.node_embedding_calls + assert store.embedding_store is None + assert store.file_chunks["a"].embedding is None + await store.close() + + run(go()) + + +def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions(): + """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="") + await store.start() + stale = chunk("a", "a.md", "alpha text") + stale.embedding = np.array([1.0], dtype=np.float16) + store.file_chunks = {"a": stale} + await store.dump() + await store.close() + + store = LocalFileStore(name="t_embedding_stale_dim", embedding_store="") + await store.start() + fake = CountingFakeEmbeddingStore() + store.embedding_store = fake + await store.load() + + assert fake.node_embedding_calls == [["a"]] + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + await store.close() + + run(go()) + + +def test_drop_stale_embedding_noops_without_embedding_store(): + """The helper should not clear embeddings when vector search is disabled.""" + + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = LocalFileStore(name="t_embedding_no_store_drop", embedding_store="") + stale = chunk("a", "a.md", "alpha text") + stale.embedding = np.array([1.0], dtype=np.float16) + + assert store._drop_stale_embedding(stale, "test") is False + assert stale.embedding.tolist() == [1.0] + + +def test_upsert_does_not_reuse_cached_embedding_with_stale_dimensions(): + """Re-upsert queues a fresh embedding when cached same-text vector has old dimensions.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = LocalFileStore(name="t_embedding_stale_reuse", embedding_store="") + await store.start() + fake = CountingFakeEmbeddingStore() + store.embedding_store = fake + + await store.upsert([(node("note.md"), [chunk("same", "note.md", "alpha text")])]) + store.file_chunks["same"].embedding = np.array([1.0], dtype=np.float16) + await store.file_graph.upsert_nodes([FileNode(path="note.md", st_mtime=1.0, chunk_ids=["same"])]) + + await store.upsert([(node("note.md"), [chunk("same", "note.md", "alpha text")])]) + + assert fake.node_embedding_calls == [["same"], ["same"]] + assert store.file_chunks["same"].embedding.tolist() == [1.0, 0.0] + await store.close() + + run(go()) + + +def test_upsert_drops_wrong_dimension_from_custom_embedding_store(): + """Wrong-dimensional embeddings from custom stores are not persisted on chunks.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = LocalFileStore(name="t_embedding_wrong_dim_custom", embedding_store="") + await store.start() + store.embedding_store = WrongDimEmbeddingStore() + + await store.upsert([(node("note.md"), [chunk("a", "note.md", "alpha text")])]) + + assert store.file_chunks["a"].embedding is None + assert await store.vector_search("alpha", 5, {}) == [] + assert store.embedding_store is None + await store.close() + + run(go()) + + +def test_upsert_reembeds_prefilled_chunk_with_stale_dimension(): + """Incoming chunks with stale embeddings are re-embedded before persistence.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = LocalFileStore(name="t_embedding_prefilled_stale", embedding_store="") + await store.start() + fake = CountingFakeEmbeddingStore() + store.embedding_store = fake + prefilled = chunk("a", "note.md", "alpha text") + prefilled.embedding = np.array([1.0], dtype=np.float16) + + await store.upsert([(node("note.md"), [prefilled])]) + + assert fake.node_embedding_calls == [["a"]] + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + await store.close() + + run(go()) + + def test_search_filter_applies_to_vector_and_keyword_results(): """Search filters apply consistently to vector and keyword results.""" @@ -229,6 +385,33 @@ def test_faiss_rebuilds_stale_sidecar_and_updates_same_id_text(): run(go()) +def test_faiss_rebuild_skips_wrong_dimension_chunks(): + """FAISS rebuild should ignore chunks whose embedding dimensions do not match.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + try: + store = FaissLocalFileStore(name="t_faiss_dim_filter", embedding_store="") + except ImportError: + pytest.skip("faiss is not installed") + await store.start() + store.embedding_store = FakeEmbeddingStore() + store.file_chunks = { + "good": chunk("good", "good.md", "alpha text"), + "bad": chunk("bad", "bad.md", "alpha text"), + } + store.file_chunks["good"].embedding = np.array([1.0, 0.0], dtype=np.float16) + store.file_chunks["bad"].embedding = np.array([1.0], dtype=np.float16) + + store._rebuild_index() + + assert set(store._id_to_row) == {"good"} + assert store._faiss_index.ntotal == 1 + await store.close() + + run(go()) + + # -- Date filter tests ------------------------------------------------------- diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py new file mode 100644 index 00000000..48181541 --- /dev/null +++ b/tests/unit/test_local_embedding_store.py @@ -0,0 +1,96 @@ +"""Regression tests for LocalEmbeddingStore dimension handling.""" + +# pylint: disable=protected-access + +import asyncio + +import numpy as np + +from reme.components.embedding_store.base_embedding_store import BaseEmbeddingStore +from reme.components.embedding_store.local_embedding_store import LocalEmbeddingStore +from reme.schema import EmbNode + + +class FakeAsEmbedding: + """Fake AgentScope embedding component.""" + + dimensions = 2 + + async def __call__(self, texts: list[str], **_kwargs): + return [[1.0] if text == "bad" else [1.0, 0.0] for text in texts] + + +class BadHealthAsEmbedding: + """Fake provider whose health probe returns the wrong dimension.""" + + dimensions = 2 + + async def __call__(self, _texts: list[str], **_kwargs): + return [[1.0]] + + +class BadNodeEmbeddingStore(BaseEmbeddingStore): + """Embedding store that returns wrong-dimensional vectors.""" + + dimensions = 2 + + async def health_check(self, timeout: float = 2.0) -> bool: + return True + + async def get_embeddings(self, input_text: list[str], **_kwargs): + return [np.array([1.0], dtype=np.float16) for _ in input_text] + + +def run(coro): + """Run an async test body.""" + return asyncio.run(coro) + + +def test_compute_batch_rejects_embeddings_with_wrong_dimension(): + """Provider results with wrong dimensions are not padded, truncated, or cached.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_dim") + store.as_embedding = FakeAsEmbedding() + store._key_suffix = f"|{store.dimensions}".encode() + + results = await store._compute_batch( + [ + (0, "ok", "ok-cache-key"), + (1, "bad", "bad-cache-key"), + ], + ) + + assert len(results) == 1 + assert results[0][0] == 0 + assert results[0][2].tolist() == [1.0, 0.0] + assert isinstance(results[0][2], np.ndarray) + + run(go()) + + +def test_base_get_node_embeddings_rejects_wrong_dimension(): + """Base node assignment should not accept wrong-dimensional vectors.""" + + async def go(): + store = BadNodeEmbeddingStore(name="t_base_embedding_dim") + node = EmbNode(text="bad") + + await store.get_node_embeddings([node]) + + assert node.embedding is None + + run(go()) + + +def test_health_check_rejects_wrong_dimension(): + """Health check should fail when the provider returns the wrong vector length.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_health_dim") + store.as_embedding = BadHealthAsEmbedding() + + assert await store.health_check() is False + assert store.is_healthy is False + + run(go())