scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-04 07:30:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 10

fix(embedding-store): keep shrinking an over-limit request until it fits

Browse source 
`LocalEmbeddingStore._call_with_retry()` shortened a rejected input once, to
half of the configured `max_input_length`. Two shapes of input-length
rejection still left the chunk unembedded, and an unembedded chunk keeps
`needs_reindex` set, so a strict reindex never completes:

- a shortened request that is still above the provider's real per-item limit
  was never shortened again, even with retry budget left;
- an input already below half the configured budget produced an identical
  "shorter" request, so nothing was re-sent at all.

The retry now derives a strictly smaller request from the text the provider
just rejected and keeps shrinking within an explicit round budget, still
bounded by the remaining retry budget. It stops on success, on an exhausted
budget, or when no smaller non-empty request exists. Source text, chunk
contents, configuration, vector-space checks, cache behaviour and the existing
error routing are untouched; the shrink step no longer sleeps, because an
input-length rejection is deterministic rather than rate limiting.

Refs: #683
This commit is contained in:
BeiMu-new 2026-10-26 22:02:32 +08:00
parent 363e497ae9
commit d17022fc32
2 changed files with 142 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
reme/components/embedding_store/local_embedding_store.py
View file

@ -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] = []

125
tests/unit/test_local_embedding_store.py
View file

@ -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]