fix(embedding): recover over-limit items with a bounded input-shortening retry (#655)

A per-item context-length rejection survives the per-item split: the item is resent unchanged and the provider keeps rejecting it, so the chunk never gets a vector and a scoped reindex can never finalize. Re-prepare that single item under a smaller character budget and retry it once, bounded by max_retries.

Only the request shrinks. The chunk keeps its full source text, the vector space is unchanged, and the shortening is logged so the reduced coverage is visible.

Shortening needs an attempt left to send. On the final attempt a shorter request cannot be delivered, so that branch now falls through to the existing input-dependent handling: the rejection stays specific to its item instead of being recorded as a terminal batch failure that drops the healthy items after it. The retry lives inside the input-dependent branch, so the same error is classified once, and the single-use budget helper is inlined.

Follow-up to the merged per-item isolation change (#579), rebased onto the retry-result refactor (#656).

reme/components/embedding_store/base_embedding_store.py

@@ -36,7 +36,7 @@ class BaseEmbeddingStore(BaseComponent):
         self.health_check_timeout = health_check_timeout
         self.is_healthy: bool = True
 
-    def _truncate(self, text: str) -> str:
+    def _truncate(self, text: str, *, max_length: int | None = None) -> str:
         """Truncate text using a CJK-aware character budget.
 
         ASCII text keeps its historical character limit. For non-ASCII text,
@@ -44,8 +44,12 @@ class BaseEmbeddingStore(BaseComponent):
         characters cost 1.5 units because they commonly consume more
         embedding tokens. The estimate reserves a 5% safety margin, and
         integer half-units avoid floating-point boundary errors.
+
+        ``max_length`` overrides the configured ``max_input_length`` for this call only, so a
+        caller can re-prepare a rejected input under a smaller budget without mutating shared
+        configuration or changing the vector space.
         """
-        limit = max(0, self.max_input_length)
+        limit = max(0, self.max_input_length if max_length is None else max_length)
         if text.isascii():
             return text[:limit]
 

reme/components/embedding_store/local_embedding_store.py

@@ -206,6 +206,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
 
     async def _call_with_retry(self, texts: list[str], **kwargs) -> _EmbeddingResult:
         """Retry failures owned by the store and isolate input-length rejections."""
+        shortened = False
         for attempt in range(self.max_retries):
             try:
                 embeddings = await self.as_embedding(texts, **kwargs)
@@ -247,6 +248,19 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
                     self.logger.exception("Embedding request failed: provider-wide error, not splitting")
                     break
                 if details.is_input_dependent:
+                    # Retry one item with half the character budget without changing its source text.
+                    # Require a remaining attempt so exhausted input failures stay non-terminal.
+                    if not shortened and len(texts) == 1 and attempt < self.max_retries - 1:
+                        shorter = self._truncate(texts[0], max_length=max(1, self.max_input_length // 2))
+                        if shorter != texts[0]:
+                            self.logger.warning(
+                                f"Embedding input rejected as too long; retrying with a smaller budget "
+                                f"(chars={len(texts[0])} -> {len(shorter)})",
+                            )
+                            texts = [shorter]
+                            shortened = True
+                            await asyncio.sleep(2**attempt)
+                            continue
                     self.logger.exception("Embedding request failed: input-dependent error")
                     self.is_healthy = False
                     if len(texts) > 1:

tests/unit/test_file_store_consistency.py

@@ -2160,3 +2160,78 @@ async def test_faiss_delete_queries_graph_once():
 
         store.file_graph.get_nodes = original_get_nodes
         await store.close()
+
+
+class OverLimitError(Exception):
+    """OpenAI-compatible 400 raised for an item over the provider's per-item window."""
+
+    status_code = 400
+    body = {"error": {"message": "the input length exceeds the context length"}}
+
+
+class OverLimitAsEmbedding:
+    """Reject any text longer than ``max_chars`` and record every text it did embed."""
+
+    dimensions = 2
+    vector_space_id = "fakespace000"
+
+    def __init__(self, max_chars: int):
+        self.max_chars = max_chars
+        self.embedded_texts: list[str] = []
+
+    def initialize_model(self):
+        """Mirror the real component's idempotent initialization hook."""
+
+    async def __call__(self, texts: list[str], **_kwargs):
+        if any(len(text) > self.max_chars for text in texts):
+            raise OverLimitError("the input length exceeds the context length")
+        self.embedded_texts.extend(texts)
+        return [[1.0, 0.0] for _ in texts]
+
+
+@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store])
+@pytest.mark.asyncio
+async def test_embedding_reindex_completes_when_a_chunk_exceeds_the_provider_limit(store_factory):
+    """An over-limit chunk must not block the embedding gate forever.
+
+    Reproduces the reported dead end: the provider rejects one chunk as longer than its
+    token window, so the chunk never gets a vector, the scoped reindex keeps raising and
+    the gate never clears. A bounded input-shortening retry has to recover that chunk
+    instead, so reindex completes and the gate opens with every vector in place.
+    """
+    with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
+        store = store_factory(name="t_reindex_over_limit_chunk")
+        await store.start()
+        long_text = "x" * 64
+        await set_chunks_with_graph(
+            store,
+            {
+                "good": chunk("good", "a.md", "alpha text"),
+                "long": chunk("long", "b.md", long_text),
+            },
+        )
+        provider = OverLimitAsEmbedding(max_chars=32)
+        embedding = LocalEmbeddingStore(
+            name="t_reindex_over_limit_chunk_embedding",
+            enable_cache=False,
+            max_input_length=64,
+        )
+        embedding.as_embedding = provider
+        store.embedding_store = embedding
+        _ensure_zvec_collection(store)
+        if isinstance(store, FaissLocalFileStore):
+            store._rebuild_index()
+        elif isinstance(store, ZvecLocalFileStore):
+            store._rebuild_collection()
+        await store.require_embedding_rebuild()
+
+        result = await store.reindex("embedding")
+
+        assert result == {"indexed": 2, "scope": "embedding"}
+        assert store._embedding_rebuild_pending is False
+        # The health probe also reaches the provider, so only the chunk texts matter here.
+        assert "alpha text" in provider.embedded_texts
+        assert "x" * 32 in provider.embedded_texts
+        assert long_text not in provider.embedded_texts
+
+        await store.close()

tests/unit/test_local_embedding_store.py

@@ -1762,3 +1762,241 @@ async def test_service_failures_preserve_the_batch_retry_budget(monkeypatch, err
     assert provider.batch_sizes == [batch_size] * attempts
     assert len(sleeps) == attempts - 1
     assert store.is_healthy is False
+
+
+# -- Bounded input-shortening recovery for over-limit items --
+
+
+class ShrinkablePerItemLimitAsEmbedding:
+    """Reject any text longer than ``max_chars``, so only a shorter request can embed it.
+
+    Mirrors a provider that enforces a token window the store's character budget cannot see:
+    the same text stays rejected no matter how often it is resent, so the only way to keep
+    its vector is to send less of it.
+    """
+
+    dimensions = 2
+    vector_space_id = "fakespace000"
+
+    def __init__(self, max_chars: int):
+        self.max_chars = max_chars
+        self.batch_sizes: list[int] = []
+        self.embedded_texts: list[str] = []
+
+    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 any(len(text) > self.max_chars for text in texts):
+            raise ContextLengthError("the input length exceeds the context length")
+        self.embedded_texts.extend(texts)
+        return [[1.0, 0.0] for _ in texts]
+
+
+class RateLimitedError(Exception):
+    """OpenAI-compatible 429 used to burn the retry budget before a length rejection."""
+
+    status_code = 429
+    code = "rate_limit_exceeded"
+
+
+class LateOverLimitAsEmbedding:
+    """Rate-limit the isolated long item, then reject it as too long on its final attempt."""
+
+    dimensions = 2
+    vector_space_id = "fakespace000"
+
+    def __init__(self, max_chars: int, single_call_limits: int = 2):
+        self.max_chars = max_chars
+        self.single_call_limits = single_call_limits
+        self.batch_sizes: list[int] = []
+        self.embedded_texts: list[str] = []
+
+    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 any(len(text) > self.max_chars for text in texts):
+            if len(texts) == 1 and self.single_call_limits:
+                self.single_call_limits -= 1
+                raise RateLimitedError("rate limit exceeded")
+            raise ContextLengthError("the input length exceeds the context length")
+        self.embedded_texts.extend(texts)
+        return [[1.0, 0.0] for _ in texts]
+
+
+class UnrelatedBadRequestError(Exception):
+    """A 400 that has nothing to do with the input length, e.g. a malformed parameter."""
+
+    status_code = 400
+
+
+class UnrelatedBadRequestAsEmbedding:
+    """Always fail with a 400 that does not name the input length."""
+
+    dimensions = 2
+    vector_space_id = "fakespace000"
+
+    def __init__(self):
+        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))
+        raise UnrelatedBadRequestError("unsupported encoding format")
+
+
+@pytest.mark.asyncio
+async def test_over_limit_item_is_recovered_by_a_shorter_request():
+    """A per-item rejection that survives isolation must be retried with less input."""
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=40)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_recovery",
+        enable_cache=False,
+        max_input_length=80,
+    )
+    store.as_embedding = provider
+    long_text = "x" * 80
+
+    result = await store._call_with_retry(["ok", long_text, "also ok"])
+
+    assert result.embeddings == [[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]]
+    assert result.terminal_failure is False
+    assert provider.batch_sizes == [3, 1, 1, 1, 1]
+    assert long_text not in provider.embedded_texts
+    assert provider.embedded_texts == ["ok", "x" * 40, "also ok"]
+
+
+@pytest.mark.asyncio
+async def test_shortened_retry_leaves_the_configured_budget_alone():
+    """Recovery is per request: the shared character budget must not shrink for everyone."""
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=40)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_budget",
+        enable_cache=False,
+        max_input_length=80,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["x" * 80])
+
+    assert result.embeddings == [[1.0, 0.0]]
+    assert store.max_input_length == 80
+
+
+@pytest.mark.asyncio
+async def test_batch_rejected_for_one_item_keeps_every_vector():
+    """The end-to-end promise: an over-limit chunk no longer costs its batch mates or itself."""
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_nodes",
+        enable_cache=False,
+        max_input_length=80,
+    )
+    store.as_embedding = ShrinkablePerItemLimitAsEmbedding(max_chars=40)
+    nodes = [EmbNode(text="first"), EmbNode(text="y" * 80), EmbNode(text="third")]
+
+    await store.get_node_embeddings(nodes)
+
+    assert all(node.embedding is not None for node in nodes)
+
+
+@pytest.mark.asyncio
+async def test_input_length_retry_is_skipped_when_the_input_cannot_shrink():
+    """A rejection that a smaller budget cannot change must not be retried at all."""
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=8)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_noop",
+        enable_cache=False,
+        max_input_length=8192,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["x" * 64])
+
+    assert result.embeddings is None
+    assert provider.batch_sizes == [1]
+
+
+@pytest.mark.asyncio
+async def test_shortening_needs_an_attempt_left_to_send():
+    """With no retry budget left the store gives up instead of resending a shorter input."""
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=40)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_bounded",
+        enable_cache=False,
+        max_input_length=80,
+        max_retries=1,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["x" * 80])
+
+    assert result.embeddings is None
+    assert result.terminal_failure is False
+    assert provider.batch_sizes == [1]
+
+
+@pytest.mark.asyncio
+async def test_final_attempt_rejection_keeps_the_healthy_siblings():
+    """An input-length rejection on the last attempt must not stop the per-item isolation.
+
+    Shortening cannot be sent without an attempt left, so that request has to fall through to
+    the input-dependent path. Otherwise the single item's rejection is recorded as a terminal
+    batch failure and every item after it is never requested -- the reported regression.
+    """
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=40)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_final_attempt",
+        enable_cache=False,
+        max_input_length=80,
+        max_retries=1,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["ok", "x" * 80, "after"])
+
+    assert result.embeddings == [[1.0, 0.0], None, [1.0, 0.0]]
+    assert result.terminal_failure is False
+    assert provider.embedded_texts == ["ok", "after"]
+    assert provider.batch_sizes == [3, 1, 1, 1]
+
+
+@pytest.mark.asyncio
+async def test_rate_limited_item_rejected_on_its_final_attempt_keeps_the_siblings():
+    """The same regression through rate limits: two 429s burn the budget before the rejection."""
+    provider = LateOverLimitAsEmbedding(max_chars=40, single_call_limits=2)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_after_rate_limits",
+        enable_cache=False,
+        max_input_length=80,
+        max_retries=3,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["ok", "x" * 80, "after"])
+
+    assert result.embeddings == [[1.0, 0.0], None, [1.0, 0.0]]
+    assert result.terminal_failure is False
+    assert provider.embedded_texts == ["ok", "after"]
+    assert provider.batch_sizes == [3, 1, 1, 1, 1, 1]
+
+
+@pytest.mark.asyncio
+async def test_unrelated_bad_request_is_not_retried_with_a_shorter_input():
+    """Only a 400 that names the input length may be answered with less input."""
+    provider = UnrelatedBadRequestAsEmbedding()
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_unrelated",
+        enable_cache=False,
+        max_input_length=80,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["x" * 80])
+
+    assert result.embeddings is None
+    assert provider.batch_sizes == [1]