diff --git a/reme/components/embedding_store/base_embedding_store.py b/reme/components/embedding_store/base_embedding_store.py index bb7703e2..f499b0eb 100644 --- a/reme/components/embedding_store/base_embedding_store.py +++ b/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] diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 99e59369..14e5cc9c 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/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: diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index 990bfad8..346a507e 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/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() diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 061d078a..ed9d705f 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/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]