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
This commit is contained in:
jinli.yl 2026-08-20 19:38:10 +08:00
parent eb3b4e5545
commit 530ef1ea95
6 changed files with 9 additions and 42 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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