From 530ef1ea9587e3c0faa1f7c7729c25c20fbb5568 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Thu, 20 Aug 2026 19:38:10 +0800 Subject: [PATCH] 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 --- .../embedding_store/base_embedding_store.py | 3 -- .../embedding_store/local_embedding_store.py | 1 - .../file_store/faiss_local_file_store.py | 3 +- .../components/file_store/local_file_store.py | 38 ++++--------------- .../file_store/zvec_local_file_store.py | 3 +- tests/unit/test_file_store_consistency.py | 3 -- 6 files changed, 9 insertions(+), 42 deletions(-) diff --git a/reme/components/embedding_store/base_embedding_store.py b/reme/components/embedding_store/base_embedding_store.py index a6f9e9cc..a8cd6fa9 100644 --- a/reme/components/embedding_store/base_embedding_store.py +++ b/reme/components/embedding_store/base_embedding_store.py @@ -43,9 +43,6 @@ class BaseEmbeddingStore(BaseComponent): self.quota_retry_delay = quota_retry_delay self.health_check_timeout = health_check_timeout self.is_healthy: bool = True - # Monotonic signal used by file stores to distinguish real provider - # recovery from cache-only results. - self.provider_success_count: int = 0 def _truncate(self, text: str) -> str: """Truncate text using a CJK-aware character budget. diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index f08ff531..81167560 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -153,7 +153,6 @@ class LocalEmbeddingStore(BaseEmbeddingStore): 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.provider_success_count += 1 self.is_healthy = True else: self.is_healthy = False diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 4122cb59..48174e4e 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -635,7 +635,6 @@ class FaissLocalFileStore(LocalFileStore): return [] query_embedding = None - provider_success_count = self._provider_success_count() was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) try: query_embedding = await self.embedding_store.get_embedding(query) @@ -647,7 +646,7 @@ class FaissLocalFileStore(LocalFileStore): f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", ) return [] - self._recover_after_real_request(provider_success_count, was_healthy, True) + 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 95cc320a..6a5294c0 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -109,31 +109,12 @@ class LocalFileStore(BaseFileStore): self.embedding_store.is_healthy = False self.logger.error(f"{self.name}: embedding unavailable, {reason}; keyword search remains active") - def _provider_success_count(self) -> int | None: - if self.embedding_store is None: - return None - value = getattr(self.embedding_store, "provider_success_count", None) - return value if isinstance(value, int) else None - - def _recover_after_real_request( - self, - previous_count: int | None, - was_healthy: bool, - valid_result: bool, - ) -> None: + 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 not valid_result: + if self.embedding_store is None or was_healthy or not getattr(self.embedding_store, "is_healthy", True): return - current_count = self._provider_success_count() - provider_succeeded = ( - current_count > previous_count if current_count is not None and previous_count is not None else True - ) - if not provider_succeeded: - return - self.embedding_store.is_healthy = True - if not was_healthy: - self.logger.info(f"{self.name}: embedding provider recovered; scheduling missing-vector backfill") - self._start_embedding_backfill(skip_health_check=True) + self.logger.info(f"{self.name}: embedding provider recovered; scheduling missing-vector backfill") + self._start_embedding_backfill(skip_health_check=True) async def resume_embedding(self, *, verified: bool = False) -> bool: """Resume a configured provider and schedule a deduplicated repair. @@ -585,7 +566,6 @@ class LocalFileStore(BaseFileStore): async def _embed_pending(self, chunks: list[FileChunk]) -> None: if not (chunks and self.embedding_store): return - provider_success_count = self._provider_success_count() was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) try: await self.embedding_store.get_node_embeddings(chunks) @@ -593,11 +573,8 @@ class LocalFileStore(BaseFileStore): self._mark_embedding_unhealthy(f"upsert: {type(e).__name__}: {e}") return self._drop_stale_embeddings(chunks, "upsert") - self._recover_after_real_request( - provider_success_count, - was_healthy, - any(chunk.embedding is not None for chunk in chunks), - ) + if any(chunk.embedding is not None for chunk in chunks): + self._recover_after_real_request(was_healthy) async def delete(self, path: str | list[str]) -> None: assert self.file_graph is not None @@ -656,7 +633,6 @@ class LocalFileStore(BaseFileStore): if self.embedding_store is None or not query or limit <= 0: return [] - provider_success_count = self._provider_success_count() was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) try: query_embedding = await self.embedding_store.get_embedding(query) @@ -670,7 +646,7 @@ class LocalFileStore(BaseFileStore): f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", ) return [] - self._recover_after_real_request(provider_success_count, was_healthy, True) + 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 3dbd3457..652795c0 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -419,7 +419,6 @@ class ZvecLocalFileStore(LocalFileStore): return [] query_embedding = None - provider_success_count = self._provider_success_count() was_healthy = bool(getattr(self.embedding_store, "is_healthy", True)) try: query_embedding = await self.embedding_store.get_embedding(query) @@ -431,7 +430,7 @@ class ZvecLocalFileStore(LocalFileStore): f"search: query embedding dimension {len(query_embedding)} != {self.embedding_store.dimensions}", ) return [] - self._recover_after_real_request(provider_success_count, was_healthy, True) + 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 83cac4e0..ba011912 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -91,7 +91,6 @@ class RecoveringEmbeddingStore(CountingFakeEmbeddingStore): def __init__(self): super().__init__() self.is_healthy = False - self.provider_success_count = 0 self.health_calls = 0 async def health_check(self, _timeout: float = 2.0) -> bool: @@ -99,12 +98,10 @@ class RecoveringEmbeddingStore(CountingFakeEmbeddingStore): return False async def get_embedding(self, input_text: str, **kwargs) -> np.ndarray: - self.provider_success_count += 1 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.provider_success_count += 1 self.is_healthy = True return await super().get_node_embeddings(nodes, **kwargs)