scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-14 09:01:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

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

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

`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

* test(embedding-store): cover a repeated rebuild of an over-limit chunk

`reindex` is what #683 was about, and while the request-level recovery had
a repeatability test, the rebuild path did not. Add the rebuild that runs
twice over the same over-limit chunk, with the embedding cache both off and
on, so a regression on the second pass fails here.

The shrinking helper keeps the reasoning behind halving the rejected text,
because that is what makes each round a real reduction.
This commit is contained in:
北沐城歌 2026-10-30 11:19:21 +08:00 • committed by GitHub
parent 3f1447c03d
commit 56e4fe843f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 198 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
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,21 @@ 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, or ``None`` when there is none.
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.
"""
if len(text) <= 1:
return None
return text[: len(text) // 2]
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] = []

55
tests/unit/test_file_store_consistency.py
View file

@ -2256,3 +2256,58 @@ async def test_embedding_reindex_completes_when_a_chunk_exceeds_the_provider_lim
assert long_text not in provider.embedded_texts
await store.close()
@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store])
@pytest.mark.parametrize("long_chars", [5000, 2400])
@pytest.mark.parametrize("enable_cache", [False, True])
@pytest.mark.asyncio
async def test_second_reindex_recovers_the_same_over_limit_chunk(store_factory, long_chars, enable_cache):
"""The repeat rebuild must clear the gate again, not only the first one.
#683 is an availability bug: reindex never finished. The test above covers the first
rebuild; this one raises the gate a second time and requires the same over-limit chunk
to be recovered again, with the embedding cache both off and on.
"""
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
store = store_factory(name=f"t_reindex_twice_{long_chars}_{enable_cache}")
await store.start()
long_text = "x" * long_chars
await set_chunks_with_graph(
store,
{
"good": chunk("good", "a.md", "alpha text"),
"long": chunk("long", "b.md", long_text),
},
)
provider = OverLimitAsEmbedding(max_chars=2048)
embedding = LocalEmbeddingStore(
name=f"t_reindex_twice_embedding_{long_chars}_{enable_cache}",
enable_cache=enable_cache,
max_input_length=8192,
)
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()
shortened = long_text
while len(shortened) > 2048:
shortened = shortened[: len(shortened) // 2]
for round_index in (1, 2):
await store.require_embedding_rebuild()
result = await store.reindex("embedding")
assert result == {"indexed": 2, "scope": "embedding"}
assert store._embedding_rebuild_pending is False
assert provider.embedded_texts.count(shortened) >= 1, f"rebuild {round_index}"
# Without a cache each rebuild has to reach the provider again; with one, the
# second rebuild is served from it and the shortened request is sent only once.
assert provider.embedded_texts.count(shortened) == (1 if enable_cache else 2)
await store.close()

130
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,90 @@ 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_request_recovery_reuses_the_same_shrinking_path():
"""Two consecutive requests on one store both recover the same item.
This is request-level repeatability only: no rebuild is performed and the gate is not
involved. The rebuild-side regression lives in
``test_file_store_consistency.py::test_second_reindex_recovers_the_same_over_limit_chunk``.
"""
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]