From 49c0e7a215d4b8d0a2a0498181c7772a83fb0b94 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Wed, 30 Sep 2026 17:06:04 +0800 Subject: [PATCH 1/6] fix(embedding): keep one bad item from dropping its whole batch A provider-side request error rejects an entire batch, so a single over-limit item used to cost the vectors of every item beside it. QwenPaw #8040 shows two poison chunks turning into exactly twenty lost vectors, while the backfill loop still reported full success. Retry a rejected multi-item request one item at a time so only the item that actually fails is left without a vector. Single-item requests keep their previous behaviour, and provider-level failures (429 / quota) keep using the existing retry paths instead of being split. Also make the backfill loop report what it actually filled, and name the failing chunk ids in the final reindex error. `processed` counts attempts, so it claimed success after a partial loss. --- .../embedding_store/local_embedding_store.py | 23 +++++ .../components/file_store/local_file_store.py | 12 ++- tests/unit/test_file_store_consistency.py | 45 +++++++++ tests/unit/test_local_embedding_store.py | 98 +++++++++++++++++++ 4 files changed, 176 insertions(+), 2 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 4fad8361..f913cb3d 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -215,10 +215,33 @@ class LocalEmbeddingStore(BaseEmbeddingStore): continue self.logger.exception("Embedding request failed") self.is_healthy = False + if len(texts) > 1: + return await self._embed_individually(texts, **kwargs) return None self.is_healthy = False return None + async def _embed_individually(self, texts: list[str], **kwargs) -> list[list[float] | None]: + """Retry a rejected batch one item at a time so one bad item cannot drop the rest. + + Provider-side request errors, such as a per-item context-length limit, reject the + whole batch and would otherwise cost every vector in it. Isolating each item bounds + the damage to the item that actually failed and keeps the remaining vectors usable. + """ + results: list[list[float] | None] = [] + for text in texts: + try: + result = await self.as_embedding([text], **kwargs) + except Exception as error: + self.logger.error( + f"Embedding request failed for a single item (chars={len(text)}): " + f"{type(error).__name__}: {error}", + ) + results.append(None) + continue + results.append(result[0] if result and len(result) == 1 else None) + return results + @staticmethod def _is_rate_limited(error: Exception) -> bool: """Recognize an OpenAI-compatible 429 response without importing a provider SDK.""" diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 243917f8..13e95378 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -29,6 +29,7 @@ _EMBEDDING_F16_DTYPE = np.dtype(" _FAILED_CHUNK_PREVIEW: + preview += f" (+{len(missing) - _FAILED_CHUNK_PREVIEW} more)" + raise RuntimeError( + f"embedding reindex incomplete: {len(missing)} chunks failed: {preview}", + ) await self._finalize_embedding_reindex() await self._dump_owned_state() if self.embedding_store is not None: @@ -530,6 +536,7 @@ class LocalFileStore(BaseFileStore): return processed = 0 + filled = 0 batch_count = 0 embedding_started_at = time.monotonic() next_percent = _PROGRESS_LOG_PERCENT_STEP @@ -538,10 +545,11 @@ class LocalFileStore(BaseFileStore): await self.embedding_store.get_node_embeddings(batch) self._drop_stale_embeddings(batch, "backfill") processed += len(batch) + filled += sum(1 for chunk in batch if chunk.embedding is not None) batch_count += 1 next_percent = self._log_progress("embedding backfill", processed, total, next_percent) self.logger.info( - f"{self.name}: embedding batches complete: processed={processed}/{total}, " + f"{self.name}: embedding batches complete: processed={processed}/{total}, filled={filled}, " f"batches={batch_count}, elapsed={time.monotonic() - embedding_started_at:.3f}s", ) except asyncio.CancelledError: diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index b7750017..9e6af475 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -118,6 +118,16 @@ class RecoveringEmbeddingStore(CountingFakeEmbeddingStore): return await super().get_node_embeddings(nodes, **kwargs) +class PartialEmbeddingStore(FakeEmbeddingStore): + """Fake provider that leaves chunks whose text contains "beta" without a vector.""" + + async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]: + for chunk_node in nodes: + if "beta" not in chunk_node.text: + chunk_node.embedding = self._embed(chunk_node.text) + return nodes + + class HealthCountingEmbeddingStore(FakeEmbeddingStore): """Fake provider that records eager health checks.""" @@ -838,6 +848,41 @@ def test_verified_resume_without_chunks_supersedes_inflight_health_check(): run(go()) +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) +def test_embedding_reindex_error_names_the_chunks_without_vectors(store_factory): + """The final reindex error must name the chunks that still have no vector.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory("t_reindex_failure_ids") + await store.start() + await set_chunks_with_graph( + store, + { + "good": chunk("good", "a.md", "alpha text"), + "bad": chunk("bad", "b.md", "beta text"), + }, + ) + store.embedding_store = PartialEmbeddingStore() + _ensure_zvec_collection(store) + if isinstance(store, FaissLocalFileStore): + store._rebuild_index() + elif isinstance(store, ZvecLocalFileStore): + store._rebuild_collection() + await store.require_embedding_rebuild() + + with pytest.raises(RuntimeError) as failure: + await store.reindex("embedding") + + message = str(failure.value) + assert "1 chunks failed" in message + assert "bad" in message + assert "good" not in message + await store.close() + + run(go()) + + @pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) def test_embedding_reindex_is_explicit_and_keeps_vectors_disabled_until_success( store_factory, diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 51b7bf63..985eba3b 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -126,6 +126,38 @@ class QuotaThenSuccessAsEmbedding: return [[1.0, 0.0] for _ in texts] +class ContextLengthError(Exception): + """OpenAI-compatible 400 raised when one item exceeds the provider's per-item window.""" + + status_code = 400 + body = {"error": {"message": "the input length exceeds the context length"}} + + +class PerItemLimitAsEmbedding: + """Reject a whole request when any single text is over the provider's per-item limit. + + Mirrors an OpenAI-compatible gateway: one over-limit item makes the provider answer + HTTP 400 for the entire batch, so a client that gives up on the batch loses every + vector in it. + """ + + dimensions = 2 + vector_space_id = "fakespace000" + max_chars = 32 + + def __init__(self): + self.batch_sizes: list[int] = [] + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + if any(len(text) > self.max_chars for text in texts): + raise ContextLengthError("the input length exceeds the context length") + return [[1.0, 0.0] for _ in texts] + + class BadNodeEmbeddingStore(BaseEmbeddingStore): """Embedding store that returns wrong-dimensional vectors.""" @@ -730,3 +762,69 @@ def test_start_ignores_cache_file_without_vector_space_tag(monkeypatch, tmp_path assert not store._cache run(go()) + + +def test_batch_failure_falls_back_to_per_item_embedding(): + """One over-limit item must cost one vector, not its whole batch.""" + + async def go(): + provider = PerItemLimitAsEmbedding() + store = LocalEmbeddingStore(name="t_local_embedding_poison_batch", enable_cache=False) + store.as_embedding = provider + + results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0]] + assert provider.batch_sizes == [3, 1, 1, 1] + + run(go()) + + +def test_poison_chunk_keeps_its_batch_mates(): + """A rejected chunk must not clear the vectors of the chunks beside it.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_poison_nodes", enable_cache=False) + store.as_embedding = PerItemLimitAsEmbedding() + nodes = [EmbNode(text="first"), EmbNode(text="y" * 64), EmbNode(text="third")] + + await store.get_node_embeddings(nodes) + + assert nodes[0].embedding is not None + assert nodes[1].embedding is None + assert nodes[2].embedding is not None + + run(go()) + + +def test_single_item_failure_still_reports_a_failed_batch(): + """A one-item request has nothing to isolate, so the existing contract stands.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_single_poison", enable_cache=False) + store.as_embedding = PerItemLimitAsEmbedding() + + assert await store._call_with_retry(["x" * 64]) is None + + run(go()) + + +def test_rate_limited_batch_is_not_split_into_per_item_requests(monkeypatch): + """A 429 is a provider-level failure; splitting it would multiply the load.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore(name="t_local_embedding_rate_limit_batch", max_retries=2) + embedding = AlwaysRateLimitedAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry(["a", "b"]) is None + assert embedding.calls == 2 + assert sleeps == [1.0] + + run(go()) From d5dc9b224dada1ba6fc45494a0978d9c42a931f6 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Thu, 8 Oct 2026 21:35:57 +0800 Subject: [PATCH 2/6] fix(embedding): stop quota failures and transient item errors from multiplying Review follow-up on the two P2 findings in #579: * quota exhaustion no longer falls into the per-item split. The store now gives up explicitly when the provider reports insufficient_quota, including after the opt-in quota retries are spent, because no single item is to blame and splitting only multiplied the rejected requests (3 items -> 4 requests). * the per-item fallback goes back through `_call_with_retry`, so a 429, connection error or timeout on an isolated item keeps the existing backoff and retry policy instead of silently dropping that item's vector. A one-item request still cannot split again. Regression tests cover both findings: a quota failure stays a single request with the default policy and with exhausted quota retries, and a transient failure on one item of a split batch still ends with every healthy sibling embedded. python -m pytest tests/unit/test_local_embedding_store.py tests/unit/test_file_store_consistency.py -q -> 83 passed, 41 skipped --- .../embedding_store/local_embedding_store.py | 24 ++-- tests/unit/test_local_embedding_store.py | 123 ++++++++++++++++++ 2 files changed, 138 insertions(+), 9 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index f913cb3d..1d9c93b6 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -213,6 +213,13 @@ class LocalEmbeddingStore(BaseEmbeddingStore): ) await asyncio.sleep(self.quota_retry_delay) continue + if self._is_insufficient_quota(error): + # Quota is a service-level failure: no single item is to blame, so splitting the + # batch cannot isolate anything and would only multiply the rejected requests. + # This also covers the case where the quota retries above are exhausted. + self.logger.exception("Embedding request failed: quota exhausted") + self.is_healthy = False + return None self.logger.exception("Embedding request failed") self.is_healthy = False if len(texts) > 1: @@ -227,19 +234,18 @@ class LocalEmbeddingStore(BaseEmbeddingStore): Provider-side request errors, such as a per-item context-length limit, reject the whole batch and would otherwise cost every vector in it. Isolating each item bounds the damage to the item that actually failed and keeps the remaining vectors usable. + Each single-item request goes back through ``_call_with_retry`` so a transient + failure (429, connection error, timeout) still gets the existing backoff and retry + policy instead of losing the vector; a one-item request never splits again. """ results: list[list[float] | None] = [] for text in texts: - try: - result = await self.as_embedding([text], **kwargs) - except Exception as error: - self.logger.error( - f"Embedding request failed for a single item (chars={len(text)}): " - f"{type(error).__name__}: {error}", - ) - results.append(None) + single = await self._call_with_retry([text], **kwargs) + if single and len(single) == 1: + results.append(single[0]) continue - results.append(result[0] if result and len(result) == 1 else None) + self.logger.error(f"Embedding request failed for a single item (chars={len(text)})") + results.append(None) return results @staticmethod diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 985eba3b..880ff8fa 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -158,6 +158,48 @@ class PerItemLimitAsEmbedding: return [[1.0, 0.0] for _ in texts] +class AlwaysInsufficientQuotaAsEmbedding: + """Always fail with quota exhaustion, recording every requested batch size.""" + + dimensions = 2 + + def __init__(self): + self.batch_sizes: list[int] = [] + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + raise InsufficientQuotaError("quota exhausted") + + +class PerItemLimitThenTransientErrorAsEmbedding: + """Reject an over-limit batch, then fail the first single-item request once. + + Mirrors the review case: the batch is rejected because one item exceeds the per-item + window, and a sibling item then hits a transient provider failure while retried alone. + """ + + dimensions = 2 + vector_space_id = "fakespace000" + max_chars = 32 + + def __init__(self, transient: Exception): + self.transient = transient + self.batch_sizes: list[int] = [] + self.transient_failures = 0 + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + if any(len(text) > self.max_chars for text in texts): + raise ContextLengthError("the input length exceeds the context length") + if len(texts) == 1 and self.transient_failures == 0: + self.transient_failures += 1 + raise self.transient + return [[1.0, 0.0] for _ in texts] + + class BadNodeEmbeddingStore(BaseEmbeddingStore): """Embedding store that returns wrong-dimensional vectors.""" @@ -828,3 +870,84 @@ def test_rate_limited_batch_is_not_split_into_per_item_requests(monkeypatch): assert sleeps == [1.0] run(go()) + + +def test_insufficient_quota_batch_is_not_split_into_per_item_requests(monkeypatch): + """Quota exhaustion is a service-level failure; splitting it would multiply the load.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore(name="t_local_embedding_quota_batch", max_retries=2) + embedding = AlwaysInsufficientQuotaAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert embedding.batch_sizes == [3] + assert store.is_healthy is False + assert not sleeps + + run(go()) + + +def test_exhausted_quota_retries_still_do_not_split_the_batch(monkeypatch): + """Once the opt-in quota retries run out, the store gives up instead of splitting.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore( + name="t_local_embedding_quota_batch_exhausted", + max_retries=3, + quota_retry_delay=5.0, + ) + embedding = AlwaysInsufficientQuotaAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry(["a", "b"]) is None + assert embedding.batch_sizes == [2, 2, 2] + assert sleeps == [5.0, 5.0] + + run(go()) + + +def test_per_item_fallback_keeps_a_transiently_failed_vector(monkeypatch): + """A 429 / connection error / timeout on one isolated item must not lose its vector.""" + + async def go(): + transients = ( + RateLimitError("Requests are too frequent"), + ConnectionError("connection reset"), + TimeoutError("request timed out"), + ) + for transient in transients: + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = PerItemLimitThenTransientErrorAsEmbedding(transient) + store = LocalEmbeddingStore( + name="t_local_embedding_fallback_transient", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0]] + assert provider.batch_sizes == [3, 1, 1, 1, 1] + assert provider.transient_failures == 1 + assert sleeps == [1.0] + + run(go()) From 4d8a3c54756952542c3025ccbd089b7fdfd8e486 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 10:42:04 +0800 Subject: [PATCH 3/6] style(test): hoist the sleeps collector out of the transient-failure loop CI / Python quality failed on d5dc9b2: the pylint hook reported tests/unit/test_local_embedding_store.py:935:16: W0640: Cell variable sleeps defined in loop (cell-var-from-loop) `fake_sleep` closes over `sleeps` while the loop rebinds it on every iteration. The collector is now created once before the loop and cleared at the top of each iteration, which keeps the per-iteration assertion (`assert sleeps == [1.0]`) exactly as strict and removes the loop-scoped rebinding. Verification (repo hook arguments): pylint -> Your code has been rated at 10.00/10, exit 0 black --check --line-length=120 --target-version=py311 -> clean flake8 --extend-ignore=E203 --max-line-length=120 -> clean --- tests/unit/test_local_embedding_store.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 880ff8fa..da75be04 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -928,11 +928,13 @@ def test_per_item_fallback_keeps_a_transiently_failed_vector(monkeypatch): ConnectionError("connection reset"), TimeoutError("request timed out"), ) - for transient in transients: - sleeps = [] + sleeps = [] - async def fake_sleep(delay): - sleeps.append(delay) + async def fake_sleep(delay): + sleeps.append(delay) + + for transient in transients: + sleeps.clear() provider = PerItemLimitThenTransientErrorAsEmbedding(transient) store = LocalEmbeddingStore( From c9448aeb19e753767d69c4fffa11a13a70e78086 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 12:21:52 +0800 Subject: [PATCH 4/6] fix(embedding): keep provider-wide failures out of the per-item split Review follow-up on the P2 raised against 4d8a3c5. `_call_with_retry` sent every unrecognized exception to `_embed_individually` whenever the request had more than one item. Only rate limiting and `insufficient_quota` were excluded, so a 401, a 403 or a 5xx still split a ten-item request into ten extra single-item requests, none of which can succeed: credentials, permissions and server outages are properties of the endpoint, not of any individual input. The split path is now gated on a dedicated classifier, `_is_provider_wide_failure()`, which recognizes 401/403 and 5xx by status code (with the OpenAI SDK class names as a fallback for exceptions that do not expose one). Transient server errors keep their retry policy through the existing backoff; the remaining provider-wide failures stop the loop and report unhealthy, like the quota branch. Input-dependent rejections such as the per-item context-length limit are unchanged and still isolate one item at a time. Regression coverage: test_provider_wide_credentials_errors_are_not_split_into_per_item_requests (401/403 -> a single request of size 10) and test_provider_wide_server_errors_are_retried_without_splitting (500 -> batch sizes [10, 10, 10] with sleeps [1.0, 2.0]). The existing context-length, quota, 429 and per-item transient-retry tests are retained. Local verification, using the repo's own hook arguments: pylint -> Your code has been rated at 10.00/10, exit 0 black -> 2 files would be left unchanged flake8 -> clean Classifier discriminator check (old vs new split decision, extracted from the production file and executed): 400 context-length -> SPLIT / SPLIT 401, 403, 500, 504 -> SPLIT / NO-SPLIT 429 and insufficient_quota keep their own earlier branches. --- .../embedding_store/local_embedding_store.py | 33 +++++++++ tests/unit/test_local_embedding_store.py | 73 +++++++++++++++++++ 2 files changed, 106 insertions(+) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 1d9c93b6..d5e26e23 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -220,6 +220,17 @@ class LocalEmbeddingStore(BaseEmbeddingStore): self.logger.exception("Embedding request failed: quota exhausted") self.is_healthy = False return None + if self._is_provider_wide_failure(error): + # Credentials, permissions and server outages describe the endpoint, not any + # single item: isolation cannot fix them, so retry the transient ones with the + # usual backoff and give up on the rest instead of multiplying the requests. + if self._is_transient_provider_failure(error) and attempt < self.max_retries - 1: + delay = 2**attempt + self.logger.warning(f"Embedding provider error; retrying in {delay:.1f}s") + await asyncio.sleep(delay) + continue + self.logger.exception("Embedding request failed: provider-wide error, not splitting") + break self.logger.exception("Embedding request failed") self.is_healthy = False if len(texts) > 1: @@ -248,6 +259,28 @@ class LocalEmbeddingStore(BaseEmbeddingStore): results.append(None) return results + @staticmethod + def _is_provider_wide_failure(error: Exception) -> bool: + """Return whether an error describes the provider rather than a single input. + + Authentication, permission and server errors belong to the endpoint or to the + credentials: no individual input can be blamed, so splitting the batch cannot isolate + anything and would only multiply the rejected requests. Rate limits and quota + exhaustion have their own branches above and never reach this classifier. + """ + status = getattr(error, "status_code", None) + if isinstance(status, int) and (status in (401, 403) or 500 <= status < 600): + return True + return type(error).__name__ in ("AuthenticationError", "PermissionDeniedError", "InternalServerError") + + @staticmethod + def _is_transient_provider_failure(error: Exception) -> bool: + """Return whether a provider-wide failure is worth retrying with the usual backoff.""" + status = getattr(error, "status_code", None) + if isinstance(status, int): + return 500 <= status < 600 + return type(error).__name__ == "InternalServerError" + @staticmethod def _is_rate_limited(error: Exception) -> bool: """Recognize an OpenAI-compatible 429 response without importing a provider SDK.""" diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index da75be04..926b872c 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -200,6 +200,38 @@ class PerItemLimitThenTransientErrorAsEmbedding: return [[1.0, 0.0] for _ in texts] +class AuthenticationError(Exception): + """OpenAI-compatible 401 error used without importing the provider SDK.""" + + status_code = 401 + + +class PermissionDeniedError(Exception): + """OpenAI-compatible 403 error used without importing the provider SDK.""" + + status_code = 403 + + +class InternalServerError(Exception): + """OpenAI-compatible 500 error used without importing the provider SDK.""" + + status_code = 500 + + +class AlwaysProviderWideFailureAsEmbedding: + """Always fail with the same provider-wide error, recording every requested batch size.""" + + dimensions = 2 + + def __init__(self, error: Exception): + self.error = error + self.batch_sizes: list[int] = [] + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + raise self.error + + class BadNodeEmbeddingStore(BaseEmbeddingStore): """Embedding store that returns wrong-dimensional vectors.""" @@ -953,3 +985,44 @@ def test_per_item_fallback_keeps_a_transiently_failed_vector(monkeypatch): assert sleeps == [1.0] run(go()) + + +def test_provider_wide_credentials_errors_are_not_split_into_per_item_requests(): + """401/403 belong to the credentials: splitting ten items only multiplies the failures.""" + + async def go(): + for error in ( + AuthenticationError("invalid api key"), + PermissionDeniedError("no access to model"), + ): + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore(name="t_local_embedding_provider_wide", max_retries=3) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10] + assert store.is_healthy is False + + run(go()) + + +def test_provider_wide_server_errors_are_retried_without_splitting(monkeypatch): + """A 500 is transient, so it is retried with the usual backoff and never split per item.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = AlwaysProviderWideFailureAsEmbedding(InternalServerError("upstream unavailable")) + store = LocalEmbeddingStore(name="t_local_embedding_provider_server", max_retries=3) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10, 10, 10] + assert sleeps == [1.0, 2.0] + assert store.is_healthy is False + + run(go()) From 055ac9dbb5462c15c9b368a6f2429e2718338c48 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 13:24:22 +0800 Subject: [PATCH 5/6] fix(embedding): retry SDK transport failures instead of splitting the batch Third review follow-up on #579. `_call_with_retry` only caught the built-in `TimeoutError`, `ConnectionError` and `OSError`. `openai.APIConnectionError` and `openai.APITimeoutError` do not inherit from any of them (they go APIError -> OpenAIError -> Exception), so they reached the generic handler: `_is_provider_wide_failure()` did not recognize them either, and a multi-item request fanned out into per-item requests that failed for exactly the same reason. The same gap meant a single-item request got no ReMe-level retry at all after a transport failure. Transport failures are now recognized by a dedicated classifier, `_is_transient_transport_failure()`, covering the built-in socket errors as well as the SDK/httpx exception types (matched across the MRO, so subclasses such as `APITimeoutError` are included) without importing the provider SDK. They retry the original request with the existing backoff and never split; the retry budget is unchanged. Regression coverage: - test_sdk_transport_errors_are_retried_without_splitting APIConnectionError / APITimeoutError -> batch sizes [10, 10, 10], sleeps [1.0, 2.0], then failure with no single-item requests. - test_sdk_transport_error_recovers_without_losing_the_batch one APITimeoutError then success -> [[1.0, 0.0]] * 10, batch sizes [10, 10]. - test_isolated_item_retries_an_sdk_transport_error context-length rejection, then APIConnectionError on one isolated item -> that item's vector is retained. The existing context-length, 401/403, 5xx, quota and rate-limit tests are retained. Local verification with the repo's own hook arguments: pylint -> Your code has been rated at 10.00/10, exit 0 black -> 2 files would be left unchanged flake8 -> clean Branch-order discriminator check (classifiers extracted from the production file and executed against each error type): case old new 400 context-length SPLIT SPLIT 401 / 403 SPLIT STOP 500 SPLIT RETRY SDK APIConnectionError SPLIT RETRY SDK APITimeoutError SPLIT RETRY builtin ConnectionError RETRY RETRY builtin TimeoutError RETRY RETRY 429 RETRY RETRY insufficient_quota STOP STOP --- .../embedding_store/local_embedding_store.py | 41 ++++++- tests/unit/test_local_embedding_store.py | 108 ++++++++++++++++++ 2 files changed, 146 insertions(+), 3 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index d5e26e23..c0f1e257 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -14,6 +14,25 @@ from ..as_embedding import BaseAsEmbedding Miss = tuple[int, str, str] # (result_index, text, cache_key) _MAX_VECTOR_SPACE_ATTEMPTS = 3 +# Transport failures raised by the OpenAI SDK and httpx. They do not inherit from the built-in +# TimeoutError / ConnectionError / OSError, so the built-in check alone lets them slip into the +# generic handler and fan out into per-item requests. Matching on class names (across the MRO, +# so subclasses such as APITimeoutError are covered) keeps this module free of SDK imports. +_TRANSPORT_FAILURE_TYPES = frozenset( + { + "APIConnectionError", + "APITimeoutError", + "ConnectError", + "ConnectTimeout", + "PoolTimeout", + "ReadError", + "ReadTimeout", + "RemoteProtocolError", + "WriteError", + "WriteTimeout", + }, +) + @R.register("local") class LocalEmbeddingStore(BaseEmbeddingStore): @@ -190,10 +209,14 @@ class LocalEmbeddingStore(BaseEmbeddingStore): result = await self.as_embedding(texts, **kwargs) if result and len(result) == len(texts): return result - except (TimeoutError, ConnectionError, OSError): - if attempt < self.max_retries - 1: - await asyncio.sleep(2**attempt) except Exception as error: + if self._is_transient_transport_failure(error): + # A connection failure or a timeout describes the endpoint, not any single + # item: retry the very same request with the usual backoff instead of + # fanning out into per-item requests that would fail the same way. + if attempt < self.max_retries - 1: + await asyncio.sleep(2**attempt) + continue if self._is_rate_limited(error): if attempt < self.max_retries - 1: delay = 2**attempt @@ -259,6 +282,18 @@ class LocalEmbeddingStore(BaseEmbeddingStore): results.append(None) return results + @staticmethod + def _is_transient_transport_failure(error: Exception) -> bool: + """Return whether an error is a transient transport failure worth retrying. + + Covers the built-in socket errors as well as the provider/httpx exceptions that describe + a connection failure or a timeout. Transport failures are provider-wide: no individual + input can be blamed, so the same request is retried and never split per item. + """ + if isinstance(error, (TimeoutError, ConnectionError, OSError)): + return True + return any(cls.__name__ in _TRANSPORT_FAILURE_TYPES for cls in type(error).__mro__) + @staticmethod def _is_provider_wide_failure(error: Exception) -> bool: """Return whether an error describes the provider rather than a single input. diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 926b872c..9e01215a 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -218,6 +218,18 @@ class InternalServerError(Exception): status_code = 500 +class APIError(Exception): + """Base of the OpenAI SDK error hierarchy, reproduced without importing the SDK.""" + + +class APIConnectionError(APIError): + """OpenAI SDK connection failure: it does not inherit from the built-in ConnectionError.""" + + +class APITimeoutError(APIConnectionError): + """OpenAI SDK timeout, a subclass of APIConnectionError.""" + + class AlwaysProviderWideFailureAsEmbedding: """Always fail with the same provider-wide error, recording every requested batch size.""" @@ -232,6 +244,27 @@ class AlwaysProviderWideFailureAsEmbedding: raise self.error +class TransientThenSuccessAsEmbedding: + """Fail the first `failures` calls with one error, then return valid embeddings.""" + + dimensions = 2 + vector_space_id = "fakespace000" + + def __init__(self, error: Exception, failures: int = 1): + self.error = error + self.failures = failures + self.batch_sizes: list[int] = [] + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + if len(self.batch_sizes) <= self.failures: + raise self.error + return [[1.0, 0.0] for _ in texts] + + class BadNodeEmbeddingStore(BaseEmbeddingStore): """Embedding store that returns wrong-dimensional vectors.""" @@ -1026,3 +1059,78 @@ def test_provider_wide_server_errors_are_retried_without_splitting(monkeypatch): assert store.is_healthy is False run(go()) + + +def test_sdk_transport_errors_are_retried_without_splitting(monkeypatch): + """SDK connection/timeout errors are provider-wide: retry the batch, never split it.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + for error in (APIConnectionError("connection reset"), APITimeoutError("timed out")): + sleeps.clear() + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore(name="t_local_embedding_sdk_transport", max_retries=3) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10, 10, 10] + assert sleeps == [1.0, 2.0] + + run(go()) + + +def test_sdk_transport_error_recovers_without_losing_the_batch(monkeypatch): + """One SDK transport failure retries the original batch and keeps every vector.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = TransientThenSuccessAsEmbedding(APITimeoutError("timed out")) + store = LocalEmbeddingStore(name="t_local_embedding_sdk_recover", max_retries=3) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + results = await store._call_with_retry([f"item{i}" for i in range(10)]) + + assert results == [[1.0, 0.0]] * 10 + assert provider.batch_sizes == [10, 10] + assert sleeps == [1.0] + + run(go()) + + +def test_isolated_item_retries_an_sdk_transport_error(monkeypatch): + """After a context-length rejection, one isolated item still retries an SDK transport error.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = PerItemLimitThenTransientErrorAsEmbedding(APIConnectionError("connection reset")) + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_isolated", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0]] + assert provider.batch_sizes == [3, 1, 1, 1, 1] + assert provider.transient_failures == 1 + assert sleeps == [1.0] + + run(go()) From c480fd9fdbce70ae9f88bfe9ae7e399e4ea6e77e Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 14:30:51 +0800 Subject: [PATCH 6/6] fix(embedding): treat 404 as a provider-wide failure too Fourth review follow-up on #579. The classifier added in c9448ae covered 401/403 and 5xx, but not 404. A missing model or a wrong embedding route answers with 404 (SDK `NotFoundError`), which no individual input can explain, so the batch still fell through to `_embed_individually()`: a ten-item request against a bad model name produced eleven requests, all failing for the same reason. 404 is now part of the provider-wide set (`_PROVIDER_WIDE_STATUS_CODES`) and `NotFoundError` joins the class-name fallback. Like 401/403 it is not transient, so the request is reported as failed once instead of being retried or split. The context-length `400` is deliberately left out of that set: an over-limit item is the one status a single input can explain by itself, so it has to stay isolatable. Regression coverage: - test_provider_wide_statuses_exclude_input_dependent_bad_request 401/403/404/500 -> provider-wide; an over-limit 400 -> not provider-wide. - test_provider_wide_not_found_errors_are_not_split 404 -> batch sizes [10], no sleeps, no single-item requests. The existing context-length isolation, 401/403, 5xx, quota, rate-limit and SDK transport tests are retained. Local verification with the repo's own hook arguments: pylint -> Your code has been rated at 10.00/10, exit 0 black -> 2 files would be left unchanged flake8 -> clean Branch-order discriminator check (classifiers extracted from the production file and executed against each error type): case old new 400 context-length SPLIT SPLIT 401 / 403 SPLIT STOP 404 NotFoundError SPLIT STOP 500 SPLIT RETRY SDK APIConnectionError SPLIT RETRY SDK APITimeoutError SPLIT RETRY builtin ConnectionError RETRY RETRY builtin TimeoutError RETRY RETRY 429 RETRY RETRY insufficient_quota STOP STOP --- .../embedding_store/local_embedding_store.py | 21 ++++++--- tests/unit/test_local_embedding_store.py | 45 +++++++++++++++++++ 2 files changed, 60 insertions(+), 6 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index c0f1e257..75b8dbdf 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -14,6 +14,8 @@ from ..as_embedding import BaseAsEmbedding Miss = tuple[int, str, str] # (result_index, text, cache_key) _MAX_VECTOR_SPACE_ATTEMPTS = 3 +_PROVIDER_WIDE_STATUS_CODES = frozenset({401, 403, 404}) + # Transport failures raised by the OpenAI SDK and httpx. They do not inherit from the built-in # TimeoutError / ConnectionError / OSError, so the built-in check alone lets them slip into the # generic handler and fan out into per-item requests. Matching on class names (across the MRO, @@ -298,15 +300,22 @@ class LocalEmbeddingStore(BaseEmbeddingStore): def _is_provider_wide_failure(error: Exception) -> bool: """Return whether an error describes the provider rather than a single input. - Authentication, permission and server errors belong to the endpoint or to the - credentials: no individual input can be blamed, so splitting the batch cannot isolate - anything and would only multiply the rejected requests. Rate limits and quota - exhaustion have their own branches above and never reach this classifier. + Authentication, permission, missing model or route and server errors belong to the + endpoint, the credentials or the deployment: no individual input can be blamed, so + splitting the batch cannot isolate anything and would only multiply the rejected + requests. Rate limits and quota exhaustion have their own branches above and never + reach this classifier. A context-length ``400`` stays input-dependent on purpose, + because that is the one status an over-limit item can explain by itself. """ status = getattr(error, "status_code", None) - if isinstance(status, int) and (status in (401, 403) or 500 <= status < 600): + if isinstance(status, int) and (status in _PROVIDER_WIDE_STATUS_CODES or 500 <= status < 600): return True - return type(error).__name__ in ("AuthenticationError", "PermissionDeniedError", "InternalServerError") + return type(error).__name__ in ( + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "InternalServerError", + ) @staticmethod def _is_transient_provider_failure(error: Exception) -> bool: diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 9e01215a..8d3053e5 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -218,6 +218,12 @@ class InternalServerError(Exception): status_code = 500 +class NotFoundError(Exception): + """OpenAI-compatible 404 error, e.g. a missing model or endpoint.""" + + status_code = 404 + + class APIError(Exception): """Base of the OpenAI SDK error hierarchy, reproduced without importing the SDK.""" @@ -1134,3 +1140,42 @@ def test_isolated_item_retries_an_sdk_transport_error(monkeypatch): assert sleeps == [1.0] run(go()) + + +def test_provider_wide_statuses_exclude_input_dependent_bad_request(): + """404 is provider-wide, but an over-limit 400 must stay input-dependent.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_classifier") + for error in ( + AuthenticationError("bad key"), + PermissionDeniedError("no access"), + NotFoundError("model_not_found"), + InternalServerError("boom"), + ): + assert store._is_provider_wide_failure(error) is True + over_limit = ContextLengthError("the input length exceeds the context length") + assert store._is_provider_wide_failure(over_limit) is False + + run(go()) + + +def test_provider_wide_not_found_errors_are_not_split(monkeypatch): + """A missing model or endpoint is provider-wide: one request, then failure, never a split.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = AlwaysProviderWideFailureAsEmbedding(NotFoundError("model_not_found")) + store = LocalEmbeddingStore(name="t_local_embedding_not_found", max_retries=3) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10] + assert not sleeps + + run(go())