diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 84d2c293..4fad8361 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -194,6 +194,15 @@ class LocalEmbeddingStore(BaseEmbeddingStore): if attempt < self.max_retries - 1: await asyncio.sleep(2**attempt) except Exception as error: + if self._is_rate_limited(error): + 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 + return None if ( self.quota_retry_delay is not None and self._is_insufficient_quota(error) @@ -210,6 +219,19 @@ class LocalEmbeddingStore(BaseEmbeddingStore): self.is_healthy = False return None + @staticmethod + def _is_rate_limited(error: Exception) -> bool: + """Recognize an OpenAI-compatible 429 response without importing a provider SDK.""" + if LocalEmbeddingStore._is_insufficient_quota(error): + return False + if getattr(error, "status_code", None) == 429: + return True + body = getattr(error, "body", None) + if not isinstance(body, dict): + return False + details = body.get("error", body) + return isinstance(details, dict) and details.get("code") == "rate_limit_exceeded" + @staticmethod def _is_insufficient_quota(error: Exception) -> bool: """Recognize OpenAI-compatible quota errors without importing a provider SDK.""" diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 757e184b..51b7bf63 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -73,9 +73,44 @@ class FakeProviderModel: class InsufficientQuotaError(Exception): """OpenAI-compatible quota error used without importing the provider SDK.""" + status_code = 429 body = {"error": {"code": "insufficient_quota"}} +class RateLimitError(Exception): + """OpenAI-compatible 429 error used without importing the provider SDK.""" + + status_code = 429 + + +class RateLimitedThenSuccessAsEmbedding: + """Fail once with a 429, then return a valid embedding.""" + + dimensions = 2 + + def __init__(self): + self.calls = 0 + + async def __call__(self, texts: list[str], **_kwargs): + self.calls += 1 + if self.calls == 1: + raise RateLimitError("Requests are too frequent") + return [[1.0, 0.0] for _ in texts] + + +class AlwaysRateLimitedAsEmbedding: + """Always fail with a 429.""" + + dimensions = 2 + + def __init__(self): + self.calls = 0 + + async def __call__(self, _texts: list[str], **_kwargs): + self.calls += 1 + raise RateLimitError("Requests are too frequent") + + class QuotaThenSuccessAsEmbedding: """Fail once for quota, then return a valid embedding.""" @@ -270,6 +305,78 @@ def test_insufficient_quota_does_not_retry_without_opt_in(monkeypatch): run(go()) +def test_rate_limit_retries_with_backoff_without_opt_in(monkeypatch): + """A 429 must retry like a network error, with no quota_retry_delay required.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore(name="t_local_embedding_rate_limit", max_retries=2) + embedding = RateLimitedThenSuccessAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + result = await store._call_with_retry(["text"]) + + assert result == [[1.0, 0.0]] + assert embedding.calls == 2 + assert sleeps == [1.0] + + run(go()) + + +def test_is_rate_limited_recognizes_status_code_or_body_and_nothing_else(): + """A 429 is recognized either by status_code or a body-embedded code, never by guesswork.""" + + class StatusCodeError(Exception): + """A 429 identified by an HTTP status code attribute.""" + + status_code = 429 + + class BodyCodeError(Exception): + """A 429 identified only by a body-embedded error code.""" + + 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")) + + +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")) + + +def test_rate_limit_exhausts_retries_and_reports_unhealthy(monkeypatch): + """Repeated 429s still give up after max_retries, unlike an unclassified error.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore(name="t_local_embedding_rate_limit_exhausted", max_retries=3) + embedding = AlwaysRateLimitedAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + result = await store._call_with_retry(["text"]) + + assert result is None + assert store.is_healthy is False + assert embedding.calls == 3 + assert sleeps == [1.0, 2.0] + + run(go()) + + def test_vector_space_id_separates_models_of_equal_dimension(): """Two models of the same width must not claim the same vector space.""" common = {"backend": "openai", "dimensions": 1024, "credential": {"base_url": "https://example.com/v1"}}