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