From c8e124876974d79da15ca8c35979ed2b21e14b20 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:58:51 +0800 Subject: [PATCH] fix: recover embeddings after transient health check failure (#471) * fix: recover embedding after transient health failure * refactor(embedding_store): remove provider_success_count and simplify health recovery logic - Deleted provider_success_count attribute and related methods across embedding and file stores - Updated _recover_after_real_request to rely solely on is_healthy flag for recovery decisions - Removed redundant counting logic for provider successes during embedding operations - Cleaned up health status management to streamline provider recovery detection - Adjusted unit tests to align with removal of provider_success_count and maintain health checks consistency * refactor(embedding_store): use default health check timeout * fix(embedding_store): ensure is_healthy remains unchanged on cache hits - Updated get_embeddings docstring to clarify cache hits must not alter is_healthy state - Improved code comment for embedding dimension matching method * fix(file_store): make embedding recovery race-safe * ci: use default CodeQL query suite * fix(file_store): preserve queued embedding rebuilds * fix(file_store): preserve verified recovery without chunks --- .github/workflows/security-codeql.yml | 1 - docs/en/memory_search.md | 8 + docs/zh/memory_search.md | 6 + .../embedding_store/base_embedding_store.py | 6 +- .../embedding_store/local_embedding_store.py | 25 +- .../file_store/faiss_local_file_store.py | 18 +- .../components/file_store/local_file_store.py | 160 +++++++-- .../file_store/zvec_local_file_store.py | 15 +- tests/unit/test_file_store_consistency.py | 317 +++++++++++++++++- tests/unit/test_local_embedding_store.py | 34 ++ 10 files changed, 541 insertions(+), 49 deletions(-) diff --git a/.github/workflows/security-codeql.yml b/.github/workflows/security-codeql.yml index 99ef6b2c..2e558a46 100644 --- a/.github/workflows/security-codeql.yml +++ b/.github/workflows/security-codeql.yml @@ -37,7 +37,6 @@ jobs: with: languages: ${{ matrix.language }} build-mode: none - queries: security-and-quality - name: Perform CodeQL analysis uses: github/codeql-action/analyze@v4 diff --git a/docs/en/memory_search.md b/docs/en/memory_search.md index 3beb80b4..bdc4692e 100644 --- a/docs/en/memory_search.md +++ b/docs/en/memory_search.md @@ -117,6 +117,14 @@ Out of the box, search therefore uses primarily BM25 plus link expansion. After `SearchStep` runs vector and keyword recall together. Additionally, switching the `file_store` `backend` from `local` to `faiss` upgrades vector retrieval from a linear scan to a FAISS HNSW index, offering faster recall at scale. +The embedding store accepts `health_check_timeout` for its startup probe. A temporary failure skips the current vector +backfill while keeping BM25 available; a later successful provider request resumes the missing-vector backfill +automatically. + +Embedded integrations that have already verified a provider can call `resume_embedding(verified=True)`. When changing +the embedding vector space, pass `rebuild=True`; persisted vectors are invalidated before a serial background rebuild, +and vector search remains unavailable until the rebuilt vectors are safely persisted. + ## How to Search The `search` Job is also configured in `default.yaml`: diff --git a/docs/zh/memory_search.md b/docs/zh/memory_search.md index 502e2bec..67fdac8c 100644 --- a/docs/zh/memory_search.md +++ b/docs/zh/memory_search.md @@ -106,6 +106,12 @@ file_store: 所以开箱搜索主要是 BM25 + 链接展开。把 `embedding_store: default` 打开后，`SearchStep` 会同时跑向量召回和关键词召回。此时若将 `file_store` 的 `backend` 从 `local` 改为 `faiss`，向量检索会从线性扫描升级为 FAISS HNSW 索引，在大规模 chunk 场景下召回效率更高。 +Embedding store 可通过 `health_check_timeout` 配置启动探测。临时失败只会跳过本次向量回填，BM25 仍可使用； +后续真实请求成功后会自动恢复缺失向量的回填。 + +已经完成真实服务验证的嵌入式集成可以调用 `resume_embedding(verified=True)`。切换 Embedding 向量空间时应同时传入 +`rebuild=True`；ReMe 会先使旧向量失效，再串行后台重建，并在新向量安全持久化前暂停向量搜索。 + ## 怎么搜索 `search` Job 也是在 `default.yaml` 中配置： diff --git a/reme/components/embedding_store/base_embedding_store.py b/reme/components/embedding_store/base_embedding_store.py index cc3cbed5..bb7703e2 100644 --- a/reme/components/embedding_store/base_embedding_store.py +++ b/reme/components/embedding_store/base_embedding_store.py @@ -25,6 +25,7 @@ class BaseEmbeddingStore(BaseComponent): max_input_length: int = 8192, max_retries: int = 3, quota_retry_delay: float | None = None, + health_check_timeout: float = 15.0, **kwargs, ): super().__init__(**kwargs) @@ -32,6 +33,7 @@ class BaseEmbeddingStore(BaseComponent): self.max_input_length = max_input_length self.max_retries = max_retries self.quota_retry_delay = quota_retry_delay + self.health_check_timeout = health_check_timeout self.is_healthy: bool = True def _truncate(self, text: str) -> str: @@ -57,7 +59,7 @@ class BaseEmbeddingStore(BaseComponent): return text @abstractmethod - async def health_check(self, timeout: float = 2.0) -> bool: + async def health_check(self, timeout: float | None = None) -> bool: """Probe the provider; sets and returns is_healthy.""" async def get_embedding(self, input_text: str, **kwargs) -> np.ndarray | None: @@ -67,7 +69,7 @@ class BaseEmbeddingStore(BaseComponent): @abstractmethod async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]: - """Get embeddings for texts.""" + """Get embeddings; cache hits must not change is_healthy.""" def _embedding_dim_matches(self, embedding: np.ndarray | None) -> bool: """Return whether an embedding matches the configured model dimension.""" diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 930d9af4..81167560 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -68,8 +68,12 @@ class LocalEmbeddingStore(BaseEmbeddingStore): async def _close(self) -> None: await self.dump() - async def health_check(self, timeout: float = 5.0) -> bool: + async def health_check(self, timeout: float | None = None) -> bool: + timeout = self.health_check_timeout if timeout is None else timeout + if not isinstance(timeout, (int, float)) or not np.isfinite(timeout) or timeout <= 0: + raise ValueError("timeout must be finite and greater than 0") tag = f"[EMBEDDING HEALTH CHECK] name={self.name} workspace_dir={self.workspace_path}" + started_at = asyncio.get_running_loop().time() try: # Provider construction may synchronously import an SDK and build # its HTTP client. Keep that one-time work outside the request @@ -82,13 +86,18 @@ class LocalEmbeddingStore(BaseEmbeddingStore): 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") + elapsed = asyncio.get_running_loop().time() - started_at + self.logger.info(f"{tag} -> OK timeout={timeout}s elapsed={elapsed:.3f}s") except asyncio.TimeoutError: self.is_healthy = False - self.logger.error(f"{tag} -> FAIL timeout({timeout}s)") - except Exception as e: + elapsed = asyncio.get_running_loop().time() - started_at + self.logger.error(f"{tag} -> FAIL timeout={timeout}s elapsed={elapsed:.3f}s error=timeout({timeout}s)") + except Exception as exc: # Provider SDKs expose many exception types. self.is_healthy = False - self.logger.error(f"{tag} -> FAIL {type(e).__name__}: {e}") + elapsed = asyncio.get_running_loop().time() - started_at + self.logger.error( + f"{tag} -> FAIL timeout={timeout}s elapsed={elapsed:.3f}s error={type(exc).__name__}: {exc}", + ) return self.is_healthy # -- Public API -- @@ -143,6 +152,10 @@ class LocalEmbeddingStore(BaseEmbeddingStore): 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}") + if out: + self.is_healthy = True + else: + self.is_healthy = False return out async def _call_with_retry(self, texts: list[str], **kwargs) -> list[list[float] | None] | None: @@ -166,7 +179,9 @@ class LocalEmbeddingStore(BaseEmbeddingStore): await asyncio.sleep(self.quota_retry_delay) continue self.logger.exception("Embedding request failed") + self.is_healthy = False return None + self.is_healthy = False return None @staticmethod diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 5fdf7846..2e02bfe2 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -252,6 +252,11 @@ class FaissLocalFileStore(LocalFileStore): self._add_to_index([c.id for c in to_add], vectors) self._compact_if_needed() + async def _reset_vector_index(self) -> None: + """Discard all vectors before rebuilding a changed vector space.""" + await self._stop_reindex_worker() + self._rebuild_index() + # -- async reindex ---------------------------------------------------- def _submit_reindex(self) -> None: @@ -625,26 +630,29 @@ class FaissLocalFileStore(LocalFileStore): # -- search ----------------------------------------------------------- async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]: + index_empty = self._faiss_index is None or self._faiss_index.ntotal == 0 + embedding_unavailable = self.embedding_store is None or self._embedding_rebuild_pending if ( - self.embedding_store is None + embedding_unavailable or not query or limit <= 0 - or self._faiss_index is None - or self._faiss_index.ntotal == 0 + or (index_empty and getattr(self.embedding_store, "is_healthy", True)) ): return [] query_embedding = None + was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) try: query_embedding = await self.embedding_store.get_embedding(query) except Exception as e: - self._disable_embedding(f"search: {type(e).__name__}: {e}") + self._mark_embedding_unhealthy(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( + self._mark_embedding_unhealthy( f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", ) return [] + await self._recover_after_real_request(was_healthy) # get_embedding above yielded control; a concurrent clear() drops the # index to None once embedding is disabled, and a reindex may have swapped diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 889a792a..3bb797f8 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -67,10 +67,14 @@ class LocalFileStore(BaseFileStore): self.file_chunks: dict[str, FileChunk] = {} self.chunks_path = self.component_metadata_path / f"file_chunks_{self.name}_{self.store_version}.jsonl.zst" self._embedding_backfill_task: asyncio.Task | None = None + self._embedding_backfill_pending: tuple[bool, bool] | None = None + self._embedding_rebuild_pending = False + self._closing = False # -- lifecycle ------------------------------------------------------------ async def _start(self) -> None: + self._closing = False started_at = time.monotonic() self.component_metadata_path.mkdir(parents=True, exist_ok=True) await super()._start() @@ -94,17 +98,45 @@ class LocalFileStore(BaseFileStore): ) async def _close(self) -> None: + self._closing = True await self._cancel_embedding_backfill() await self.dump() self.file_chunks.clear() await super()._close() - def _disable_embedding(self, reason: str) -> None: - """Drop embedding after a runtime failure; keyword search still works.""" + def _mark_embedding_unhealthy(self, reason: str) -> None: + """Record a temporary provider failure while preserving the component.""" if self.embedding_store is None: return - self.logger.error(f"{self.name}: embedding disabled, {reason}") - self.embedding_store = None + self.embedding_store.is_healthy = False + self.logger.error(f"{self.name}: embedding unavailable, {reason}; keyword search remains active") + + async def _recover_after_real_request(self, was_healthy: bool) -> None: + """Schedule repair when a real, non-cache provider request recovers.""" + if self.embedding_store is None or was_healthy or not getattr(self.embedding_store, "is_healthy", True): + return + self.logger.info(f"{self.name}: embedding provider recovered; scheduling missing-vector backfill") + await self.resume_embedding(verified=True) + + async def resume_embedding(self, *, verified: bool = False, rebuild: bool = False) -> bool: + """Resume a configured provider and schedule a deduplicated repair. + + Embedded applications may pass ``verified=True`` after they have already + made a successful real provider request, avoiding a redundant ping. Pass + ``rebuild=True`` when the active vector space changed; existing vectors + are derived data and are discarded before a full background rebuild. + """ + if self.embedding_store is None or self._closing: + return False + if verified: + self.embedding_store.is_healthy = True + if rebuild: + await self._prepare_embedding_rebuild() + if not self.file_chunks: + self._embedding_rebuild_pending = False + return True + self._start_embedding_backfill(skip_health_check=verified, rebuild=rebuild) + return True def _embedding_dim_matches(self, embedding: np.ndarray | None) -> bool: """Return whether an index embedding matches the active embedding model.""" @@ -228,29 +260,40 @@ class LocalFileStore(BaseFileStore): return self._drop_stale_embeddings(self.file_chunks.values(), "load") - def _start_embedding_backfill(self) -> None: + def _start_embedding_backfill(self, *, skip_health_check: bool = False, rebuild: bool = False) -> None: """Schedule startup embedding repair without delaying component readiness.""" started_at = time.monotonic() + if self._closing: + self.logger.info(f"{self.name}: embedding backfill skipped: reason=closing") + return if not self.embedding_store: self.logger.info( f"{self.name}: embedding backfill skipped: reason=embedding_disabled, " f"elapsed={time.monotonic() - started_at:.3f}s", ) return + if self._embedding_backfill_task is not None and not self._embedding_backfill_task.done(): + pending_verified = skip_health_check or bool( + self._embedding_backfill_pending and self._embedding_backfill_pending[0], + ) + pending_rebuild = rebuild or bool( + self._embedding_backfill_pending and self._embedding_backfill_pending[1], + ) + if pending_verified or pending_rebuild: + self._embedding_backfill_pending = (pending_verified, pending_rebuild) + self.logger.info( + f"{self.name}: embedding backfill scheduling skipped: reason=already_running, " + f"elapsed={time.monotonic() - started_at:.3f}s", + ) + return if not self.file_chunks: self.logger.info( f"{self.name}: embedding backfill skipped: reason=no_chunks, " f"elapsed={time.monotonic() - started_at:.3f}s", ) return - if self._embedding_backfill_task is not None and not self._embedding_backfill_task.done(): - self.logger.info( - f"{self.name}: embedding backfill scheduling skipped: reason=already_running, " - f"elapsed={time.monotonic() - started_at:.3f}s", - ) - return self._embedding_backfill_task = asyncio.create_task( - self._backfill_missing_embeddings(), + self._run_embedding_backfill(skip_health_check=skip_health_check, rebuild=rebuild), name=f"embedding-backfill:{self.name}", ) self.logger.info( @@ -258,10 +301,37 @@ class LocalFileStore(BaseFileStore): f"elapsed={time.monotonic() - started_at:.3f}s", ) + async def _run_embedding_backfill(self, *, skip_health_check: bool, rebuild: bool) -> None: + """Run one repair and honor a verified request queued behind it.""" + current_task = asyncio.current_task() + try: + if rebuild: + # A task that was already running when rebuild was requested + # may have written a stale provider result after the first + # invalidation. Clear once more at the queue boundary. + await self._prepare_embedding_rebuild() + await self._backfill_missing_embeddings(skip_health_check=skip_health_check) + finally: + if self._embedding_backfill_task is current_task: + self._embedding_backfill_task = None + pending = self._embedding_backfill_pending + self._embedding_backfill_pending = None + if pending is not None and not self._closing and self.embedding_store is not None: + pending_verified, pending_rebuild = pending + if pending_verified: + self.embedding_store.is_healthy = True + if pending_rebuild: + self._embedding_rebuild_pending = True + self._start_embedding_backfill( + skip_health_check=pending_verified, + rebuild=pending_rebuild, + ) + async def _cancel_embedding_backfill(self) -> None: """Cancel and collect the startup repair task during component shutdown.""" task = self._embedding_backfill_task self._embedding_backfill_task = None + self._embedding_backfill_pending = None if task is None: return if not task.done(): @@ -285,9 +355,22 @@ class LocalFileStore(BaseFileStore): next_percent += _PROGRESS_LOG_PERCENT_STEP return next_percent - async def _backfill_missing_embeddings(self) -> None: + async def _backfill_missing_embeddings(self, *, skip_health_check: bool = False) -> None: """Background-repair persisted chunks that do not have usable vectors.""" started_at = time.monotonic() + try: + await self._backfill_missing_embeddings_inner(skip_health_check=skip_health_check, started_at=started_at) + finally: + if self._embedding_rebuild_pending and not self._closing: + try: + await self._after_embedding_backfill() + await self.dump() + self._embedding_rebuild_pending = False + except Exception: + self.logger.exception(f"{self.name}: failed to finalize embedding rebuild") + + async def _backfill_missing_embeddings_inner(self, *, skip_health_check: bool, started_at: float) -> None: + """Perform one backfill pass; the caller owns rebuild finalization.""" if not self.embedding_store or not self.file_chunks: self.logger.info( f"{self.name}: embedding backfill finished without work: " @@ -313,18 +396,18 @@ class LocalFileStore(BaseFileStore): batch_size = max(1, int(getattr(self.embedding_store, "max_batch_size", 10))) self.logger.info(f"{self.name}: embedding backfill started: total={total}, batch_size={batch_size}") try: - health_check_started_at = time.monotonic() - is_healthy = await self.embedding_store.health_check() - self.logger.info( - f"{self.name}: embedding health check complete: healthy={is_healthy}, " - f"elapsed={time.monotonic() - health_check_started_at:.3f}s", - ) - if not is_healthy: - self._disable_embedding("backfill health check failed") - self.logger.warning( - f"{self.name}: embedding backfill failed: processed=0/{total}, reason=health check failed", + if not skip_health_check: + health_check_started_at = time.monotonic() + is_healthy = await self.embedding_store.health_check() + self.logger.info( + f"{self.name}: embedding health check complete: healthy={is_healthy}, " + f"elapsed={time.monotonic() - health_check_started_at:.3f}s", ) - return + if not is_healthy: + self.logger.warning( + f"{self.name}: embedding backfill skipped: processed=0/{total}, reason=health check failed", + ) + return processed = 0 batch_count = 0 @@ -348,8 +431,9 @@ class LocalFileStore(BaseFileStore): f"{total}, elapsed={elapsed:.2f}s", ) raise + except Exception as e: - self._disable_embedding(f"backfill: {type(e).__name__}: {e}") + self._mark_embedding_unhealthy(f"backfill: {type(e).__name__}: {e}") elapsed = time.monotonic() - started_at self.logger.exception( f"{self.name}: embedding backfill failed: processed={processed if 'processed' in locals() else 0}/" @@ -362,13 +446,24 @@ class LocalFileStore(BaseFileStore): self.logger.info( f"{self.name}: embedding backfill complete: filled={filled}/{total}, elapsed={elapsed:.2f}s", ) - if filled: + if filled and not self._embedding_rebuild_pending: try: await self._after_embedding_backfill() await self.dump() except Exception: self.logger.exception(f"{self.name}: failed to persist completed embedding backfill") + async def _prepare_embedding_rebuild(self) -> None: + """Invalidate and persist vectors from the previous vector space.""" + self._embedding_rebuild_pending = True + for chunk in self.file_chunks.values(): + chunk.embedding = None + await self._reset_vector_index() + await self.dump() + + async def _reset_vector_index(self) -> None: + """Drop a derived vector index before rebuilding a changed vector space.""" + async def _after_embedding_backfill(self) -> None: """Backend hook for refreshing derived vector indexes after backfill.""" @@ -540,12 +635,15 @@ class LocalFileStore(BaseFileStore): async def _embed_pending(self, chunks: list[FileChunk]) -> None: if not (chunks and self.embedding_store): return + was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) try: await self.embedding_store.get_node_embeddings(chunks) except Exception as e: - self._disable_embedding(f"upsert: {type(e).__name__}: {e}") + self._mark_embedding_unhealthy(f"upsert: {type(e).__name__}: {e}") return self._drop_stale_embeddings(chunks, "upsert") + if any(chunk.embedding is not None for chunk in chunks): + await self._recover_after_real_request(was_healthy) async def delete(self, path: str | list[str]) -> None: assert self.file_graph is not None @@ -601,21 +699,23 @@ class LocalFileStore(BaseFileStore): # -- 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: + if self.embedding_store is None or self._embedding_rebuild_pending or not query or limit <= 0: return [] + was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) try: query_embedding = await self.embedding_store.get_embedding(query) except Exception as e: - self._disable_embedding(f"search: {type(e).__name__}: {e}") + self._mark_embedding_unhealthy(f"search: {type(e).__name__}: {e}") return [] if query_embedding is None: return [] if not self._embedding_dim_matches(query_embedding): - self._disable_embedding( + self._mark_embedding_unhealthy( f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", ) return [] + await self._recover_after_real_request(was_healthy) top: list[tuple[float, int, FileChunk]] = [] candidates: list[FileChunk] = [] diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py index 80129eed..4b234a3e 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -176,6 +176,10 @@ class ZvecLocalFileStore(LocalFileStore): ] self._upsert_docs(to_add) + async def _reset_vector_index(self) -> None: + """Discard all vectors before rebuilding a changed vector space.""" + self._collection = self._create_collection() + # -- maintenance ------------------------------------------------------ async def optimize_index(self) -> None: @@ -412,20 +416,25 @@ class ZvecLocalFileStore(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 limit <= 0 or self._collection is None or not self._indexed_ids: + if self.embedding_store is None or self._embedding_rebuild_pending or not query or limit <= 0: + return [] + index_empty = self._collection is None or not self._indexed_ids + if index_empty and getattr(self.embedding_store, "is_healthy", True): return [] query_embedding = None + was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) try: query_embedding = await self.embedding_store.get_embedding(query) except Exception as e: - self._disable_embedding(f"search: {type(e).__name__}: {e}") + self._mark_embedding_unhealthy(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( + self._mark_embedding_unhealthy( f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", ) return [] + await self._recover_after_real_request(was_healthy) # get_embedding above yielded control; a concurrent clear() may have # swapped or dropped the collection. Re-read before dereferencing. diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index cd61cbdd..6c49b4ab 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -16,6 +16,7 @@ import pytest 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.components.embedding_store import LocalEmbeddingStore from reme.schema import FileChunk, FileNode from reme.utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst @@ -67,6 +68,7 @@ class CountingFakeEmbeddingStore(FakeEmbeddingStore): def __init__(self): self.node_embedding_calls: list[list[str]] = [] + self.is_healthy = True async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]: self.node_embedding_calls.append([node.id for node in nodes]) @@ -76,10 +78,35 @@ class CountingFakeEmbeddingStore(FakeEmbeddingStore): class UnhealthyCountingEmbeddingStore(CountingFakeEmbeddingStore): """Fake embedding store that fails the backfill health gate.""" + def __init__(self): + super().__init__() + self.is_healthy = False + async def health_check(self, _timeout: float = 2.0) -> bool: return False +class RecoveringEmbeddingStore(CountingFakeEmbeddingStore): + """Fake provider that starts unhealthy and records real recoveries.""" + + def __init__(self): + super().__init__() + self.is_healthy = False + self.health_calls = 0 + + async def health_check(self, _timeout: float = 2.0) -> bool: + self.health_calls += 1 + return False + + async def get_embedding(self, input_text: str, **kwargs) -> np.ndarray: + self.is_healthy = True + return await super().get_embedding(input_text, **kwargs) + + async def get_node_embeddings(self, nodes: list[FileChunk], **kwargs) -> list[FileChunk]: + self.is_healthy = True + return await super().get_node_embeddings(nodes, **kwargs) + + class HealthCountingEmbeddingStore(FakeEmbeddingStore): """Fake provider that records eager health checks.""" @@ -104,6 +131,44 @@ class BlockingEmbeddingStore(FakeEmbeddingStore): return await super().get_node_embeddings(nodes, **kwargs) +class CancellationResistantHealthStore(CountingFakeEmbeddingStore): + """Startup probe that completes stale after cancellation is requested.""" + + def __init__(self): + super().__init__() + self.is_healthy = True + self.health_started = asyncio.Event() + self.release_health = asyncio.Event() + + async def health_check(self, _timeout: float = 2.0) -> bool: + self.health_started.set() + try: + await self.release_health.wait() + except asyncio.CancelledError: + await self.release_health.wait() + self.is_healthy = False + return False + + +class DelayedOldVectorStore(CountingFakeEmbeddingStore): + """First batch returns an old-space vector after rebuild was requested.""" + + def __init__(self): + super().__init__() + self.first_batch_started = asyncio.Event() + self.release_first_batch = asyncio.Event() + + async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]: + self.node_embedding_calls.append([node.id for node in nodes]) + if len(self.node_embedding_calls) == 1: + self.first_batch_started.set() + await self.release_first_batch.wait() + for chunk_node in nodes: + chunk_node.embedding = np.array([0.0, 1.0], dtype=np.float16) + return nodes + return await FakeEmbeddingStore.get_node_embeddings(self, nodes) + + class WrongDimEmbeddingStore(FakeEmbeddingStore): """Fake embedding store that returns vectors with the wrong dimension.""" @@ -153,6 +218,15 @@ def _new_local_store(name, **kwargs): return LocalFileStore(name=name, embedding_store="", **kwargs) +def _new_faiss_store(name, **kwargs): + """Construct a FAISS store when the optional backend is installed.""" + try: + store = FaissLocalFileStore(name=name, embedding_store="", **kwargs) + except ImportError: + pytest.skip("faiss is not installed") + return store + + def _new_zvec_store(name, **kwargs): """Construct a zvec store with embedding disabled at bind time.""" try: @@ -593,7 +667,7 @@ def test_background_embedding_backfill_uses_provider_batch_size(): def test_load_skips_backfill_when_embedding_health_check_fails(): - """Background backfill disables embeddings before batching when the provider is unhealthy.""" + """Background backfill preserves an unhealthy provider for later recovery.""" async def go(): with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): @@ -610,13 +684,249 @@ def test_load_skips_backfill_when_embedding_health_check_fails(): await store._embedding_backfill_task assert not fake.node_embedding_calls - assert store.embedding_store is None + assert store.embedding_store is fake + assert fake.is_healthy is False assert store.file_chunks["a"].embedding is None await store.close() run(go()) +def test_verified_resume_supersedes_inflight_startup_health_check(): + """A stale startup probe cannot consume or overwrite verified recovery.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_embedding_verified_resume_race") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) + fake = CancellationResistantHealthStore() + store.embedding_store = fake + store._start_embedding_backfill() + startup_task = store._embedding_backfill_task + await fake.health_started.wait() + + recovery = asyncio.create_task(store.resume_embedding(verified=True)) + await asyncio.sleep(0) + assert await recovery is True + assert store._embedding_backfill_pending == (True, False) + fake.release_health.set() + + await startup_task + assert store._embedding_backfill_task is not startup_task + if store._embedding_backfill_task is not None: + await store._embedding_backfill_task + assert fake.is_healthy is True + assert fake.node_embedding_calls == [["a"]] + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + await store.close() + + run(go()) + + +def test_verified_resume_without_chunks_supersedes_inflight_health_check(): + """A newer verified state survives a stale probe even after chunks are cleared.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_embedding_empty_verified_resume_race") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) + fake = CancellationResistantHealthStore() + store.embedding_store = fake + store._start_embedding_backfill() + startup_task = store._embedding_backfill_task + await fake.health_started.wait() + + await store.clear() + assert await store.resume_embedding(verified=True) is True + assert store._embedding_backfill_pending == (True, False) + fake.release_health.set() + + await startup_task + assert fake.is_healthy is True + assert not fake.node_embedding_calls + assert store._embedding_backfill_task is None + await store.close() + + run(go()) + + +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) +def test_verified_rebuild_discards_same_dimension_vectors_before_backfill(store_factory): + """A changed vector space never searches compatible-shaped stale vectors.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory("t_embedding_verified_rebuild") + await store.start() + stale = chunk("a", "a.md", "alpha text") + stale.embedding = np.array([0.0, 1.0], dtype=np.float16) + await set_chunks_with_graph(store, {"a": stale}) + fake = CountingFakeEmbeddingStore() + fake.is_healthy = False + store.embedding_store = fake + if isinstance(store, FaissLocalFileStore): + store._rebuild_index() + elif isinstance(store, ZvecLocalFileStore): + store._rebuild_collection() + + assert await store.resume_embedding(verified=True, rebuild=True) is True + assert store._embedding_rebuild_pending is True + assert store.file_chunks["a"].embedding is None + assert await store.vector_search("alpha", 5, {}) == [] + + await store._embedding_backfill_task + assert store._embedding_rebuild_pending is False + assert fake.node_embedding_calls == [["a"]] + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert [item.id for item in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_verified_rebuild_discards_late_result_from_previous_vector_space(): + """A queued rebuild clears old-space vectors written by an in-flight batch.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_embedding_verified_rebuild_race") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) + fake = DelayedOldVectorStore() + store.embedding_store = fake + store._start_embedding_backfill(skip_health_check=True) + old_task = store._embedding_backfill_task + await fake.first_batch_started.wait() + + assert await store.resume_embedding(verified=True, rebuild=True) is True + assert store._embedding_backfill_pending == (True, True) + fake.release_first_batch.set() + + await old_task + if store._embedding_backfill_task is not None: + await store._embedding_backfill_task + assert fake.node_embedding_calls == [["a"], ["a"]] + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert store._embedding_rebuild_pending is False + await store.close() + + run(go()) + + +def test_unverified_rebuild_is_queued_behind_inflight_backfill(): + """An unverified rebuild request cannot be lost while another batch is running.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_embedding_unverified_rebuild_race") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) + fake = DelayedOldVectorStore() + store.embedding_store = fake + store._start_embedding_backfill(skip_health_check=True) + old_task = store._embedding_backfill_task + await fake.first_batch_started.wait() + + assert await store.resume_embedding(rebuild=True) is True + assert store._embedding_backfill_pending == (False, True) + fake.release_first_batch.set() + + await old_task + if store._embedding_backfill_task is not None: + await store._embedding_backfill_task + assert fake.node_embedding_calls == [["a"], ["a"]] + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert store._embedding_rebuild_pending is False + await store.close() + + run(go()) + + +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) +def test_clear_during_scheduled_rebuild_finishes_rebuild_state(store_factory): + """Clearing all chunks before the worker scan must not disable vector search forever.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory("t_embedding_clear_during_rebuild") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) + store.embedding_store = CountingFakeEmbeddingStore() + + assert await store.resume_embedding(verified=True, rebuild=True) is True + task = store._embedding_backfill_task + await store.clear() + if task is not None: + await task + + assert store.file_chunks == {} + assert store._embedding_rebuild_pending is False + await store.close() + + run(go()) + + +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) +def test_search_recovery_schedules_backfill_without_another_health_check(store_factory): + """A successful real search request repairs historical missing vectors.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory(name="t_embedding_search_recovery") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")}) + fake = RecoveringEmbeddingStore() + store.embedding_store = fake + + assert await store.vector_search("alpha", 5, {}) == [] + await store._embedding_backfill_task + + assert fake.is_healthy is True + assert fake.health_calls == 0 + assert fake.node_embedding_calls == [["a"]] + assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0] + assert [item.id for item in await store.vector_search("alpha", 5, {})] == ["a"] + await store.close() + + run(go()) + + +def test_cache_only_search_does_not_mark_provider_recovered(): + """Cached vectors do not prove that the remote provider is available.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + embedding_store = LocalEmbeddingStore(name="t_embedding_cache_only_recovery") + embedding_store.as_embedding = type( + "CachedProvider", + (), + { + "dimensions": 2, + "vector_space_id": "cached-provider", + "__call__": lambda self, texts, **_kwargs: asyncio.sleep( + 0, + result=[[1.0, 0.0] for _ in texts], + ), + }, + )() + await embedding_store.get_embedding("alpha") + embedding_store.is_healthy = False + + store = LocalFileStore(name="t_embedding_cache_only_recovery", embedding_store="") + await store.start() + await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "historical text")}) + store.embedding_store = embedding_store + + assert await store.vector_search("alpha", 5, {}) == [] + assert embedding_store.is_healthy is False + assert store._embedding_backfill_task is None + await store.close() + + run(go()) + + @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.""" @@ -692,7 +1002,8 @@ def test_upsert_drops_wrong_dimension_from_custom_embedding_store(): assert store.file_chunks["a"].embedding is None assert await store.vector_search("alpha", 5, {}) == [] - assert store.embedding_store is None + assert store.embedding_store is not None + assert store.embedding_store.is_healthy is False await store.close() run(go()) diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 5ac1d925..713afaa8 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -39,6 +39,23 @@ class BadHealthAsEmbedding: return [[1.0]] +class FailingHealthAsEmbedding: + """Fake provider that records a failed health-check attempt.""" + + dimensions = 2 + vector_space_id = "fakespace000" + + def __init__(self): + self.calls = 0 + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, _texts: list[str], **_kwargs): + self.calls += 1 + raise ConnectionError("not ready") + + class FakeProviderModel: """Stand-in for a constructed AgentScope embedding model object.""" @@ -186,6 +203,23 @@ def test_health_check_starts_timeout_after_provider_initialization(monkeypatch): run(go()) +def test_health_check_makes_one_attempt(): + """A failed startup probe does not add hidden retries.""" + + async def go(): + provider = FailingHealthAsEmbedding() + store = LocalEmbeddingStore( + name="t_local_embedding_health_retry", + health_check_timeout=3.0, + ) + store.as_embedding = provider + + assert await store.health_check() is False + assert provider.calls == 1 + + run(go()) + + def test_insufficient_quota_waits_sixty_seconds_before_retry(monkeypatch): """Quota exhaustion uses the dedicated delay before ReMe retries."""