Merge d17022fc32 into 2767836b21

reme/components/embedding_store/local_embedding_store.py

@@ -15,6 +15,10 @@ from ..as_embedding import BaseAsEmbedding
 Miss = tuple[int, str, str]  # (result_index, text, cache_key)
 _MAX_VECTOR_SPACE_ATTEMPTS = 3
 
+# Each input-length rejection may be retried with a strictly smaller request, up to this
+# many rounds (also bounded by the remaining retry budget).
+_MAX_INPUT_SHRINK_ROUNDS = 3
+
 
 # Match SDK/httpx transport exceptions without importing optional providers.
 _TRANSPORT_FAILURE_TYPES = frozenset(
@@ -206,7 +210,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
+        shrink_rounds = 0
         for attempt in range(self.max_retries):
             try:
                 embeddings = await self.as_embedding(texts, **kwargs)
@@ -248,18 +252,22 @@ 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]:
+                    # Retry the item with a strictly smaller request derived from the text the
+                    # provider just rejected. A remaining attempt is required so exhausted input
+                    # failures stay non-terminal, and the shrinking rounds are bounded.
+                    can_shrink = (
+                        len(texts) == 1 and attempt < self.max_retries - 1 and shrink_rounds < _MAX_INPUT_SHRINK_ROUNDS
+                    )
+                    if can_shrink:
+                        shorter = self._shrink_rejected_input(texts[0])
+                        if shorter is not None:
+                            shrink_rounds += 1
                             self.logger.warning(
-                                f"Embedding input rejected as too long; retrying with a smaller budget "
-                                f"(chars={len(texts[0])} -> {len(shorter)})",
+                                f"Embedding input rejected as too long; retrying with a smaller request "
+                                f"(chars={len(texts[0])} -> {len(shorter)}, "
+                                f"round {shrink_rounds}/{_MAX_INPUT_SHRINK_ROUNDS})",
                             )
                             texts = [shorter]
-                            shortened = True
-                            await asyncio.sleep(2**attempt)
                             continue
                     self.logger.exception("Embedding request failed: input-dependent error")
                     self.is_healthy = False
@@ -271,6 +279,25 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
         self.is_healthy = False
         return _EmbeddingResult(terminal_failure=True)
 
+    @staticmethod
+    def _shrink_rejected_input(text: str) -> str | None:
+        """Return a strictly smaller request built from the text the provider rejected.
+
+        The rejected text is the only reliable starting point: a budget derived from the
+        configured ``max_input_length`` can still sit above the provider's real per-item
+        limit, or leave an already short input untouched so nothing is re-sent. Halving the
+        *current* request keeps every round an actual reduction, and the caller bounds the
+        number of rounds so an impossible limit cannot spin.
+
+        Returns ``None`` when no smaller non-empty request can be produced.
+        """
+        if len(text) <= 1:
+            return None
+        target = max(1, len(text) // 2)
+        if target >= len(text):
+            return None
+        return text[:target]
+
     async def _embed_individually(self, texts: list[str], **kwargs) -> _EmbeddingResult:
         """Keep healthy siblings, stopping when a failure is no longer specific to an input."""
         embeddings: list[list[float] | None] = []

tests/unit/test_local_embedding_store.py

@@ -993,13 +993,13 @@ async def test_batch_failure_falls_back_to_per_item_embedding():
 
     results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings
 
-    assert results == [[1.0, 0.0], None, [1.0, 0.0]]
-    assert provider.batch_sizes == [3, 1, 1, 1]
+    assert results == [[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]]
+    assert provider.batch_sizes == [3, 1, 1, 1, 1]
 
 
 @pytest.mark.asyncio
 async def test_poison_chunk_keeps_its_batch_mates():
-    """A rejected chunk must not clear the vectors of the chunks beside it."""
+    """An over-limit chunk is recovered from a shorter request without touching its siblings."""
     store = LocalEmbeddingStore(name="t_local_embedding_poison_nodes", enable_cache=False)
     store.as_embedding = PerItemLimitAsEmbedding()
     nodes = [EmbNode(text="first"), EmbNode(text="y" * 64), EmbNode(text="third")]
@@ -1007,17 +1007,19 @@ async def test_poison_chunk_keeps_its_batch_mates():
     await store.get_node_embeddings(nodes)
 
     assert nodes[0].embedding is not None
-    assert nodes[1].embedding is None
+    assert nodes[1].embedding is not None
     assert nodes[2].embedding is not None
 
 
 @pytest.mark.asyncio
 async def test_single_item_failure_still_reports_a_failed_batch():
-    """A one-item request has nothing to isolate, so the existing contract stands."""
+    """A one-item request that cannot be shortened keeps the non-terminal contract."""
     store = LocalEmbeddingStore(name="t_local_embedding_single_poison", enable_cache=False)
-    store.as_embedding = PerItemLimitAsEmbedding()
+    provider = PerItemLimitAsEmbedding()
+    provider.max_chars = 0  # no prefix can satisfy this limit
+    store.as_embedding = provider
 
-    result = await store._call_with_retry(["x" * 64])
+    result = await store._call_with_retry(["x"])
     assert result.embeddings is None
     assert result.terminal_failure is False
 
@@ -1107,8 +1109,8 @@ async def test_per_item_fallback_keeps_a_rate_limited_vector(monkeypatch):
 
     results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings
 
-    assert results == [[1.0, 0.0], None, [1.0, 0.0]]
-    assert provider.batch_sizes == [3, 1, 1, 1, 1]
+    assert results == [[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]]
+    assert provider.batch_sizes == [3, 1, 1, 1, 1, 1]
     assert provider.transient_failures == 1
     assert sleeps == [1.0]
 
@@ -1174,8 +1176,8 @@ async def test_a_provider_without_a_retry_policy_keeps_the_store_backoff(monkeyp
 
         results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings
 
-        assert results == [[1.0, 0.0], None, [1.0, 0.0]]
-        assert provider.batch_sizes == [3, 1, 1, 1, 1]
+        assert results == [[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]]
+        assert provider.batch_sizes == [3, 1, 1, 1, 1, 1]
         assert sleeps == [1.0]
 
 
@@ -1260,9 +1262,9 @@ async def test_an_isolated_item_still_recovers_through_the_sdk(monkeypatch):
 
     results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings
 
-    assert results == [[0.0, 1.0], None, [0.0, 1.0]]
-    assert provider.calls == 4
-    assert endpoint.requests == 5
+    assert results == [[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]]
+    assert provider.calls == 5
+    assert endpoint.requests == 6
 
 
 @pytest.mark.asyncio
@@ -1463,9 +1465,9 @@ async def test_isolation_keeps_going_while_only_inputs_fail():
 
     result = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"])
 
-    assert result.embeddings == [[1.0, 0.0], None, [1.0, 0.0], [1.0, 0.0]]
+    assert result.embeddings == [[1.0, 0.0], [1.0, 0.0], [1.0, 0.0], [1.0, 0.0]]
     assert result.terminal_failure is False
-    assert provider.batch_sizes == [4, 1, 1, 1, 1]
+    assert provider.batch_sizes == [4, 1, 1, 1, 1, 1]
 
 
 @pytest.mark.asyncio
@@ -1557,8 +1559,8 @@ async def test_context_length_rejection_reported_as_http_413_is_isolated():
 
     results = (await store._call_with_retry(["ok", "x" * 64, "also ok"])).embeddings
 
-    assert results == [[1.0, 0.0], None, [1.0, 0.0]]
-    assert provider.batch_sizes == [3, 1, 1, 1]
+    assert results == [[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]]
+    assert provider.batch_sizes == [3, 1, 1, 1, 1]
 
 
 @pytest.mark.asyncio
@@ -1906,16 +1908,17 @@ async def test_batch_rejected_for_one_item_keeps_every_vector():
 
 @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."""
+    """A one-character input cannot be shrunk, so no second request is sent."""
     provider = ShrinkablePerItemLimitAsEmbedding(max_chars=8)
     store = LocalEmbeddingStore(
         name="t_local_embedding_shrink_noop",
         enable_cache=False,
         max_input_length=8192,
     )
+    provider.max_chars = 0
     store.as_embedding = provider
 
-    result = await store._call_with_retry(["x" * 64])
+    result = await store._call_with_retry(["x"])
 
     assert result.embeddings is None
     assert provider.batch_sizes == [1]
@@ -2000,3 +2003,85 @@ async def test_unrelated_bad_request_is_not_retried_with_a_shorter_input():
 
     assert result.embeddings is None
     assert provider.batch_sizes == [1]
+
+
+@pytest.mark.asyncio
+async def test_shortened_request_that_is_still_too_long_shrinks_again():
+    """A shorter request that the provider still rejects must be shortened once more.
+
+    Regression for the first reported case: the half-budget retry landed above the
+    provider's real per-item limit and the store stopped after a single attempt.
+    """
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=2048)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_twice",
+        enable_cache=False,
+        max_input_length=8192,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["x" * 5000])
+
+    assert result.embeddings == [[1.0, 0.0]]
+    assert result.terminal_failure is False
+    assert provider.embedded_texts == ["x" * 1250]
+    assert provider.batch_sizes == [1, 1, 1]
+
+
+@pytest.mark.asyncio
+async def test_input_below_the_half_budget_is_still_shortened():
+    """An input already below half the configured budget still gets its own smaller request.
+
+    Regression for the second reported case: ``_truncate(text, max_length=max_input_length
+    // 2)`` returned the text unchanged, so nothing was re-sent even though the provider
+    reported an input-length error.
+    """
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=2048)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_below_half",
+        enable_cache=False,
+        max_input_length=8192,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["x" * 2400])
+
+    assert result.embeddings == [[1.0, 0.0]]
+    assert provider.batch_sizes == [1, 1]
+    assert provider.embedded_texts == ["x" * 1200]
+
+
+@pytest.mark.asyncio
+async def test_shrinking_stops_when_no_smaller_request_can_fit():
+    """The shrinking budget is bounded: a limit no prefix can satisfy must not spin."""
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=0)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_bounded_rounds",
+        enable_cache=False,
+        max_input_length=8192,
+    )
+    store.as_embedding = provider
+
+    result = await store._call_with_retry(["x" * 400])
+
+    assert result.embeddings is None
+    assert len(provider.batch_sizes) == store.max_retries
+
+
+@pytest.mark.asyncio
+async def test_repeated_recovery_reuses_the_same_shrinking_path():
+    """A second rebuild must recover the same item without any new configuration."""
+    provider = ShrinkablePerItemLimitAsEmbedding(max_chars=2048)
+    store = LocalEmbeddingStore(
+        name="t_local_embedding_shrink_repeat",
+        enable_cache=False,
+        max_input_length=8192,
+    )
+    store.as_embedding = provider
+
+    first = await store._call_with_retry(["x" * 5000])
+    second = await store._call_with_retry(["x" * 5000])
+
+    assert first.embeddings == [[1.0, 0.0]]
+    assert second.embeddings == [[1.0, 0.0]]
+    assert provider.embedded_texts == ["x" * 1250, "x" * 1250]