scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-10-23 05:35:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 9

fix(embedding): classify every provider status shape and split only input errors

Browse source 
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.
This commit is contained in:
BeiMu-new 2026-10-12 17:44:21 +08:00
parent eba753f783
commit 5e70ec0476
2 changed files with 295 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

112
reme/components/embedding_store/local_embedding_store.py
View file

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

194
tests/unit/test_local_embedding_store.py
View file

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