diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 86d94177..4ccd76cf 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -91,9 +91,31 @@ class LocalFileStore(BaseFileStore): self.file_chunks[chunk.id] = chunk self.logger.info(f"Loaded {len(self.file_chunks)} chunks from {self.chunks_path}") 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}") + 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: + return + + missing = [chunk for chunk in self.file_chunks.values() if chunk.text and chunk.embedding is None] + if not missing: + return + + self.logger.info(f"{self.name}: backfilling embeddings for {len(missing)} chunks") + try: + await self.embedding_store.get_node_embeddings(missing) + except Exception as e: + self._disable_embedding(f"backfill: {type(e).__name__}: {e}") + return + + 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") + await self.dump() + async def _sync_keyword_index_from_chunks(self) -> None: """Repair keyword index when its persisted state does not match chunks.""" if not self.keyword_index or not self.file_chunks: diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index 6ddd1a97..8d42c070 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -140,6 +140,39 @@ def test_same_chunk_id_with_changed_text_gets_new_embedding(): run(go()) +def test_load_backfills_missing_embeddings_from_persisted_chunks(): + """Loading old chunks after enabling embeddings backfills and persists vectors.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + await store.start() + await store.upsert( + [ + (node("a.md"), [chunk("a", "a.md", "alpha text")]), + (node("b.md"), [chunk("b", "b.md", "fresh beta text")]), + ], + ) + await store.close() + + store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + await store.start() + store.embedding_store = FakeEmbeddingStore() + await store.load() + + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert store.file_chunks["b"].embedding.tolist() == [0.0, 1.0] + await store.close() + + store = LocalFileStore(name="t_embedding_backfill", embedding_store="") + 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] + await store.close() + + run(go()) + + def test_search_filter_applies_to_vector_and_keyword_results(): """Search filters apply consistently to vector and keyword results."""