From 69e3db8deb22221c39afb6c0b1f24e7cccec81bd Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 21 Oct 2026 21:15:27 +0800 Subject: [PATCH] refactor(embedding): simplify retry results and error classification (#656) --- .../embedding_store/_embedding_types.py | 109 +++++ .../embedding_store/local_embedding_store.py | 406 ++---------------- tests/unit/test_local_embedding_store.py | 146 ++++--- 3 files changed, 228 insertions(+), 433 deletions(-) create mode 100644 reme/components/embedding_store/_embedding_types.py diff --git a/reme/components/embedding_store/_embedding_types.py b/reme/components/embedding_store/_embedding_types.py new file mode 100644 index 00000000..69aebebe --- /dev/null +++ b/reme/components/embedding_store/_embedding_types.py @@ -0,0 +1,109 @@ +"""Internal embedding results and normalized provider errors.""" + +import re +from dataclasses import dataclass + +_PROVIDER_WIDE_STATUS_CODES = frozenset({401, 403, 404}) + + +# AgentScope's DashScope wrapper serializes the response into a RuntimeError. +_STATUS_IN_MESSAGE = re.compile(r'"status_code"\s*:\s*(\d{3})') + + +_CODE_IN_MESSAGE = re.compile(r'"code"\s*:\s*"([^"]+)"') + + +_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", +) + + +@dataclass(frozen=True) +class _EmbeddingResult: + """Keep recovered vectors together with the failure state that ends isolation.""" + + embeddings: list[list[float] | None] | None = None + terminal_failure: bool = False + + +@dataclass(frozen=True) +class _EmbeddingError: + """Normalize provider error shapes once before applying the retry policy.""" + + status: int | None + code: str | None + text: str + error_type: str + + @classmethod + def from_exception(cls, error: Exception) -> "_EmbeddingError": + """Read attributes, structured bodies, and serialized DashScope responses.""" + message = str(error) + body = getattr(error, "body", None) + status = None + for value in ( + getattr(error, "status_code", None), + getattr(error, "code", None), + getattr(getattr(error, "response", None), "status_code", None), + ): + # Integer codes are HTTP statuses; strings name business errors. + if isinstance(value, int) and not isinstance(value, bool): + status = value + break + if status is None and (match := _STATUS_IN_MESSAGE.search(message)): + status = int(match.group(1)) + + code = getattr(error, "code", None) + if not isinstance(code, str) or not code: + details = body.get("error", body) if isinstance(body, dict) else None + code = details.get("code") if isinstance(details, dict) else None + if not isinstance(code, str) or not code: + match = _CODE_IN_MESSAGE.search(message) + code = match.group(1) if match else None + text = f"{message} {body}" if isinstance(body, dict) else message + return cls(status, code, text.lower(), type(error).__name__) + + @property + def is_insufficient_quota(self) -> bool: + """Quota exhaustion takes precedence over a generic HTTP 429.""" + return self.code == "insufficient_quota" + + @property + def is_rate_limited(self) -> bool: + """Recognize rate limits while preserving the separate quota policy.""" + return not self.is_insufficient_quota and (self.code == "rate_limit_exceeded" or self.status == 429) + + @property + def is_provider_wide(self) -> bool: + """Credentials, missing routes, and server outages cannot be isolated per input.""" + return ( + self.status is not None and (self.status in _PROVIDER_WIDE_STATUS_CODES or 500 <= self.status < 600) + ) or self.error_type in ("AuthenticationError", "PermissionDeniedError", "NotFoundError", "InternalServerError") + + @property + def is_input_dependent(self) -> bool: + """Only input-length rejections may split; a 413 always requires a text signal.""" + if self.status not in (None, 400, 413): + return False + if self.status != 413 and self.error_type in _INPUT_DEPENDENT_ERROR_TYPES: + return True + return any(marker in self.text for marker in _INPUT_DEPENDENT_MARKERS) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index abb08caa..99e59369 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -2,12 +2,12 @@ import asyncio import hashlib -import re from collections import OrderedDict from pathlib import Path import numpy as np +from ._embedding_types import _EmbeddingError, _EmbeddingResult from .base_embedding_store import BaseEmbeddingStore from ..component_registry import R from ..as_embedding import BaseAsEmbedding @@ -15,12 +15,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, -# so subclasses such as APITimeoutError are covered) keeps this module free of SDK imports. +# Match SDK/httpx transport exceptions without importing optional providers. _TRANSPORT_FAILURE_TYPES = frozenset( { "APIConnectionError", @@ -36,60 +32,9 @@ _TRANSPORT_FAILURE_TYPES = frozenset( }, ) -# Transport errors an SDK retries on its own. A positive numeric retry budget is not evidence: -# AgentScope's Ollama model keeps the inherited ``max_retries=3`` next to an empty policy, and -# the number never buys a retry. Only a client that both has a budget and retries this class of -# error owns the failure; otherwise the store keeps its backoff. +# A client retry budget only applies to the transport exceptions it retries. _SDK_TRANSPORT_RETRY_TYPES = frozenset({"APIConnectionError", "APITimeoutError"}) -# 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})') - -# 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. -_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", - "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): @@ -238,17 +183,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore): async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]: texts = [text for _, text, _ in batch] - # 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, - ) + result = await self._call_with_retry(texts, **kwargs) + embeddings = result.embeddings if not embeddings or len(embeddings) != len(texts): return [] out: list[tuple[int, str, np.ndarray]] = [] @@ -264,221 +200,91 @@ 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 and not terminal_failure: - self.is_healthy = True - else: - self.is_healthy = False + # Retain partial vectors without hiding the outage that ended isolation. + self.is_healthy = bool(out) and not result.terminal_failure return out - 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. - - Transport failures are provider-wide: no individual input can explain them, so the - request is never split per item. Which layer retries them depends on the provider, not - on where the request came from: when the embedding model below already retries the - transport exceptions it raises (AgentScope's OpenAI model does, and the OpenAI SDK - retries each of those attempts), the store answers once and leaves that budget there. - A provider whose model has no retry policy -- Ollama inherits an empty retryable tuple - and its SDK turns ``httpx.ConnectError`` into a plain ``ConnectionError`` -- gets the - store's backoff instead, because otherwise a single dropped connection would cost the - vector. - - ``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, 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: - """Record that this failure ends the work for the whole batch.""" - if terminal_failure is not None: - terminal_failure[:] = [True] - + async def _call_with_retry(self, texts: list[str], **kwargs) -> _EmbeddingResult: + """Retry failures owned by the store and isolate input-length rejections.""" for attempt in range(self.max_retries): try: - result = await self.as_embedding(texts, **kwargs) - if result and len(result) == len(texts): - return result + embeddings = await self.as_embedding(texts, **kwargs) + if embeddings and len(embeddings) == len(texts): + return _EmbeddingResult(embeddings) except Exception as error: if self._is_transport_failure(error): - # A connection failure or a timeout describes the endpoint, not any single - # item, so this request is never split per item. Whether the store retries - # it is a property of the provider: when the model below has its own retry - # policy, that layer (and the SDK under it) already answers the failure -- - # a batch they answer after 12 attempts used to reach the provider 36 times - # with this store retrying on top of them, and an isolated item used to - # multiply the same budget a third time. + # Preserve recovery where the provider has no matching retry policy; + # otherwise leave its budget intact instead of adding another loop. if self._provider_owns_transport_retry(error): self.logger.exception( "Embedding request failed: transport error owned by the provider, not splitting", ) - self.is_healthy = False - stop() break if attempt < self.max_retries - 1: await asyncio.sleep(2**attempt) continue - self.logger.exception( - "Embedding request failed: transport error after exhausting retries", - ) - self.is_healthy = False - stop() + self.logger.exception("Embedding request failed: transport error after exhausting retries") break - if self._is_rate_limited(error): + + details = _EmbeddingError.from_exception(error) + if details.is_rate_limited: if attempt < self.max_retries - 1: delay = 2**attempt self.logger.warning(f"Embedding rate limited; retrying in {delay:.1f}s") await asyncio.sleep(delay) 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 - and self._is_insufficient_quota(error) - and attempt < self.max_retries - 1 - ): - self.logger.warning( - f"Embedding quota exhausted; retrying in {self.quota_retry_delay:.1f}s", - ) - 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 - stop() - 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 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. 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 - 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. + if details.is_insufficient_quota: + if self.quota_retry_delay is not None and attempt < self.max_retries - 1: + self.logger.warning(f"Embedding quota exhausted; retrying in {self.quota_retry_delay:.1f}s") + await asyncio.sleep(self.quota_retry_delay) + continue + self.logger.exception("Embedding request failed: quota exhausted") + break + if details.is_provider_wide: + # Server-error retries belong to the wrapper/SDK, not to item isolation. + self.logger.exception("Embedding request failed: provider-wide error, not splitting") + break + if details.is_input_dependent: self.logger.exception("Embedding request failed: input-dependent error") self.is_healthy = False if len(texts) > 1: - 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 - # with one request per item. + return await self._embed_individually(texts, **kwargs) + return _EmbeddingResult() self.logger.exception("Embedding request failed: unrecognized error, not splitting") - stop() break self.is_healthy = False - stop() - return None + return _EmbeddingResult(terminal_failure=True) - 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 - 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 that a rate limit and - a transient transport failure still get the retries that layer owns 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] = [] + async def _embed_individually(self, texts: list[str], **kwargs) -> _EmbeddingResult: + """Keep healthy siblings, stopping when a failure is no longer specific to an input.""" + embeddings: list[list[float] | None] = [] 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]) + result = await self._call_with_retry([text], **kwargs) + if result.embeddings: + embeddings.append(result.embeddings[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. 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] + embeddings.append(None) + if result.terminal_failure: 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 + embeddings.extend([None] * (len(texts) - index - 1)) + return _EmbeddingResult(embeddings, terminal_failure=True) + return _EmbeddingResult(embeddings) @staticmethod 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 request is never split per item. Who retries them is decided - per provider by ``_provider_owns_transport_retry`` rather than per call site. - """ + """Recognize transport errors that must never trigger per-item isolation.""" if isinstance(error, (TimeoutError, ConnectionError, OSError)): return True return any(cls.__name__ in _TRANSPORT_FAILURE_TYPES for cls in type(error).__mro__) def _provider_owns_transport_retry(self, error: Exception) -> bool: - """Return whether a lower layer retries *this* transport failure. - - Retry ownership is a property of the provider and of the error, not of the call site. - A model whose ``_get_retryable_exceptions()`` covers the error owns the retry, and so - does an SDK that retries that class of error. Everything else must keep the store's - backoff, or one dropped connection costs the vector. - - Two shapes are deliberately not treated as ownership. AgentScope's Ollama embedding - model declares an empty policy but keeps the inherited ``max_retries=3``: its loop - re-raises every exception outside the empty tuple, so the number never buys a retry. - AgentScope's DashScope model declares ``(RuntimeError,)``, which does not cover a - connection error raised underneath it. A nonempty policy is therefore checked against - the error that actually occurred rather than trusted as a whole. - - A model that exposes no policy at all -- test doubles and custom providers -- is only - trusted to an SDK that has a positive retry budget *and* retries this class of - transport error. Otherwise the store answers, as the base implementation did. - """ + """Require both a matching exception policy and an enabled retry budget.""" model = getattr(self.as_embedding, "model", None) if model is None: return False @@ -510,13 +316,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): @staticmethod def _sdk_owns_transport_retry(model, error: Exception) -> bool: - """Return whether the SDK under the model retries this transport error itself. - - The numeric ``max_retries`` only counts once the client also identifies this error as - retryable. The OpenAI client retries ``APIConnectionError`` / ``APITimeoutError``; - Ollama's client has no retry budget at all, and AgentScope's DashScope path calls the - module directly, so neither has a client to consult. - """ + """Check the SDK's budget and supported transport errors independently of the model.""" if not LocalEmbeddingStore._is_transport_failure(error): return False client = getattr(model, "client", None) @@ -527,118 +327,6 @@ class LocalEmbeddingStore(BaseEmbeddingStore): return False return any(cls.__name__ in _SDK_TRANSPORT_RETRY_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. - - 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 = 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", - "PermissionDeniedError", - "NotFoundError", - "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. 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: - 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) - 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 _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.""" - code = LocalEmbeddingStore._error_code(error) - if code == "insufficient_quota": - return False - if code == "rate_limit_exceeded": - return True - return LocalEmbeddingStore._http_status(error) == 429 - - @staticmethod - def _is_insufficient_quota(error: Exception) -> bool: - """Recognize OpenAI-compatible quota errors without importing a provider SDK.""" - return LocalEmbeddingStore._error_code(error) == "insufficient_quota" - def _validate_dim(self, emb: np.ndarray) -> bool: """Return whether an embedding exactly matches the configured dimension.""" return len(emb) == self.dimensions diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 912c1789..23444703 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -11,6 +11,7 @@ import numpy as np import openai from reme.components.as_embedding import DashScopeAsEmbedding, OllamaAsEmbedding, OpenAIAsEmbedding +from reme.components.embedding_store._embedding_types import _EmbeddingError from reme.components.embedding_store.base_embedding_store import BaseEmbeddingStore from reme.components.embedding_store.local_embedding_store import LocalEmbeddingStore from reme.schema import EmbNode @@ -566,7 +567,7 @@ def test_insufficient_quota_waits_sixty_seconds_before_retry(monkeypatch): store.as_embedding = embedding monkeypatch.setattr(asyncio, "sleep", fake_sleep) - result = await store._call_with_retry(["text"]) + result = (await store._call_with_retry(["text"])).embeddings assert result == [[1.0, 0.0]] assert embedding.calls == 2 @@ -589,7 +590,7 @@ def test_insufficient_quota_does_not_retry_without_opt_in(monkeypatch): store.as_embedding = embedding monkeypatch.setattr(asyncio, "sleep", fake_sleep) - result = await store._call_with_retry(["text"]) + result = (await store._call_with_retry(["text"])).embeddings assert result is None assert embedding.calls == 1 @@ -612,7 +613,7 @@ def test_rate_limit_retries_with_backoff_without_opt_in(monkeypatch): store.as_embedding = embedding monkeypatch.setattr(asyncio, "sleep", fake_sleep) - result = await store._call_with_retry(["text"]) + result = (await store._call_with_retry(["text"])).embeddings assert result == [[1.0, 0.0]] assert embedding.calls == 2 @@ -634,16 +635,16 @@ def test_is_rate_limited_recognizes_status_code_or_body_and_nothing_else(): body = {"error": {"code": "rate_limit_exceeded"}} - assert LocalEmbeddingStore._is_rate_limited(StatusCodeError()) - assert LocalEmbeddingStore._is_rate_limited(BodyCodeError()) - assert not LocalEmbeddingStore._is_rate_limited(ValueError("unrelated")) + assert _EmbeddingError.from_exception(StatusCodeError()).is_rate_limited + assert _EmbeddingError.from_exception(BodyCodeError()).is_rate_limited + assert not _EmbeddingError.from_exception(ValueError("unrelated")).is_rate_limited def test_is_rate_limited_defers_to_insufficient_quota_on_status_code_429(): """An insufficient_quota error carrying status_code=429 is not the generic rate-limit case.""" - assert not LocalEmbeddingStore._is_rate_limited(InsufficientQuotaError("quota exhausted")) - assert LocalEmbeddingStore._is_insufficient_quota(InsufficientQuotaError("quota exhausted")) + assert not _EmbeddingError.from_exception(InsufficientQuotaError("quota exhausted")).is_rate_limited + assert _EmbeddingError.from_exception(InsufficientQuotaError("quota exhausted")).is_insufficient_quota def test_rate_limit_exhausts_retries_and_reports_unhealthy(monkeypatch): @@ -660,7 +661,7 @@ def test_rate_limit_exhausts_retries_and_reports_unhealthy(monkeypatch): store.as_embedding = embedding monkeypatch.setattr(asyncio, "sleep", fake_sleep) - result = await store._call_with_retry(["text"]) + result = (await store._call_with_retry(["text"])).embeddings assert result is None assert store.is_healthy is False @@ -1033,7 +1034,7 @@ def test_batch_failure_falls_back_to_per_item_embedding(): 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"]) + results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings assert results == [[1.0, 0.0], None, [1.0, 0.0]] assert provider.batch_sizes == [3, 1, 1, 1] @@ -1065,7 +1066,9 @@ def test_single_item_failure_still_reports_a_failed_batch(): 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 + result = await store._call_with_retry(["x" * 64]) + assert result.embeddings is None + assert result.terminal_failure is False run(go()) @@ -1084,7 +1087,7 @@ def test_rate_limited_batch_is_not_split_into_per_item_requests(monkeypatch): store.as_embedding = embedding monkeypatch.setattr(asyncio, "sleep", fake_sleep) - assert await store._call_with_retry(["a", "b"]) is None + assert (await store._call_with_retry(["a", "b"])).embeddings is None assert embedding.calls == 2 assert sleeps == [1.0] @@ -1105,7 +1108,7 @@ def test_insufficient_quota_batch_is_not_split_into_per_item_requests(monkeypatc store.as_embedding = embedding monkeypatch.setattr(asyncio, "sleep", fake_sleep) - assert await store._call_with_retry(["a", "b", "c"]) is None + assert (await store._call_with_retry(["a", "b", "c"])).embeddings is None assert embedding.batch_sizes == [3] assert store.is_healthy is False assert not sleeps @@ -1131,7 +1134,7 @@ def test_exhausted_quota_retries_still_do_not_split_the_batch(monkeypatch): store.as_embedding = embedding monkeypatch.setattr(asyncio, "sleep", fake_sleep) - assert await store._call_with_retry(["a", "b"]) is None + assert (await store._call_with_retry(["a", "b"])).embeddings is None assert embedding.batch_sizes == [2, 2, 2] assert sleeps == [5.0, 5.0] @@ -1163,7 +1166,7 @@ def test_per_item_fallback_keeps_a_rate_limited_vector(monkeypatch): 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"])).embeddings assert results == [[1.0, 0.0], None, [1.0, 0.0]] assert provider.batch_sizes == [3, 1, 1, 1, 1] @@ -1185,7 +1188,7 @@ def test_provider_wide_credentials_errors_are_not_split_into_per_item_requests() 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 (await store._call_with_retry([f"item{i}" for i in range(10)])).embeddings is None assert provider.batch_sizes == [10] assert store.is_healthy is False @@ -1206,7 +1209,7 @@ def test_provider_wide_server_errors_are_not_retried_without_splitting(monkeypat 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 (await store._call_with_retry([f"item{i}" for i in range(10)])).embeddings is None # 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] @@ -1237,7 +1240,7 @@ def test_sdk_transport_errors_are_provider_wide_and_not_retried_here(monkeypatch 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 (await store._call_with_retry([f"item{i}" for i in range(10)])).embeddings is None # 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. @@ -1279,7 +1282,7 @@ def test_a_provider_without_a_retry_policy_keeps_the_store_backoff(monkeypatch): 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"])).embeddings assert results == [[1.0, 0.0], None, [1.0, 0.0]] assert provider.batch_sizes == [3, 1, 1, 1, 1] @@ -1306,7 +1309,7 @@ def test_a_persistent_transport_error_is_not_multiplied_by_the_store(monkeypatch ) store.as_embedding = provider - assert await store._call_with_retry([f"item{i}" for i in range(3)]) is None + assert (await store._call_with_retry([f"item{i}" for i in range(3)])).embeddings is None assert provider.calls == 1 assert endpoint.requests == 12 assert store.is_healthy is False @@ -1335,7 +1338,7 @@ def test_a_transient_transport_error_still_recovers_through_the_sdk(monkeypatch) ) store.as_embedding = provider - results = await store._call_with_retry([f"item{i}" for i in range(3)]) + results = (await store._call_with_retry([f"item{i}" for i in range(3)])).embeddings assert results == [[0.0, 1.0]] * 3 assert provider.calls == 1 @@ -1374,7 +1377,7 @@ def test_an_isolated_item_still_recovers_through_the_sdk(monkeypatch): ) store.as_embedding = provider - results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings assert results == [[0.0, 1.0], None, [0.0, 1.0]] assert provider.calls == 4 @@ -1387,16 +1390,15 @@ 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 + assert _EmbeddingError.from_exception(error).is_provider_wide is True over_limit = ContextLengthError("the input length exceeds the context length") - assert store._is_provider_wide_failure(over_limit) is False + assert _EmbeddingError.from_exception(over_limit).is_provider_wide is False run(go()) @@ -1415,7 +1417,7 @@ def test_provider_wide_not_found_errors_are_not_split(monkeypatch): 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 (await store._call_with_retry([f"item{i}" for i in range(10)])).embeddings is None assert provider.batch_sizes == [10] assert not sleeps @@ -1489,26 +1491,22 @@ class DashScopeRuntimeError(RuntimeError): 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 + assert _EmbeddingError.from_exception(AuthenticationError("bad key")).status == 401 + assert _EmbeddingError.from_exception(NotFoundError("model_not_found")).status == 404 + assert _EmbeddingError.from_exception(GenAiClientError(401)).status == 401 + assert _EmbeddingError.from_exception(GenAiServerError(503)).status == 503 + assert _EmbeddingError.from_exception(WrappedTransportError(404)).status == 404 + assert _EmbeddingError.from_exception(DashScopeRuntimeError.for_status(401)).status == 401 + assert _EmbeddingError.from_exception(ValueError("no status here")).status 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 + assert _EmbeddingError.from_exception(StringCodeError("invalid key")).status is None + assert _EmbeddingError.from_exception(StringCodeError("invalid key")).is_insufficient_quota is False + assert _EmbeddingError.from_exception(InsufficientQuotaError("quota exhausted")).status == 429 + assert _EmbeddingError.from_exception(InsufficientQuotaError("quota exhausted")).is_insufficient_quota is True def test_google_genai_service_failures_are_not_split_into_per_item_requests(monkeypatch): @@ -1541,7 +1539,7 @@ def test_google_genai_service_failures_are_not_split_into_per_item_requests(monk ) store.as_embedding = provider - assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert (await store._call_with_retry([f"item{i}" for i in range(10)])).embeddings is None assert provider.batch_sizes == expected_batches run(go()) @@ -1572,7 +1570,7 @@ def test_dashscope_runtime_errors_are_not_split_into_per_item_requests(monkeypat ) store.as_embedding = provider - assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert (await store._call_with_retry([f"item{i}" for i in range(10)])).embeddings is None assert provider.batch_sizes == expected_batches run(go()) @@ -1586,7 +1584,7 @@ def test_unrecognized_errors_are_not_split_into_per_item_requests(): 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 (await store._call_with_retry(["a", "b", "c"])).embeddings is None assert provider.batch_sizes == [3] run(go()) @@ -1595,19 +1593,17 @@ def test_unrecognized_errors_are_not_split_into_per_item_requests(): 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( + _EmbeddingError.from_exception( ContextLengthError("the input length exceeds the context length"), - ) + ).is_input_dependent 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 + assert _EmbeddingError.from_exception(ContextLengthError("too long")).is_input_dependent is True + assert _EmbeddingError.from_exception(AuthenticationError("bad key")).is_input_dependent is False + assert _EmbeddingError.from_exception(InternalServerError("boom")).is_input_dependent is False + assert _EmbeddingError.from_exception(GenAiClientError(400)).is_input_dependent is False + assert _EmbeddingError.from_exception(RuntimeError("provider said no")).is_input_dependent is False def test_bad_request_without_an_input_length_signal_is_not_split(): @@ -1618,7 +1614,7 @@ def test_bad_request_without_an_input_length_signal_is_not_split(): 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 (await store._call_with_retry(["a", "b", "c"])).embeddings is None assert provider.batch_sizes == [3] run(go()) @@ -1640,9 +1636,10 @@ def test_isolation_stops_when_a_provider_wide_failure_hits(): ) store.as_embedding = provider - results = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"]) + result = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"]) - assert results == [[1.0, 0.0], None, None, None] + assert result.embeddings == [[1.0, 0.0], None, None, None] + assert result.terminal_failure is True # 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] @@ -1700,9 +1697,10 @@ def test_isolation_keeps_going_while_only_inputs_fail(): ) store.as_embedding = provider - results = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"]) + result = 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 result.embeddings == [[1.0, 0.0], None, [1.0, 0.0], [1.0, 0.0]] + assert result.terminal_failure is False assert provider.batch_sizes == [4, 1, 1, 1, 1] run(go()) @@ -1725,8 +1723,8 @@ def test_dashscope_serialized_quota_errors_keep_their_quota_policy(monkeypatch): 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 + assert _EmbeddingError.from_exception(error).is_insufficient_quota is True + assert _EmbeddingError.from_exception(error).is_rate_limited is False provider = AlwaysProviderWideFailureAsEmbedding(error) store = LocalEmbeddingStore( @@ -1737,7 +1735,7 @@ def test_dashscope_serialized_quota_errors_keep_their_quota_policy(monkeypatch): ) store.as_embedding = provider - assert await store._call_with_retry(["a", "b"]) is None + assert (await store._call_with_retry(["a", "b"])).embeddings is None assert provider.batch_sizes == [2, 2, 2] assert sleeps == [7.0, 7.0] @@ -1799,7 +1797,7 @@ def test_context_length_rejection_reported_as_http_413_is_isolated(): ) store.as_embedding = provider - results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings assert results == [[1.0, 0.0], None, [1.0, 0.0]] assert provider.batch_sizes == [3, 1, 1, 1] @@ -1821,7 +1819,7 @@ def test_bare_http_413_without_an_input_signal_is_not_split(): 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 (await store._call_with_retry(["a", "b", "c"])).embeddings is None assert provider.batch_sizes == [3] run(go()) @@ -1830,18 +1828,18 @@ def test_bare_http_413_without_an_input_signal_is_not_split(): 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( + _EmbeddingError.from_exception( ContextLength413Error("the input length exceeds the context length"), - ) + ).is_input_dependent 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 + assert _EmbeddingError.from_exception(ContextLength413Error("too long")).is_input_dependent is True + assert ( + _EmbeddingError.from_exception(RequestBodyTooLargeError("request entity too large")).is_input_dependent is False + ) + assert _EmbeddingError.from_exception(ContextLengthError("too long")).is_input_dependent is True + assert _EmbeddingError.from_exception(GenAiClientError(400)).is_input_dependent is False class _NoRetryPolicyModel: @@ -1904,7 +1902,7 @@ def test_a_provider_with_an_empty_retry_policy_recovers_a_connection_error(monke ) store.as_embedding = provider - results = await store._call_with_retry(["a", "b", "c"]) + results = (await store._call_with_retry(["a", "b", "c"])).embeddings assert results == [[1.0, 0.0]] * 3 assert provider.batch_sizes == [3, 3] @@ -1939,7 +1937,7 @@ def test_a_policy_that_does_not_cover_the_error_keeps_the_store_backoff(monkeypa ) store.as_embedding = provider - results = await store._call_with_retry(["a", "b", "c"]) + results = (await store._call_with_retry(["a", "b", "c"])).embeddings assert results == [[1.0, 0.0]] * 3 assert provider.batch_sizes == [3, 3] @@ -1976,7 +1974,7 @@ def test_an_isolated_item_does_not_multiply_sdk_transport_retries(monkeypatch): ) store.as_embedding = provider - results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings assert results == [None, None, None] assert provider.calls == 2