fix(embedding): reject mismatched embedding dimensions (#330)

* fix(embedding): enforce strict dimension matching for embeddings

- Add _embedding_dim_matches method to validate embedding dimensions
- Reject embeddings with mismatched dimensions instead of padding/truncating
- Drop stale embeddings with wrong dimensions during loading and upsert operations
- Disable embedding store when query dimensions don't match configured dimensions
- Fail health checks when embedding dimensions don't match expected values
- Skip chunks with wrong dimensions during FAISS index rebuild
- Add comprehensive tests for dimension validation behavior

* refactor(file_store): simplify conditional checks in vector search and test assertions

- Combine multiple conditionals into single check for empty FAISS index
- Replace explicit empty list comparison with boolean check for node embedding calls
- Maintain same functional behavior while improving code readability

* fix(embedding): harden dimension validation helpers

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

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

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

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

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

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