fix(embedding): retry 429 rate-limit errors instead of dropping the batch (#525)

* fix(embedding): retry 429 rate-limit errors instead of dropping the batch

openai.RateLimitError is not a TimeoutError/ConnectionError/OSError, so
_call_with_retry's except Exception branch caught it and returned None on
the first attempt with zero backoff. Add _is_rate_limited, mirroring the
existing _is_insufficient_quota duck-typed check, and retry a 429 with the
same exponential backoff used for network errors.

* fix(embedding): insufficient_quota errors carrying status_code=429 no longer bypass quota handling

_is_rate_limited() checked status_code == 429 first, so an OpenAI-compatible
insufficient_quota error (which also carries status_code=429) matched the
generic rate-limit branch before the code=insufficient_quota check ever ran.
That meant a real quota exhaustion retried on the wrong backoff (or not at
all, when quota_retry_delay is unset) instead of the dedicated quota_retry_delay
wait.

_is_rate_limited() now defers to _is_insufficient_quota() first. The existing
quota test double now sets status_code=429 to match the real OpenAI error
shape, which is what exposes the regression without the fix.

Signed-off-by: Amir Fathi <amirfathi.me@gmail.com>

* fix(embedding): log rate-limit retries

---------

Signed-off-by: Amir Fathi <amirfathi.me@gmail.com>
Co-authored-by: jinli.yl <jinli.yl@alibaba-inc.com>

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."""

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"}}