From ab96ab680a0157d0e07c63a7301aefd9541a5378 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Wed, 30 Sep 2026 17:06:04 +0800 Subject: [PATCH 01/12] 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 e5448f65..dcbe6177 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -43,6 +43,7 @@ _SEARCH_FILTER_RESERVED_KEYS = frozenset( "strict_date_filter", }, ) +_FAILED_CHUNK_PREVIEW = 5 @R.register("local") @@ -440,7 +441,12 @@ class LocalFileStore(BaseFileStore): if chunk.text and not self._embedding_dim_matches(chunk.embedding) ] if missing: - raise RuntimeError(f"embedding reindex incomplete: {len(missing)} chunks failed") + preview = ", ".join(missing[:_FAILED_CHUNK_PREVIEW]) + if len(missing) > _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: @@ -544,6 +550,7 @@ class LocalFileStore(BaseFileStore): return processed = 0 + filled = 0 batch_count = 0 embedding_started_at = time.monotonic() next_percent = _PROGRESS_LOG_PERCENT_STEP @@ -552,10 +559,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 bc47906f79e41e918bd618372b1d34fd760e0524 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Thu, 8 Oct 2026 21:35:57 +0800 Subject: [PATCH 02/12] 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 36eae4f033537176bc8296c26a8ac200e2631931 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 10:42:04 +0800 Subject: [PATCH 03/12] 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 254d895f965ae33a849721799d6745be1091bd9e Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 12:21:52 +0800 Subject: [PATCH 04/12] 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 930fddc4d25f3977fb05a80c27faada3800bf1ce Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 13:24:22 +0800 Subject: [PATCH 05/12] 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 eba753f7837878917b82ad5ba09addf0a7ce0bd2 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 14:30:51 +0800 Subject: [PATCH 06/12] 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()) From 5e70ec04765a67fe2905ea0f7936f0b97aa8ec09 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Mon, 12 Oct 2026 17:44:21 +0800 Subject: [PATCH 07/12] fix(embedding): classify every provider status shape and split only input errors The per-item fallback still multiplied service failures into one request per item. Both gaps come from the fallback this PR adds, so they are fixed here. Status extraction - `_http_status()` reads the status from `status_code`, an integer `code` (Google GenAI exposes the HTTP status there), `response.status_code` (httpx-style wrappers), and the structured `status_code` embedded in the message (AgentScope's DashScope embedding turns every non-200 response into a plain `RuntimeError`). - Only integers are read from `code`. A string `code` names an error (`insufficient_quota`, `rate_limit_exceeded`, `InvalidApiKey`); reading it as a status would misclassify quota and rate-limit errors. - `_is_provider_wide_failure`, `_is_transient_provider_failure`, `_is_rate_limited` and `_is_input_dependent_failure` all classify through it. Splitting - The generic fallback splits only when `_is_input_dependent_failure()` matches: a context-length rejection, the one error a single over-long item explains by itself. Anything unrecognized is provider-wide and is never answered with one request per item. - This narrows the isolation this PR adds. An input-dependent error that carries neither a status nor an input-length signal is no longer split, so it costs its batch instead of multiplying the requests. Erring towards the batch keeps one provider failure from becoming N failures and losing the healthy vectors anyway. - `_embed_individually()` drops the redundant `len(single) == 1` check: the `_call_with_retry([text])` contract only returns a result of the expected length. Verified against the real exception classes: `google.genai.errors.APIError.code` is an `int`, and a real DashScope response yields a `RuntimeError` whose message carries `"status_code"`. Regression coverage uses the real classes when available and mirrors their shape otherwise. --- .../embedding_store/local_embedding_store.py | 112 +++++++++- tests/unit/test_local_embedding_store.py | 194 ++++++++++++++++++ 2 files changed, 295 insertions(+), 11 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 75b8dbdf..b8e64ae0 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -2,6 +2,7 @@ import asyncio import hashlib +import re from collections import OrderedDict from pathlib import Path @@ -35,6 +36,41 @@ _TRANSPORT_FAILURE_TYPES = frozenset( }, ) +# The status a provider reports can live in more than one place. ``status_code`` is the +# OpenAI-compatible attribute, Google GenAI exposes an integer ``code``, httpx-style wrappers +# keep it on ``response``, and AgentScope's DashScope implementation raises a plain +# ``RuntimeError`` whose message embeds the structured response, so the status has to be read +# back out of the message. A string ``code`` names an error (``insufficient_quota``, +# ``rate_limit_exceeded``, ``InvalidApiKey``) and is never a status. +_STATUS_IN_MESSAGE = re.compile(r'"status_code"\s*:\s*(\d{3})') + +# A rejected batch can only be paid back one item at a time when the provider blames an input. +# A context-length rejection is the one status an over-long item explains by itself, so it is +# the only error allowed to fan out into per-item requests. +_INPUT_DEPENDENT_STATUS_CODES = frozenset({400}) + +_INPUT_DEPENDENT_ERROR_TYPES = frozenset( + { + "ContextLengthError", + "ContextWindowExceededError", + "InputLengthError", + "TokenLimitError", + }, +) + +_INPUT_DEPENDENT_MARKERS = ( + "context length", + "context_length", + "context window", + "input length", + "input is too long", + "maximum context", + "reduce the length", + "token limit", + "too many tokens", + "too long", +) + @R.register("local") class LocalEmbeddingStore(BaseEmbeddingStore): @@ -256,11 +292,19 @@ class LocalEmbeddingStore(BaseEmbeddingStore): 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: - return await self._embed_individually(texts, **kwargs) - return None + if self._is_input_dependent_failure(error): + # A rejection a single input explains by itself is the one case the per-item + # split can resolve: isolating it keeps the remaining vectors usable. + self.logger.exception("Embedding request failed: input-dependent error") + self.is_healthy = False + if len(texts) > 1: + return await self._embed_individually(texts, **kwargs) + return None + # Anything else counts as provider-wide. Splitting cannot recover a vector unless + # the failing input is identified, so an unrecognized error is never answered + # with one request per item. + self.logger.exception("Embedding request failed: unrecognized error, not splitting") + break self.is_healthy = False return None @@ -277,7 +321,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): results: list[list[float] | None] = [] for text in texts: single = await self._call_with_retry([text], **kwargs) - if single and len(single) == 1: + if single: results.append(single[0]) continue self.logger.error(f"Embedding request failed for a single item (chars={len(text)})") @@ -296,6 +340,26 @@ class LocalEmbeddingStore(BaseEmbeddingStore): return True return any(cls.__name__ in _TRANSPORT_FAILURE_TYPES for cls in type(error).__mro__) + @staticmethod + def _http_status(error: Exception) -> int | None: + """Return the HTTP status an exception carries, across the provider SDK shapes. + + ``status_code`` covers the OpenAI-compatible SDKs, an integer ``code`` covers Google + GenAI, ``response.status_code`` covers httpx-style wrappers, and the structured + ``status_code`` embedded in the message covers AgentScope's DashScope ``RuntimeError``. + Only integers are accepted: a string ``code`` is a provider error name such as + ``insufficient_quota``, and reading it as a status would misclassify those errors. + """ + for attribute in ("status_code", "code"): + value = getattr(error, attribute, None) + if isinstance(value, int) and not isinstance(value, bool): + return value + value = getattr(getattr(error, "response", None), "status_code", None) + if isinstance(value, int) and not isinstance(value, bool): + return value + match = _STATUS_IN_MESSAGE.search(str(error)) + return int(match.group(1)) if match else None + @staticmethod def _is_provider_wide_failure(error: Exception) -> bool: """Return whether an error describes the provider rather than a single input. @@ -307,8 +371,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore): 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 _PROVIDER_WIDE_STATUS_CODES or 500 <= status < 600): + status = LocalEmbeddingStore._http_status(error) + if status is not None and (status in _PROVIDER_WIDE_STATUS_CODES or 500 <= status < 600): return True return type(error).__name__ in ( "AuthenticationError", @@ -317,11 +381,37 @@ class LocalEmbeddingStore(BaseEmbeddingStore): "InternalServerError", ) + @staticmethod + def _is_input_dependent_failure(error: Exception) -> bool: + """Return whether a single input can be blamed for the rejected request. + + Only a context-length rejection qualifies: it is the one error an over-long item + explains by itself, and therefore the only one a per-item retry can resolve. Anything + else stays on the provider-wide path, so a client never answers a service failure with + one request per item. + """ + status = LocalEmbeddingStore._http_status(error) + if status is not None and status not in _INPUT_DEPENDENT_STATUS_CODES: + return False + if type(error).__name__ in _INPUT_DEPENDENT_ERROR_TYPES: + return True + text = LocalEmbeddingStore._error_text(error) + return any(marker in text for marker in _INPUT_DEPENDENT_MARKERS) + + @staticmethod + def _error_text(error: Exception) -> str: + """Flatten an exception and its OpenAI-compatible body into one lowercase string.""" + parts = [str(error)] + body = getattr(error, "body", None) + if isinstance(body, dict): + parts.append(str(body)) + return " ".join(parts).lower() + @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): + status = LocalEmbeddingStore._http_status(error) + if status is not None: return 500 <= status < 600 return type(error).__name__ == "InternalServerError" @@ -330,7 +420,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): """Recognize an OpenAI-compatible 429 response without importing a provider SDK.""" if LocalEmbeddingStore._is_insufficient_quota(error): return False - if getattr(error, "status_code", None) == 429: + if LocalEmbeddingStore._http_status(error) == 429: return True body = getattr(error, "body", None) if not isinstance(body, dict): diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 8d3053e5..982629a4 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -3,6 +3,7 @@ # pylint: disable=protected-access import asyncio +import json from types import SimpleNamespace import numpy as np @@ -1179,3 +1180,196 @@ def test_provider_wide_not_found_errors_are_not_split(monkeypatch): assert not sleeps run(go()) + + +class GenAiClientError(Exception): + """Google GenAI client error: the HTTP status is an integer ``code``, not ``status_code``.""" + + def __init__(self, code: int, message: str = "client error"): + super().__init__(message) + self.code = code + + +class GenAiServerError(Exception): + """Google GenAI server error: the HTTP status is an integer ``code``, not ``status_code``.""" + + def __init__(self, code: int, message: str = "server error"): + super().__init__(message) + self.code = code + + +class StringCodeError(Exception): + """OpenAI-compatible error whose ``code`` names the error instead of carrying a status.""" + + code = "invalid_api_key" + + +class BadRequestError(Exception): + """OpenAI-compatible 400 that names no input, such as a malformed payload.""" + + status_code = 400 + + +class ResponseLike: + """httpx-style wrapper that keeps the status on a ``response`` attribute.""" + + def __init__(self, status_code: int): + self.status_code = status_code + + +class WrappedTransportError(Exception): + """Error that only exposes its status through ``response.status_code``.""" + + def __init__(self, status_code: int): + super().__init__("wrapped") + self.response = ResponseLike(status_code) + + +class DashScopeRuntimeError(RuntimeError): + """AgentScope turns every non-200 DashScope response into a plain ``RuntimeError``. + + The structured response is embedded in the message instead of being kept as an attribute, + so a client that only reads exception attributes cannot classify it. + """ + + @staticmethod + def for_status(status_code: int, code: str = "InvalidApiKey") -> "DashScopeRuntimeError": + """Build the exception AgentScope raises for a non-200 DashScope response.""" + response = { + "status_code": status_code, + "request_id": "req-0", + "code": code, + "message": "request rejected", + "output": None, + "usage": None, + } + return DashScopeRuntimeError(f"Failed to get embedding from DashScope API: {json.dumps(response)}") + + +def test_http_status_reads_every_supported_exception_shape(): + """Read the status from status_code, an integer code, response.status_code or the message.""" + + store = LocalEmbeddingStore(name="t_local_embedding_status_shapes") + + assert store._http_status(AuthenticationError("bad key")) == 401 + assert store._http_status(NotFoundError("model_not_found")) == 404 + assert store._http_status(GenAiClientError(401)) == 401 + assert store._http_status(GenAiServerError(503)) == 503 + assert store._http_status(WrappedTransportError(404)) == 404 + assert store._http_status(DashScopeRuntimeError.for_status(401)) == 401 + assert store._http_status(ValueError("no status here")) is None + + +def test_http_status_ignores_a_string_code(): + """A string ``code`` names the error; reading it as a status would misclassify quota errors.""" + + store = LocalEmbeddingStore(name="t_local_embedding_status_string_code") + + assert store._http_status(StringCodeError("invalid key")) is None + assert store._is_insufficient_quota(StringCodeError("invalid key")) is False + assert store._http_status(InsufficientQuotaError("quota exhausted")) == 429 + assert store._is_insufficient_quota(InsufficientQuotaError("quota exhausted")) is True + + +def test_google_genai_service_failures_are_not_split_into_per_item_requests(monkeypatch): + """A Gemini 401 / 429 / 503 describes the service, so the batch is never split per item.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + for error, expected_batches in ( + (GenAiClientError(401), [10]), + (GenAiServerError(503), [10, 10, 10]), + (GenAiClientError(429), [10, 10, 10]), + ): + sleeps.clear() + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore( + name="t_local_embedding_genai", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == expected_batches + + run(go()) + + +def test_dashscope_runtime_errors_are_not_split_into_per_item_requests(monkeypatch): + """A DashScope failure arrives as a plain RuntimeError; the status lives in its message.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + for status, expected_batches in ((401, [10]), (503, [10, 10, 10]), (429, [10, 10, 10])): + sleeps.clear() + provider = AlwaysProviderWideFailureAsEmbedding(DashScopeRuntimeError.for_status(status)) + store = LocalEmbeddingStore( + name="t_local_embedding_dashscope", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == expected_batches + + run(go()) + + +def test_unrecognized_errors_are_not_split_into_per_item_requests(): + """Only an input-dependent error may split; anything unrecognized stays a batch failure.""" + + async def go(): + provider = AlwaysProviderWideFailureAsEmbedding(RuntimeError("provider said no")) + store = LocalEmbeddingStore(name="t_local_embedding_unrecognized", max_retries=3) + store.as_embedding = provider + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert provider.batch_sizes == [3] + + run(go()) + + +def test_input_dependent_classifier_pins_both_halves_of_the_boundary(): + """The isolation must keep working for the one error a single item can explain.""" + + store = LocalEmbeddingStore(name="t_local_embedding_input_dependent") + + assert ( + store._is_input_dependent_failure( + ContextLengthError("the input length exceeds the context length"), + ) + is True + ) + assert store._is_input_dependent_failure(ContextLengthError("too long")) is True + assert store._is_input_dependent_failure(AuthenticationError("bad key")) is False + assert store._is_input_dependent_failure(InternalServerError("boom")) is False + assert store._is_input_dependent_failure(GenAiClientError(400)) is False + assert store._is_input_dependent_failure(RuntimeError("provider said no")) is False + + +def test_bad_request_without_an_input_length_signal_is_not_split(): + """A bare 400 does not name an input, so the batch is not divided over it.""" + + async def go(): + provider = AlwaysProviderWideFailureAsEmbedding(BadRequestError("malformed payload")) + store = LocalEmbeddingStore(name="t_local_embedding_bare_400", max_retries=2) + store.as_embedding = provider + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert provider.batch_sizes == [3] + + run(go()) From f95ffc30ef5a04a1ef6719ffea3a1281aac358fb Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Tue, 13 Oct 2026 12:24:13 +0800 Subject: [PATCH 08/12] fix(embedding): stop per-item isolation on a provider-wide failure and read serialized codes - `_embed_individually` now stops asking for the remaining items as soon as one isolated item fails provider-wide (credentials, missing model, exhausted quota, exhausted transient retries). No remaining input can explain that failure, so continuing only multiplies the rejected requests. The vectors already recovered are kept and the items that were never answered stay `None`. - `_call_with_retry` reports that outcome through an optional `terminal_failure` out-parameter, so its return type and every existing call site stay unchanged. - `_error_code` reads the business error code from the attribute, the structured body, or the serialized message. AgentScope's DashScope embedding raises a plain `RuntimeError` with the response embedded in the message, so a quota error was classified as a rate limit and never used the configured `quota_retry_delay`. Regression tests cover both halves of the isolation boundary (a provider-wide failure stops the loop; an input-dependent failure does not) and the serialized DashScope quota path with quota retries enabled. --- .../embedding_store/local_embedding_store.py | 90 +++++++++++--- tests/unit/test_local_embedding_store.py | 115 ++++++++++++++++++ 2 files changed, 188 insertions(+), 17 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index b8e64ae0..fb653a1d 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -44,6 +44,12 @@ _TRANSPORT_FAILURE_TYPES = frozenset( # ``rate_limit_exceeded``, ``InvalidApiKey``) and is never a status. _STATUS_IN_MESSAGE = re.compile(r'"status_code"\s*:\s*(\d{3})') +# The business error code travels the same way the status does: a string ``code`` on the +# OpenAI-compatible SDKs, ``error.code`` inside a structured body, and a serialized +# ``"code"`` in the message for AgentScope's DashScope ``RuntimeError``. Reading only the +# attribute loses the DashScope case and misclassifies a quota error as a rate limit. +_CODE_IN_MESSAGE = re.compile(r'"code"\s*:\s*"([^"]+)"') + # A rejected batch can only be paid back one item at a time when the provider blames an input. # A context-length rejection is the one status an over-long item explains by itself, so it is # the only error allowed to fan out into per-item requests. @@ -241,7 +247,27 @@ class LocalEmbeddingStore(BaseEmbeddingStore): self.is_healthy = False return out - async def _call_with_retry(self, texts: list[str], **kwargs) -> list[list[float] | None] | None: + async def _call_with_retry( + self, + texts: list[str], + *, + terminal_failure: list[bool] | None = None, + **kwargs, + ) -> list[list[float] | None] | None: + """Answer a batch, retrying transient failures and isolating input ones. + + ``terminal_failure`` is an optional out-parameter. When the caller passes a list, + it is set to ``[True]`` if this call gave up on a failure that no remaining item + can explain (a provider-wide error, exhausted quota, or exhausted transient + retries). The per-item isolation loop uses it to stop instead of asking for every + remaining item and multiplying the rejected requests. + """ + + def stop() -> None: + """Record that this failure ends the work for the whole batch.""" + if terminal_failure is not None: + terminal_failure[:] = [True] + for attempt in range(self.max_retries): try: result = await self.as_embedding(texts, **kwargs) @@ -263,6 +289,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): continue self.logger.exception("Embedding request failed after exhausting rate-limit retries") self.is_healthy = False + stop() return None if ( self.quota_retry_delay is not None @@ -280,6 +307,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): # This also covers the case where the quota retries above are exhausted. self.logger.exception("Embedding request failed: quota exhausted") self.is_healthy = False + stop() return None if self._is_provider_wide_failure(error): # Credentials, permissions and server outages describe the endpoint, not any @@ -291,6 +319,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): await asyncio.sleep(delay) continue self.logger.exception("Embedding request failed: provider-wide error, not splitting") + stop() break if self._is_input_dependent_failure(error): # A rejection a single input explains by itself is the one case the per-item @@ -304,8 +333,10 @@ class LocalEmbeddingStore(BaseEmbeddingStore): # the failing input is identified, so an unrecognized error is never answered # with one request per item. self.logger.exception("Embedding request failed: unrecognized error, not splitting") + stop() break self.is_healthy = False + stop() return None async def _embed_individually(self, texts: list[str], **kwargs) -> list[list[float] | None]: @@ -319,13 +350,24 @@ class LocalEmbeddingStore(BaseEmbeddingStore): policy instead of losing the vector; a one-item request never splits again. """ results: list[list[float] | None] = [] - for text in texts: - single = await self._call_with_retry([text], **kwargs) + for index, text in enumerate(texts): + terminal: list[bool] = [] + single = await self._call_with_retry([text], terminal_failure=terminal, **kwargs) if single: results.append(single[0]) continue self.logger.error(f"Embedding request failed for a single item (chars={len(text)})") results.append(None) + if terminal and terminal[0]: + # No remaining item can explain a provider-wide failure, so asking for the + # rest would only multiply the rejected requests. Keep the vectors already + # recovered and leave the items that were never answered as None. + self.logger.error( + f"Stopping per-item isolation after {index + 1}/{len(texts)} items: " + "the failure is not specific to any input", + ) + results.extend([None] * (len(texts) - index - 1)) + break return results @staticmethod @@ -415,29 +457,43 @@ class LocalEmbeddingStore(BaseEmbeddingStore): return 500 <= status < 600 return type(error).__name__ == "InternalServerError" + @staticmethod + def _error_code(error: Exception) -> str | None: + """Return the provider's business error code, across the supported error shapes. + + ``code`` covers the OpenAI-compatible SDKs when the provider puts the error name + there, ``error.code`` inside ``body`` covers the structured OpenAI payload, and the + serialized message covers AgentScope's DashScope ``RuntimeError`` — the same shape + ``_http_status`` already reads the status back out of. Without the last case a + DashScope quota error is indistinguishable from an ordinary 429. + """ + code = getattr(error, "code", None) + if isinstance(code, str) and code: + return code + body = getattr(error, "body", None) + if isinstance(body, dict): + details = body.get("error", body) + if isinstance(details, dict): + inner = details.get("code") + if isinstance(inner, str) and inner: + return inner + match = _CODE_IN_MESSAGE.search(str(error)) + return match.group(1) if match else None + @staticmethod def _is_rate_limited(error: Exception) -> bool: """Recognize an OpenAI-compatible 429 response without importing a provider SDK.""" - if LocalEmbeddingStore._is_insufficient_quota(error): + code = LocalEmbeddingStore._error_code(error) + if code == "insufficient_quota": return False - if LocalEmbeddingStore._http_status(error) == 429: + if code == "rate_limit_exceeded": return True - body = getattr(error, "body", None) - if not isinstance(body, dict): - return False - details = body.get("error", body) - return isinstance(details, dict) and details.get("code") == "rate_limit_exceeded" + return LocalEmbeddingStore._http_status(error) == 429 @staticmethod def _is_insufficient_quota(error: Exception) -> bool: """Recognize OpenAI-compatible quota errors without importing a provider SDK.""" - if getattr(error, "code", None) == "insufficient_quota": - return True - body = getattr(error, "body", None) - if not isinstance(body, dict): - return False - details = body.get("error", body) - return isinstance(details, dict) and details.get("code") == "insufficient_quota" + return LocalEmbeddingStore._error_code(error) == "insufficient_quota" def _validate_dim(self, emb: np.ndarray) -> bool: """Return whether an embedding exactly matches the configured dimension.""" diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 982629a4..2ffaaf77 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -201,6 +201,35 @@ class PerItemLimitThenTransientErrorAsEmbedding: return [[1.0, 0.0] for _ in texts] +class PerItemLimitThenQuotaAsEmbedding: + """Reject an over-limit batch, recover the first isolated item, then exhaust the quota. + + Mirrors the review case: isolation starts, the first single-item request comes back, and + every later call then fails provider-wide. Asking for the remaining items would only + multiply the rejected requests, because no single item explains a quota failure. + """ + + dimensions = 2 + vector_space_id = "fakespace000" + max_chars = 32 + + def __init__(self): + self.batch_sizes: list[int] = [] + self.singles = 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") + self.singles += 1 + if self.singles > 1: + raise InsufficientQuotaError("quota exhausted") + return [[1.0, 0.0] for _ in texts] + + class AuthenticationError(Exception): """OpenAI-compatible 401 error used without importing the provider SDK.""" @@ -1373,3 +1402,89 @@ def test_bad_request_without_an_input_length_signal_is_not_split(): assert provider.batch_sizes == [3] run(go()) + + +def test_isolation_stops_when_a_provider_wide_failure_hits(): + """Once isolation meets a provider-wide failure, the remaining items are not requested. + + The vectors already recovered stay in the result: only the items that were never answered + are dropped, and no further request is sent for them. + """ + + async def go(): + provider = PerItemLimitThenQuotaAsEmbedding() + store = LocalEmbeddingStore( + name="t_local_embedding_isolation_stops", + max_retries=2, + enable_cache=False, + ) + store.as_embedding = provider + + results = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"]) + + assert results == [[1.0, 0.0], None, None, None] + # One rejected batch, the recovered first item, the over-limit item, then the quota + # failure. The fourth item is never requested once the quota failure is seen. + assert provider.batch_sizes == [4, 1, 1, 1] + + run(go()) + + +def test_isolation_keeps_going_while_only_inputs_fail(): + """An input-dependent failure is local to its item, so isolation must continue. + + This pins the other half of the boundary: stopping on every isolated failure would lose + the vectors of the items that are perfectly embeddable on their own. + """ + + async def go(): + provider = PerItemLimitAsEmbedding() + store = LocalEmbeddingStore( + name="t_local_embedding_isolation_keeps_going", + max_retries=2, + enable_cache=False, + ) + store.as_embedding = provider + + results = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0], [1.0, 0.0]] + assert provider.batch_sizes == [4, 1, 1, 1, 1] + + run(go()) + + +def test_dashscope_serialized_quota_errors_keep_their_quota_policy(monkeypatch): + """A serialized DashScope quota error must classify as quota, not as a rate limit. + + AgentScope embeds the structured response in the message, so the business error code has to + be read back out of it. Without that the opt-in quota retry policy is never applied and the + error is retried with the generic backoff instead. + """ + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + error = DashScopeRuntimeError.for_status(429, code="insufficient_quota") + assert LocalEmbeddingStore._is_insufficient_quota(error) is True + assert LocalEmbeddingStore._is_rate_limited(error) is False + + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore( + name="t_local_embedding_dashscope_quota", + max_retries=3, + quota_retry_delay=7.0, + enable_cache=False, + ) + store.as_embedding = provider + + assert await store._call_with_retry(["a", "b"]) is None + assert provider.batch_sizes == [2, 2, 2] + assert sleeps == [7.0, 7.0] + + run(go()) From eca11c9e04e6b9af33c236a367661aa69326c70d Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Tue, 13 Oct 2026 14:03:32 +0800 Subject: [PATCH 09/12] fix(embedding): treat an input-length 413 as a per-item rejection --- .../embedding_store/local_embedding_store.py | 22 +++- tests/unit/test_local_embedding_store.py | 100 ++++++++++++++++++ 2 files changed, 118 insertions(+), 4 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index fb653a1d..40d002fc 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -55,6 +55,13 @@ _CODE_IN_MESSAGE = re.compile(r'"code"\s*:\s*"([^"]+)"') # the only error allowed to fan out into per-item requests. _INPUT_DEPENDENT_STATUS_CODES = frozenset({400}) +# A 413 is ambiguous. Providers such as Ollama v0.40.3 normalize an embedding input-limit +# rejection to 413, but a bare 413 can also describe an unrelated request-body limit that no +# single input explains. It therefore only qualifies as input-dependent when the error itself +# names the input length, which is why it is checked through the marker list instead of being +# accepted outright. +_INPUT_DEPENDENT_AMBIGUOUS_STATUS_CODES = frozenset({413}) + _INPUT_DEPENDENT_ERROR_TYPES = frozenset( { "ContextLengthError", @@ -428,13 +435,20 @@ class LocalEmbeddingStore(BaseEmbeddingStore): """Return whether a single input can be blamed for the rejected request. Only a context-length rejection qualifies: it is the one error an over-long item - explains by itself, and therefore the only one a per-item retry can resolve. Anything - else stays on the provider-wide path, so a client never answers a service failure with - one request per item. + explains by itself, and therefore the only one a per-item retry can resolve. It can + arrive as a ``400`` or, from a provider that normalizes the rejection (Ollama reports + ``413``), as an ambiguous status that is accepted only when the error names the input + itself. Anything else stays on the provider-wide path, so a client never answers a + service failure with one request per item. """ status = LocalEmbeddingStore._http_status(error) if status is not None and status not in _INPUT_DEPENDENT_STATUS_CODES: - return False + if status not in _INPUT_DEPENDENT_AMBIGUOUS_STATUS_CODES: + return False + # The status alone does not name an input for an ambiguous code, so the error text + # has to carry the signal; otherwise this is a request-level rejection. + text = LocalEmbeddingStore._error_text(error) + return any(marker in text for marker in _INPUT_DEPENDENT_MARKERS) if type(error).__name__ in _INPUT_DEPENDENT_ERROR_TYPES: return True text = LocalEmbeddingStore._error_text(error) diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 2ffaaf77..b8e5ac6d 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -1488,3 +1488,103 @@ def test_dashscope_serialized_quota_errors_keep_their_quota_policy(monkeypatch): assert sleeps == [7.0, 7.0] run(go()) + + +class ContextLength413Error(Exception): + """An over-limit input reported as HTTP 413 instead of 400. + + Ollama v0.40.3 normalizes an embedding input-limit rejection to 413 and propagates that + status through the embedding endpoint, so the same over-long item that arrives as a 400 + from an OpenAI-compatible gateway is a 413 for that provider. + """ + + status_code = 413 + body = {"error": {"message": "the input length exceeds the context length"}} + + +class RequestBodyTooLargeError(Exception): + """A bare 413 that never names an input: the request body itself is too large.""" + + status_code = 413 + + +class PerItemLimitAs413Embedding: + """Reject a whole request with HTTP 413 when any single text is over the per-item limit.""" + + 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 ContextLength413Error("the input length exceeds the context length") + return [[1.0, 0.0] for _ in texts] + + +def test_context_length_rejection_reported_as_http_413_is_isolated(): + """An over-limit input carried by a 413 must still reach the per-item isolation path. + + Ollama normalizes an embedding input-limit rejection to 413, so a classifier that only + accepts 400 rejects the batch instead of isolating the item that caused it, and every + healthy vector in the request is lost. + """ + + async def go(): + provider = PerItemLimitAs413Embedding() + store = LocalEmbeddingStore( + name="t_local_embedding_413_isolated", + max_retries=2, + 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_bare_http_413_without_an_input_signal_is_not_split(): + """A 413 that names no input stays provider-wide, so the batch is not divided over it. + + The status alone is ambiguous: it can describe an unrelated request-body limit, which no + single input explains, so only a 413 that blames the input itself may fan out per item. + """ + + async def go(): + provider = AlwaysProviderWideFailureAsEmbedding( + RequestBodyTooLargeError("request entity too large"), + ) + store = LocalEmbeddingStore(name="t_local_embedding_bare_413", max_retries=2) + store.as_embedding = provider + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert provider.batch_sizes == [3] + + run(go()) + + +def test_input_dependent_classifier_accepts_a_413_only_with_an_input_signal(): + """Pin both halves of the 413 boundary, and that the 400 behaviour is unchanged.""" + + store = LocalEmbeddingStore(name="t_local_embedding_413_classifier") + + assert ( + store._is_input_dependent_failure( + ContextLength413Error("the input length exceeds the context length"), + ) + is True + ) + assert store._is_input_dependent_failure(ContextLength413Error("too long")) is True + assert store._is_input_dependent_failure(RequestBodyTooLargeError("request entity too large")) is False + assert store._is_input_dependent_failure(ContextLengthError("too long")) is True + assert store._is_input_dependent_failure(GenAiClientError(400)) is False From db1665db67e830fed6710df85db7828a8c9e3fec Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Wed, 14 Oct 2026 21:13:20 +0800 Subject: [PATCH 10/12] fix(embedding): keep the health state after a partial success `_compute_batch()` set `is_healthy = True` whenever the batch produced at least one vector. Per-item isolation can recover an early item and then meet a provider-wide failure (exhausted quota, authentication, exhausted transient retries), which `_call_with_retry()` correctly records as unhealthy before it stops. The recovered vector then overwrote that state, so a request that ended in an outage reported the provider as healthy. The terminal-failure flag is now threaded from `_embed_individually()` back to `_compute_batch()`. The vectors recovered before the failure are still returned and cached; only the health state now follows the failure. This matters because `LocalFileStore._recover_after_real_request()` schedules missing-vector backfill off the unhealthy-to-healthy transition: a hidden outage would either trigger backfill while the provider is still unavailable or suppress a later genuine recovery. Adds a regression test through `get_embeddings()` that checks both the retained partial vector and the final unhealthy state. The existing isolation test covers `_call_with_retry()` directly, so it never saw the overwrite. --- .../embedding_store/local_embedding_store.py | 40 ++++++++++++++++--- tests/unit/test_local_embedding_store.py | 34 ++++++++++++++++ 2 files changed, 68 insertions(+), 6 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 40d002fc..142d9365 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -232,7 +232,17 @@ class LocalEmbeddingStore(BaseEmbeddingStore): async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]: texts = [text for _, text, _ in batch] - embeddings = await self._call_with_retry(texts, **kwargs) + # A provider-wide failure can end a request after some of its vectors were already + # recovered. Those vectors are still worth keeping, but they must not report the + # provider as healthy again: the unhealthy-to-healthy transition is what schedules + # missing-vector backfill, so a hidden outage would either trigger it while the + # provider is still unavailable or suppress a later genuine recovery. + terminal_failure: list[bool] = [] + embeddings = await self._call_with_retry( + texts, + terminal_failure=terminal_failure, + **kwargs, + ) if not embeddings or len(embeddings) != len(texts): return [] out: list[tuple[int, str, np.ndarray]] = [] @@ -248,7 +258,7 @@ 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: + if out and not terminal_failure: self.is_healthy = True else: self.is_healthy = False @@ -267,7 +277,9 @@ class LocalEmbeddingStore(BaseEmbeddingStore): it is set to ``[True]`` if this call gave up on a failure that no remaining item can explain (a provider-wide error, exhausted quota, or exhausted transient retries). The per-item isolation loop uses it to stop instead of asking for every - remaining item and multiplying the rejected requests. + remaining item and multiplying the rejected requests, and it forwards the same + flag to its caller so a request that recovered some vectors before the failure is + still not reported as healthy. """ def stop() -> None: @@ -334,7 +346,11 @@ class LocalEmbeddingStore(BaseEmbeddingStore): self.logger.exception("Embedding request failed: input-dependent error") self.is_healthy = False if len(texts) > 1: - return await self._embed_individually(texts, **kwargs) + return await self._embed_individually( + texts, + terminal_failure=terminal_failure, + **kwargs, + ) return None # Anything else counts as provider-wide. Splitting cannot recover a vector unless # the failing input is identified, so an unrecognized error is never answered @@ -346,7 +362,13 @@ class LocalEmbeddingStore(BaseEmbeddingStore): stop() return None - async def _embed_individually(self, texts: list[str], **kwargs) -> list[list[float] | None]: + async def _embed_individually( + self, + texts: list[str], + *, + terminal_failure: list[bool] | None = None, + **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 @@ -355,6 +377,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore): 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. + When the loop stops on a provider-wide failure, ``terminal_failure`` is set for the + caller so the vectors recovered before it are not read as a healthy provider. """ results: list[list[float] | None] = [] for index, text in enumerate(texts): @@ -368,7 +392,11 @@ class LocalEmbeddingStore(BaseEmbeddingStore): if terminal and terminal[0]: # No remaining item can explain a provider-wide failure, so asking for the # rest would only multiply the rejected requests. Keep the vectors already - # recovered and leave the items that were never answered as None. + # recovered and leave the items that were never answered as None. The batch + # caller has to learn about it as well: a partial success must not report + # the provider as healthy again. + if terminal_failure is not None: + terminal_failure[:] = [True] self.logger.error( f"Stopping per-item isolation after {index + 1}/{len(texts)} items: " "the failure is not specific to any input", diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index b8e5ac6d..e93f81e0 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -1430,6 +1430,40 @@ def test_isolation_stops_when_a_provider_wide_failure_hits(): run(go()) +def test_partial_success_keeps_a_failed_provider_unhealthy(): + """Recovered vectors must not overwrite the health state of a provider-wide failure. + + Isolation can recover an early item and then meet a provider-wide failure such as + exhausted quota. ``_call_with_retry`` marks the store unhealthy and stops, but the + vectors recovered before it are still returned. Treating that partial success as a + healthy provider would hide the outage, and it would break the unhealthy-to-healthy + transition the file store uses to schedule missing-vector backfill. The existing + isolation test drives ``_call_with_retry`` directly, so it never sees the overwrite. + """ + + async def go(): + provider = PerItemLimitThenQuotaAsEmbedding() + store = LocalEmbeddingStore( + name="t_local_embedding_partial_health", + max_retries=2, + enable_cache=False, + ) + store.as_embedding = provider + + vectors = await store.get_embeddings(["ok", "x" * 64, "also ok", "still ok"]) + + # The vector recovered before the outage is kept ... + assert vectors[0] is not None + assert vectors[1] is None and vectors[2] is None and vectors[3] is None + # ... and the provider-wide failure still decides the health state. + assert store.is_healthy is False + # One rejected batch, the recovered first item, the over-limit item, then the quota + # failure: the fourth item is never requested. + assert provider.batch_sizes == [4, 1, 1, 1] + + run(go()) + + def test_isolation_keeps_going_while_only_inputs_fail(): """An input-dependent failure is local to its item, so isolation must continue. From 6dcc7664a332ced4cbfc35bd3b85df518b1a0163 Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 16 Oct 2026 00:43:07 +0800 Subject: [PATCH 11/12] fix(embedding): leave provider 5xx retries to the layers that own them A store-level retry for a provider 5xx stacked on top of the retries the embedding wrapper and the provider SDK already perform, so a batch the SDK answered after 12 attempts became 36. Retry ownership stays with those layers; the store answers a provider-wide 5xx once and gives up. Transport failures keep their retry because they describe a connection that no lower layer has reported, and the opt-in quota policy is untouched. Tests updated to pin the new boundary: a 503 is answered once for the OpenAI-shaped, Gemini-shaped and DashScope-shaped errors, while a 429 keeps its retry from the rate-limit branch. --- .../embedding_store/local_embedding_store.py | 26 ++++++++---------- tests/unit/test_local_embedding_store.py | 27 +++++++++++++------ 2 files changed, 30 insertions(+), 23 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 142d9365..d5a9d03d 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -330,13 +330,17 @@ class LocalEmbeddingStore(BaseEmbeddingStore): 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 + # single item: isolation cannot fix them, so give up on the whole batch + # instead of multiplying the rejected requests. + # + # A 5xx is deliberately *not* retried here. The embedding wrapper and the + # provider SDK already retry transient server errors, and every extra layer + # that retries multiplies the attempts and the waiting time of the layers + # below it -- a batch the SDK answered after 12 attempts became 36 once this + # store retried on top of it. Retry ownership therefore stays with those + # layers: a caller that wants a longer budget configures it there instead of + # stacking another silent loop here. Transport failures keep their retry + # above because they describe a connection that no lower layer has reported. self.logger.exception("Embedding request failed: provider-wide error, not splitting") stop() break @@ -491,14 +495,6 @@ class LocalEmbeddingStore(BaseEmbeddingStore): parts.append(str(body)) return " ".join(parts).lower() - @staticmethod - def _is_transient_provider_failure(error: Exception) -> bool: - """Return whether a provider-wide failure is worth retrying with the usual backoff.""" - status = LocalEmbeddingStore._http_status(error) - if status is not None: - return 500 <= status < 600 - return type(error).__name__ == "InternalServerError" - @staticmethod def _error_code(error: Exception) -> str | None: """Return the provider's business error code, across the supported error shapes. diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index e93f81e0..87d71a37 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -1075,8 +1075,8 @@ def test_provider_wide_credentials_errors_are_not_split_into_per_item_requests() 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.""" +def test_provider_wide_server_errors_are_not_retried_without_splitting(monkeypatch): + """A 500 is left to the layers that own retries, so this store asks exactly once.""" async def go(): sleeps = [] @@ -1090,8 +1090,10 @@ def test_provider_wide_server_errors_are_retried_without_splitting(monkeypatch): 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] + # One request only: the wrapper and the SDK already retry a 5xx, and retrying here + # as well multiplied their attempts (12 -> 36 in the reported reproduction). + assert provider.batch_sizes == [10] + assert sleeps == [] assert store.is_healthy is False run(go()) @@ -1301,7 +1303,12 @@ def test_http_status_ignores_a_string_code(): def test_google_genai_service_failures_are_not_split_into_per_item_requests(monkeypatch): - """A Gemini 401 / 429 / 503 describes the service, so the batch is never split per item.""" + """A Gemini 401 / 429 / 503 describes the service, so the batch is never split per item. + + A 503 is answered once: the wrapper and the SDK own the retry for a 5xx, and retrying + here as well multiplied their attempts. A 429 keeps its retry because the rate-limit + branch above this classifier is the one place that reacts to it. + """ async def go(): sleeps = [] @@ -1313,7 +1320,7 @@ def test_google_genai_service_failures_are_not_split_into_per_item_requests(monk for error, expected_batches in ( (GenAiClientError(401), [10]), - (GenAiServerError(503), [10, 10, 10]), + (GenAiServerError(503), [10]), (GenAiClientError(429), [10, 10, 10]), ): sleeps.clear() @@ -1332,7 +1339,11 @@ def test_google_genai_service_failures_are_not_split_into_per_item_requests(monk def test_dashscope_runtime_errors_are_not_split_into_per_item_requests(monkeypatch): - """A DashScope failure arrives as a plain RuntimeError; the status lives in its message.""" + """A DashScope failure arrives as a plain RuntimeError; the status lives in its message. + + As with the Gemini shapes, a 503 is answered once (the lower layers own that retry) and + a 429 keeps the retry from the rate-limit branch. + """ async def go(): sleeps = [] @@ -1342,7 +1353,7 @@ def test_dashscope_runtime_errors_are_not_split_into_per_item_requests(monkeypat monkeypatch.setattr(asyncio, "sleep", fake_sleep) - for status, expected_batches in ((401, [10]), (503, [10, 10, 10]), (429, [10, 10, 10])): + for status, expected_batches in ((401, [10]), (503, [10]), (429, [10, 10, 10])): sleeps.clear() provider = AlwaysProviderWideFailureAsEmbedding(DashScopeRuntimeError.for_status(status)) store = LocalEmbeddingStore( From 13cff385dbc6ffa7a072b7810e65c1fca76611df Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 16 Oct 2026 21:29:08 +0800 Subject: [PATCH 12/12] fix(embedding): make transport retry ownership explicit The store retried a connection failure or a timeout on top of the embedding wrapper and the provider SDK, which already retry both, so a persistent timeout on a batch reached the provider 36 times instead of 12: the OpenAI model retries a batch max_retries + 1 times and the SDK retries each of those attempts. Transport failures stay provider-wide -- the request is never split per item -- but the retry is now owned explicitly instead of implicitly: * a whole-batch request leaves it to the layers below, which already retry it; * an isolated single-item attempt keeps the store's backoff, because that attempt is the last chance to keep the one vector the isolation step is recovering. Rate limiting is unchanged. The transport cases now pin both budgets, and three tests drive a real OpenAIAsEmbedding wrapper and a real openai.AsyncClient over an httpx.MockTransport, so the provider attempt count is measured rather than assumed. --- .../embedding_store/local_embedding_store.py | 54 +++- tests/unit/test_local_embedding_store.py | 296 +++++++++++++++--- 2 files changed, 284 insertions(+), 66 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index d5a9d03d..7a074f5d 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -268,11 +268,17 @@ class LocalEmbeddingStore(BaseEmbeddingStore): self, texts: list[str], *, + isolated: bool = False, terminal_failure: list[bool] | None = None, **kwargs, ) -> list[list[float] | None] | None: """Answer a batch, retrying transient failures and isolating input ones. + ``isolated`` marks a single-item attempt made while isolating a rejected batch. Such an + attempt keeps the store's transport backoff, because it is the last chance to keep that + one vector; every other request leaves transport retries to the embedding wrapper and the + provider SDK, which already own them. + ``terminal_failure`` is an optional out-parameter. When the caller passes a list, it is set to ``[True]`` if this call gave up on a failure that no remaining item can explain (a provider-wide error, exhausted quota, or exhausted transient @@ -293,13 +299,22 @@ class LocalEmbeddingStore(BaseEmbeddingStore): if result and len(result) == len(texts): return result except Exception as error: - if self._is_transient_transport_failure(error): + if self._is_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 + # item, so this request is never split per item. Which layer retries it is + # explicit: an isolated single-item attempt keeps the store's backoff, while + # any other request is left to the layers below, which already own that retry + # -- the OpenAI model retries a batch max_retries + 1 times and the SDK retries + # each of those attempts, so a batch they answer after 12 attempts used to + # reach the provider 36 times with this store retrying on top of them. + if isolated: + if attempt < self.max_retries - 1: + await asyncio.sleep(2**attempt) + continue + self.logger.exception("Embedding request failed: transport error, not splitting") + self.is_healthy = False + stop() + break if self._is_rate_limited(error): if attempt < self.max_retries - 1: delay = 2**attempt @@ -339,8 +354,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore): # below it -- a batch the SDK answered after 12 attempts became 36 once this # store retried on top of it. Retry ownership therefore stays with those # layers: a caller that wants a longer budget configures it there instead of - # stacking another silent loop here. Transport failures keep their retry - # above because they describe a connection that no lower layer has reported. + # stacking another silent loop here. A transport failure is left to them for + # the same reason; it only skips the per-item split. self.logger.exception("Embedding request failed: provider-wide error, not splitting") stop() break @@ -378,16 +393,22 @@ 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. + Each single-item request goes back through ``_call_with_retry`` with ``isolated=True`` so + a rate limit and a transient transport failure still get the store's backoff instead of + losing the vector; a provider-wide failure ends the work for the whole batch rather than + being asked for every remaining item, and a one-item request never splits again. When the loop stops on a provider-wide failure, ``terminal_failure`` is set for the caller so the vectors recovered before it are not read as a healthy provider. """ results: list[list[float] | None] = [] for index, text in enumerate(texts): terminal: list[bool] = [] - single = await self._call_with_retry([text], terminal_failure=terminal, **kwargs) + single = await self._call_with_retry( + [text], + isolated=True, + terminal_failure=terminal, + **kwargs, + ) if single: results.append(single[0]) continue @@ -410,12 +431,15 @@ class LocalEmbeddingStore(BaseEmbeddingStore): return results @staticmethod - def _is_transient_transport_failure(error: Exception) -> bool: - """Return whether an error is a transient transport failure worth retrying. + def _is_transport_failure(error: Exception) -> bool: + """Return whether an error describes a transport failure rather than a rejected input. 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. + input can be blamed, so the request is never split per item. Who retries them is explicit + in the transport branch of ``_call_with_retry``: an isolated single-item attempt keeps the + store's backoff because it is the last chance to keep that vector, while every other + request leaves the retry to the layers below, which already own it. """ if isinstance(error, (TimeoutError, ConnectionError, OSError)): return True diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 87d71a37..1d9ed6e7 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -6,7 +6,9 @@ import asyncio import json from types import SimpleNamespace +import httpx import numpy as np +import openai from reme.components.as_embedding import DashScopeAsEmbedding, OllamaAsEmbedding, OpenAIAsEmbedding from reme.components.embedding_store.base_embedding_store import BaseEmbeddingStore @@ -301,6 +303,108 @@ class TransientThenSuccessAsEmbedding: return [[1.0, 0.0] for _ in texts] +_REAL_OPENAI_ASYNC_CLIENT = openai.AsyncClient +_OFFLINE_OPENAI_BASE_URL = "http://127.0.0.1:9/v1" + + +class CountingOpenAIWrapper(OpenAIAsEmbedding): + """Real OpenAI wrapper that records how often the store asked it for embeddings.""" + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.calls = 0 + + async def __call__(self, inputs, **kwargs): + self.calls += 1 + return await super().__call__(inputs, **kwargs) + + +class OfflineOpenAIEndpoint: + """Real openai.AsyncClient wired to an offline httpx transport that counts provider attempts. + + Gives the wrapper-and-SDK layers their real code path, so a regression test can assert how + many attempts the provider would see instead of trusting the store's own loop. + """ + + def __init__(self, handler, sdk_max_retries: int = 2) -> None: + self.requests = 0 + self.handler = handler + self.client = _REAL_OPENAI_ASYNC_CLIENT( + api_key="test-key", + base_url=_OFFLINE_OPENAI_BASE_URL, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(self._handle)), + max_retries=sdk_max_retries, + ) + + def _handle(self, request): + self.requests += 1 + return self.handler(request) + + def install(self, monkeypatch) -> "OfflineOpenAIEndpoint": + """Route every openai.AsyncClient built during this test to the offline client.""" + monkeypatch.setattr(openai, "AsyncClient", lambda **kwargs: self.client) + return self + + +def offline_openai_endpoint(handler) -> OfflineOpenAIEndpoint: + """Build an offline endpoint whose transport follows ``handler``.""" + return OfflineOpenAIEndpoint(handler) + + +def openai_wrapper(name: str) -> CountingOpenAIWrapper: + """Build a real OpenAI embedding wrapper pointed at the offline endpoint.""" + return CountingOpenAIWrapper( + name=name, + model="text-embedding-v4", + backend="openai", + dimensions=2, + credential={"base_url": _OFFLINE_OPENAI_BASE_URL, "api_key": "test-key"}, + retry_delay=0, + pass_dimensions=False, + ) + + +def embedding_payload(request) -> httpx.Response: + """Answer an embeddings request the way the provider would.""" + body = json.loads(request.content.decode("utf-8")) + inputs = body["input"] + if isinstance(inputs, str): + inputs = [inputs] + return httpx.Response( + 200, + json={ + "object": "list", + "data": [ + {"object": "embedding", "index": position, "embedding": [0.0, 1.0]} for position in range(len(inputs)) + ], + "model": body.get("model", "text-embedding-v4"), + "usage": {"prompt_tokens": len(inputs), "total_tokens": len(inputs)}, + }, + ) + + +def context_length_rejection() -> httpx.Response: + """Answer a batch the way an OpenAI-compatible provider reports an over-long input.""" + return httpx.Response( + 400, + json={ + "error": { + "message": ( + "This model's maximum context length is 8191 tokens. " "Please reduce the length of the input." + ), + "type": "invalid_request_error", + "param": None, + "code": "context_length_exceeded", + }, + }, + ) + + +def read_timeout(request) -> httpx.Response: + """Fail one attempt the way a read timeout reaches the SDK.""" + raise httpx.ReadTimeout("simulated read timeout", request=request) + + class BadNodeEmbeddingStore(BaseEmbeddingStore): """Embedding store that returns wrong-dimensional vectors.""" @@ -1020,38 +1124,37 @@ def test_exhausted_quota_retries_still_do_not_split_the_batch(monkeypatch): 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.""" +def test_per_item_fallback_keeps_a_rate_limited_vector(monkeypatch): + """A 429 on one isolated item still gets the store's backoff and keeps its vector. + + Rate limiting is the failure class the store still owns: no lower layer reports it as a + provider-wide error worth its own budget, so the retry lives here. Connection errors and + timeouts are left to the wrapper and the SDK instead; see the transport tests above. + """ async def go(): - transients = ( - RateLimitError("Requests are too frequent"), - ConnectionError("connection reset"), - TimeoutError("request timed out"), - ) sleeps = [] async def fake_sleep(delay): sleeps.append(delay) - for transient in transients: - sleeps.clear() + provider = PerItemLimitThenTransientErrorAsEmbedding( + RateLimitError("Requests are too frequent"), + ) + store = LocalEmbeddingStore( + name="t_local_embedding_fallback_rate_limit", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) - 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"]) - 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] + 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()) @@ -1099,8 +1202,12 @@ def test_provider_wide_server_errors_are_not_retried_without_splitting(monkeypat 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.""" +def test_sdk_transport_errors_are_provider_wide_and_not_retried_here(monkeypatch): + """SDK connection/timeout errors are provider-wide: never split, and not retried here. + + The wrapper and the SDK already retry a transport failure, so the store answers once and + leaves the retry budget to those layers instead of multiplying their attempts. + """ async def go(): sleeps = [] @@ -1117,14 +1224,23 @@ def test_sdk_transport_errors_are_retried_without_splitting(monkeypatch): 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] + # One request only: the wrapper and the SDK already retry a transport failure + # (4 wrapper attempts, each retried by the SDK), so another loop here only + # multiplies them -- a persistent timeout used to cost 36 attempts instead of 12. + assert provider.batch_sizes == [10] + assert sleeps == [] + assert store.is_healthy is False 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.""" +def test_isolated_item_keeps_the_store_backoff_for_a_transport_error(monkeypatch): + """An isolated single-item attempt still gets the store's backoff and keeps its vector. + + This fake provider has no wrapper or SDK beneath it, so it pins the store's own budget: the + isolated attempt is the last chance to keep that one vector. A whole-batch transport failure + is left to the layers below instead; see the earlier transport tests. + """ async def go(): sleeps = [] @@ -1132,44 +1248,122 @@ def test_sdk_transport_error_recovers_without_losing_the_batch(monkeypatch): 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) + for transient in ( + APIConnectionError("connection reset"), + ConnectionError("connection reset"), + TimeoutError("request timed out"), + ): + sleeps.clear() - results = await store._call_with_retry([f"item{i}" for i in range(10)]) + provider = PerItemLimitThenTransientErrorAsEmbedding(transient) + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_isolated", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) - assert results == [[1.0, 0.0]] * 10 - assert provider.batch_sizes == [10, 10] - assert sleeps == [1.0] + 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 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.""" +def test_a_persistent_transport_error_is_not_multiplied_by_the_store(monkeypatch): + """A timeout is asked once: the wrapper and the SDK own the retry budget. + + The provider sees the 4 wrapper attempts x 3 SDK attempts those layers make, while the + store's max_retries=3 no longer multiplies them -- a persistent timeout used to cost 36 + provider attempts instead of 12. + """ async def go(): - sleeps = [] - - async def fake_sleep(delay): - sleeps.append(delay) - - provider = PerItemLimitThenTransientErrorAsEmbedding(APIConnectionError("connection reset")) + endpoint = offline_openai_endpoint(read_timeout).install(monkeypatch) + provider = openai_wrapper("t_openai_transport_budget") store = LocalEmbeddingStore( - name="t_local_embedding_sdk_isolated", + name="t_local_embedding_sdk_transport_budget", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(3)]) is None + assert provider.calls == 1 + assert endpoint.requests == 12 + assert store.is_healthy is False + + run(go()) + + +def test_a_transient_transport_error_still_recovers_through_the_sdk(monkeypatch): + """One timeout is answered by the SDK's own retry: the batch survives and never splits.""" + + async def go(): + state = {"timeouts": 1} + + def handler(request): + if state["timeouts"] > 0: + state["timeouts"] -= 1 + return read_timeout(request) + return embedding_payload(request) + + endpoint = offline_openai_endpoint(handler).install(monkeypatch) + provider = openai_wrapper("t_openai_transport_recover") + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_transport_recover", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + results = await store._call_with_retry([f"item{i}" for i in range(3)]) + + assert results == [[0.0, 1.0]] * 3 + assert provider.calls == 1 + assert endpoint.requests == 2 + + run(go()) + + +def test_an_isolated_item_still_recovers_through_the_sdk(monkeypatch): + """After a context-length rejection, a timeout on one remaining item is still recovered. + + The store asks the wrapper once per item; the retry that saves that item belongs to the SDK. + """ + + async def go(): + state = {"timeout_pending": True} + + def handler(request): + body = json.loads(request.content.decode("utf-8")) + inputs = body["input"] + if isinstance(inputs, str): + inputs = [inputs] + if any(len(item) > 32 for item in inputs): + return context_length_rejection() + if len(inputs) == 1 and state["timeout_pending"]: + state["timeout_pending"] = False + return read_timeout(request) + return embedding_payload(request) + + endpoint = offline_openai_endpoint(handler).install(monkeypatch) + provider = openai_wrapper("t_openai_transport_isolated") + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_transport_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] + assert results == [[0.0, 1.0], None, [0.0, 1.0]] + assert provider.calls == 4 + assert endpoint.requests == 5 run(go())