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 10

refactor(embedding): simplify retry results and error classification (#656)

Browse source
This commit is contained in:
jinliyl 2026-10-21 21:15:27 +08:00 • committed by GitHub
parent a4f606c6a8
commit 69e3db8deb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 228 additions and 433 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

109
reme/components/embedding_store/_embedding_types.py Normal file
View file

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

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

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

146
tests/unit/test_local_embedding_store.py
View file

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