mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
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:
parent
eb3b4e5545
commit
530ef1ea95
6 changed files with 9 additions and 42 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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] = []
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue