fix(embedding): retry SDK transport failures instead of splitting the batch

Third review follow-up on #579. `_call_with_retry` only caught the built-in `TimeoutError`, `ConnectionError` and `OSError`. `openai.APIConnectionError` and `openai.APITimeoutError` do not inherit from any of them (they go APIError -> OpenAIError -> Exception), so they reached the generic handler: `_is_provider_wide_failure()` did not recognize them either, and a multi-item request fanned out into per-item requests that failed for exactly the same reason. The same gap meant a single-item request got no ReMe-level retry at all after a transport failure. Transport failures are now recognized by a dedicated classifier, `_is_transient_transport_failure()`, covering the built-in socket errors as well as the SDK/httpx exception types (matched across the MRO, so subclasses such as `APITimeoutError` are included) without importing the provider SDK. They retry the original request with the existing backoff and never split; the retry budget is unchanged. Regression coverage: - test_sdk_transport_errors_are_retried_without_splitting APIConnectionError / APITimeoutError -> batch sizes [10, 10, 10], sleeps [1.0, 2.0], then failure with no single-item requests. - test_sdk_transport_error_recovers_without_losing_the_batch one APITimeoutError then success -> [[1.0, 0.0]] * 10, batch sizes [10, 10]. - test_isolated_item_retries_an_sdk_transport_error context-length rejection, then APIConnectionError on one isolated item -> that item's vector is retained. The existing context-length, 401/403, 5xx, quota and rate-limit tests are retained. Local verification with the repo's own hook arguments: pylint -> Your code has been rated at 10.00/10, exit 0 black -> 2 files would be left unchanged flake8 -> clean Branch-order discriminator check (classifiers extracted from the production file and executed against each error type): case old new 400 context-length SPLIT SPLIT 401 / 403 SPLIT STOP 500 SPLIT RETRY SDK APIConnectionError SPLIT RETRY SDK APITimeoutError SPLIT RETRY builtin ConnectionError RETRY RETRY builtin TimeoutError RETRY RETRY 429 RETRY RETRY insufficient_quota STOP STOP
2026-10-24 05:47:32 +00:00 · 2026-10-09 13:24:22 +08:00 · 2026-10-09 13:24:22 +08:00 · 930fddc4d2
commit 930fddc4d2
parent 254d895f96
2 changed files with 146 additions and 3 deletions
--- a/reme/components/embedding_store/local_embedding_store.py
+++ b/reme/components/embedding_store/local_embedding_store.py
@ -14,6 +14,25 @@ from ..as_embedding import BaseAsEmbedding
 Miss = tuple[int, str, str]  # (result_index, text, cache_key)
 _MAX_VECTOR_SPACE_ATTEMPTS = 3

+# 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.
+_TRANSPORT_FAILURE_TYPES = frozenset(
+    {
+        "APIConnectionError",
+        "APITimeoutError",
+        "ConnectError",
+        "ConnectTimeout",
+        "PoolTimeout",
+        "ReadError",
+        "ReadTimeout",
+        "RemoteProtocolError",
+        "WriteError",
+        "WriteTimeout",
+    },
+)
+

@R.register("local")
 class LocalEmbeddingStore(BaseEmbeddingStore):
@ -190,10 +209,14 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
                result = await self.as_embedding(texts, **kwargs)
                if result and len(result) == len(texts):
                    return result
-            except (TimeoutError, ConnectionError, OSError):
-                if attempt < self.max_retries - 1:
-                    await asyncio.sleep(2**attempt)
            except Exception as error:
+                if self._is_transient_transport_failure(error):
+                    # A connection failure or a timeout describes the endpoint, not any single
+                    # item: retry the very same request with the usual backoff instead of
+                    # fanning out into per-item requests that would fail the same way.
+                    if attempt < self.max_retries - 1:
+                        await asyncio.sleep(2**attempt)
+                    continue
                if self._is_rate_limited(error):
                    if attempt < self.max_retries - 1:
                        delay = 2**attempt
@ -259,6 +282,18 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
            results.append(None)
        return results

+    @staticmethod
+    def _is_transient_transport_failure(error: Exception) -> bool:
+        """Return whether an error is a transient transport failure worth retrying.
+
+        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 same request is retried and never split per item.
+        """
+        if isinstance(error, (TimeoutError, ConnectionError, OSError)):
+            return True
+        return any(cls.__name__ in _TRANSPORT_FAILURE_TYPES for cls in type(error).__mro__)
+
    @staticmethod
    def _is_provider_wide_failure(error: Exception) -> bool:
        """Return whether an error describes the provider rather than a single input.
--- a/tests/unit/test_local_embedding_store.py
+++ b/tests/unit/test_local_embedding_store.py
@ -218,6 +218,18 @@ class InternalServerError(Exception):
    status_code = 500


+class APIError(Exception):
+    """Base of the OpenAI SDK error hierarchy, reproduced without importing the SDK."""
+
+
+class APIConnectionError(APIError):
+    """OpenAI SDK connection failure: it does not inherit from the built-in ConnectionError."""
+
+
+class APITimeoutError(APIConnectionError):
+    """OpenAI SDK timeout, a subclass of APIConnectionError."""
+
+
 class AlwaysProviderWideFailureAsEmbedding:
    """Always fail with the same provider-wide error, recording every requested batch size."""

@ -232,6 +244,27 @@ class AlwaysProviderWideFailureAsEmbedding:
        raise self.error


+class TransientThenSuccessAsEmbedding:
+    """Fail the first `failures` calls with one error, then return valid embeddings."""
+
+    dimensions = 2
+    vector_space_id = "fakespace000"
+
+    def __init__(self, error: Exception, failures: int = 1):
+        self.error = error
+        self.failures = failures
+        self.batch_sizes: list[int] = []
+
+    def initialize_model(self):
+        """Mirror the real component's idempotent initialization hook."""
+
+    async def __call__(self, texts: list[str], **_kwargs):
+        self.batch_sizes.append(len(texts))
+        if len(self.batch_sizes) <= self.failures:
+            raise self.error
+        return [[1.0, 0.0] for _ in texts]
+
+
 class BadNodeEmbeddingStore(BaseEmbeddingStore):
    """Embedding store that returns wrong-dimensional vectors."""

@ -1026,3 +1059,78 @@ def test_provider_wide_server_errors_are_retried_without_splitting(monkeypatch):
        assert store.is_healthy is False

    run(go())
+
+
+def test_sdk_transport_errors_are_retried_without_splitting(monkeypatch):
+    """SDK connection/timeout errors are provider-wide: retry the batch, never split it."""
+
+    async def go():
+        sleeps = []
+
+        async def fake_sleep(delay):
+            sleeps.append(delay)
+
+        monkeypatch.setattr(asyncio, "sleep", fake_sleep)
+
+        for error in (APIConnectionError("connection reset"), APITimeoutError("timed out")):
+            sleeps.clear()
+            provider = AlwaysProviderWideFailureAsEmbedding(error)
+            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 provider.batch_sizes == [10, 10, 10]
+            assert sleeps == [1.0, 2.0]
+
+    run(go())
+
+
+def test_sdk_transport_error_recovers_without_losing_the_batch(monkeypatch):
+    """One SDK transport failure retries the original batch and keeps every vector."""
+
+    async def go():
+        sleeps = []
+
+        async def fake_sleep(delay):
+            sleeps.append(delay)
+
+        provider = TransientThenSuccessAsEmbedding(APITimeoutError("timed out"))
+        store = LocalEmbeddingStore(name="t_local_embedding_sdk_recover", max_retries=3)
+        store.as_embedding = provider
+        monkeypatch.setattr(asyncio, "sleep", fake_sleep)
+
+        results = await store._call_with_retry([f"item{i}" for i in range(10)])
+
+        assert results == [[1.0, 0.0]] * 10
+        assert provider.batch_sizes == [10, 10]
+        assert sleeps == [1.0]
+
+    run(go())
+
+
+def test_isolated_item_retries_an_sdk_transport_error(monkeypatch):
+    """After a context-length rejection, one isolated item still retries an SDK transport error."""
+
+    async def go():
+        sleeps = []
+
+        async def fake_sleep(delay):
+            sleeps.append(delay)
+
+        provider = PerItemLimitThenTransientErrorAsEmbedding(APIConnectionError("connection reset"))
+        store = LocalEmbeddingStore(
+            name="t_local_embedding_sdk_isolated",
+            max_retries=3,
+            enable_cache=False,
+        )
+        store.as_embedding = provider
+        monkeypatch.setattr(asyncio, "sleep", fake_sleep)
+
+        results = await store._call_with_retry(["ok", "x" * 64, "also ok"])
+
+        assert results == [[1.0, 0.0], None, [1.0, 0.0]]
+        assert provider.batch_sizes == [3, 1, 1, 1, 1]
+        assert provider.transient_failures == 1
+        assert sleeps == [1.0]
+
+    run(go())