mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
fix: recover embedding after transient health failure
This commit is contained in:
parent
f44f52d919
commit
eb3b4e5545
13 changed files with 284 additions and 38 deletions
|
|
@ -117,6 +117,10 @@ 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.
|
||||
|
||||
## How to Search
|
||||
|
||||
The `search` Job is also configured in `default.yaml`:
|
||||
|
|
|
|||
|
|
@ -106,6 +106,9 @@ file_store:
|
|||
所以开箱搜索主要是 BM25 + 链接展开。把 `embedding_store: default` 打开后,`SearchStep` 会同时跑向量召回和关键词召回。此时若将
|
||||
`file_store` 的 `backend` 从 `local` 改为 `faiss`,向量检索会从线性扫描升级为 FAISS HNSW 索引,在大规模 chunk 场景下召回效率更高。
|
||||
|
||||
Embedding store 可通过 `health_check_timeout` 配置启动探测。临时失败只会跳过本次向量回填,BM25 仍可使用;
|
||||
后续真实请求成功后会自动恢复缺失向量的回填。
|
||||
|
||||
## 怎么搜索
|
||||
|
||||
`search` Job 也是在 `default.yaml` 中配置:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""Base embedding store with abstract interface for caching and retrieval."""
|
||||
|
||||
from abc import abstractmethod
|
||||
import math
|
||||
import unicodedata
|
||||
|
||||
import numpy as np
|
||||
|
|
@ -25,14 +26,26 @@ class BaseEmbeddingStore(BaseComponent):
|
|||
max_input_length: int = 8192,
|
||||
max_retries: int = 3,
|
||||
quota_retry_delay: float | None = None,
|
||||
health_check_timeout: float = 5.0,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
if (
|
||||
isinstance(health_check_timeout, bool)
|
||||
or not isinstance(health_check_timeout, (int, float))
|
||||
or not math.isfinite(health_check_timeout)
|
||||
or health_check_timeout <= 0
|
||||
):
|
||||
raise ValueError("health_check_timeout must be finite and greater than 0")
|
||||
self.max_batch_size = max_batch_size
|
||||
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
|
||||
# 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.
|
||||
|
|
@ -57,7 +70,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:
|
||||
|
|
|
|||
|
|
@ -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,11 @@ 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.provider_success_count += 1
|
||||
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 +180,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
|
||||
|
|
|
|||
|
|
@ -625,26 +625,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
|
||||
if (
|
||||
self.embedding_store is None
|
||||
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
|
||||
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)
|
||||
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 []
|
||||
self._recover_after_real_request(provider_success_count, was_healthy, True)
|
||||
|
||||
# get_embedding above yielded control; a concurrent clear() drops the
|
||||
# index to None once embedding is disabled, and a reindex may have swapped
|
||||
|
|
|
|||
|
|
@ -67,10 +67,12 @@ 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._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 +96,57 @@ 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")
|
||||
|
||||
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:
|
||||
"""Schedule repair when a real, non-cache provider request recovers."""
|
||||
if self.embedding_store is None or not valid_result:
|
||||
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)
|
||||
|
||||
async def resume_embedding(self, *, verified: 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.
|
||||
"""
|
||||
if self.embedding_store is None or self._closing:
|
||||
return False
|
||||
if verified:
|
||||
self.embedding_store.is_healthy = True
|
||||
self._start_embedding_backfill(skip_health_check=verified)
|
||||
return True
|
||||
|
||||
def _embedding_dim_matches(self, embedding: np.ndarray | None) -> bool:
|
||||
"""Return whether an index embedding matches the active embedding model."""
|
||||
|
|
@ -228,9 +270,12 @@ 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) -> 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, "
|
||||
|
|
@ -250,7 +295,7 @@ class LocalFileStore(BaseFileStore):
|
|||
)
|
||||
return
|
||||
self._embedding_backfill_task = asyncio.create_task(
|
||||
self._backfill_missing_embeddings(),
|
||||
self._backfill_missing_embeddings(skip_health_check=skip_health_check),
|
||||
name=f"embedding-backfill:{self.name}",
|
||||
)
|
||||
self.logger.info(
|
||||
|
|
@ -285,7 +330,7 @@ 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()
|
||||
if not self.embedding_store or not self.file_chunks:
|
||||
|
|
@ -313,18 +358,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
|
||||
|
|
@ -349,7 +394,7 @@ class LocalFileStore(BaseFileStore):
|
|||
)
|
||||
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}/"
|
||||
|
|
@ -540,12 +585,19 @@ 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)
|
||||
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")
|
||||
self._recover_after_real_request(
|
||||
provider_success_count,
|
||||
was_healthy,
|
||||
any(chunk.embedding is not None for chunk in chunks),
|
||||
)
|
||||
|
||||
async def delete(self, path: str | list[str]) -> None:
|
||||
assert self.file_graph is not None
|
||||
|
|
@ -604,18 +656,21 @@ 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)
|
||||
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 []
|
||||
self._recover_after_real_request(provider_success_count, was_healthy, True)
|
||||
|
||||
top: list[tuple[float, int, FileChunk]] = []
|
||||
candidates: list[FileChunk] = []
|
||||
|
|
|
|||
|
|
@ -412,20 +412,26 @@ 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 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
|
||||
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)
|
||||
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 []
|
||||
self._recover_after_real_request(provider_success_count, was_healthy, True)
|
||||
|
||||
# get_embedding above yielded control; a concurrent clear() may have
|
||||
# swapped or dropped the collection. Re-read before dereferencing.
|
||||
|
|
|
|||
|
|
@ -497,6 +497,7 @@ components:
|
|||
default:
|
||||
backend: local
|
||||
as_embedding: default
|
||||
health_check_timeout: 5.0
|
||||
|
||||
as_llm:
|
||||
default:
|
||||
|
|
|
|||
|
|
@ -453,6 +453,7 @@ components:
|
|||
# as_embedding: default
|
||||
# max_retries: 3
|
||||
# quota_retry_delay: 60.0
|
||||
# health_check_timeout: 5.0
|
||||
|
||||
file_graph:
|
||||
default:
|
||||
|
|
|
|||
|
|
@ -740,6 +740,7 @@ components:
|
|||
# default:
|
||||
# backend: local
|
||||
# as_embedding: default
|
||||
# health_check_timeout: 5.0
|
||||
|
||||
as_llm:
|
||||
default:
|
||||
|
|
|
|||
|
|
@ -490,6 +490,7 @@ components:
|
|||
default:
|
||||
backend: local
|
||||
as_embedding: default
|
||||
health_check_timeout: 5.0
|
||||
|
||||
as_llm:
|
||||
default:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -76,10 +77,38 @@ 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.provider_success_count = 0
|
||||
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.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)
|
||||
|
||||
|
||||
class HealthCountingEmbeddingStore(FakeEmbeddingStore):
|
||||
"""Fake provider that records eager health checks."""
|
||||
|
||||
|
|
@ -153,6 +182,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 +631,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 +648,73 @@ 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())
|
||||
|
||||
|
||||
@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 +790,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())
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import asyncio
|
|||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from reme.components.as_embedding import DashScopeAsEmbedding, OllamaAsEmbedding, OpenAIAsEmbedding
|
||||
from reme.components.embedding_store.base_embedding_store import BaseEmbeddingStore
|
||||
|
|
@ -39,6 +40,20 @@ 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
|
||||
|
||||
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."""
|
||||
|
||||
|
|
@ -183,9 +198,37 @@ def test_health_check_starts_timeout_after_provider_initialization(monkeypatch):
|
|||
assert await store.health_check(timeout=5.0) is True
|
||||
assert events == ["initialized", "remote request"]
|
||||
|
||||
|
||||
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())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "message"),
|
||||
[
|
||||
({"health_check_timeout": 0}, "health_check_timeout"),
|
||||
({"health_check_timeout": float("inf")}, "health_check_timeout"),
|
||||
],
|
||||
)
|
||||
def test_health_check_config_rejects_invalid_values(kwargs, message):
|
||||
"""Invalid probe policies fail during component construction."""
|
||||
with pytest.raises(ValueError, match=message):
|
||||
LocalEmbeddingStore(name="t_local_embedding_invalid_health_config", **kwargs)
|
||||
|
||||
|
||||
def test_insufficient_quota_waits_sixty_seconds_before_retry(monkeypatch):
|
||||
"""Quota exhaustion uses the dedicated delay before ReMe retries."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue