diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index d5e26e23..c0f1e257 100644 --- 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. diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 926b872c..9e01215a 100644 --- 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())