fix(embedding): guarantee truncation progress and truncate on re-embed

The retry slice in _truncate_to_token_limit always recomputed the same
keep count, so a slice whose re-encode grew past the cap could oscillate.
Decrement keep after each unsuccessful retry.

Document re-embed in the reconciler used the default on_oversize="raise",
so one oversize document failed every other document in the batch.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Aakash Kattelu 2026-08-12 18:03:10 -04:00
parent 144ea2a07b
commit 5b55813585
3 changed files with 21 additions and 6 deletions

View File

@ -313,17 +313,23 @@ class _EmbeddingClient:
"""Return a prefix of `text` whose re-encoded token count fits the cap.
`decode(ids[:n])` can re-encode to more than `n` tokens at BPE /
pre-tokenizer boundaries, so we re-encode after each slice and trim
again until the verified count fits. If a slice does not shrink the
count we drop one extra token so the loop cannot spin.
pre-tokenizer boundaries (a slice landing mid-character decodes to
replacement chars that cost tokens of their own), so we re-encode
after each slice and trim again until the verified count fits. Each
retry keeps strictly fewer tokens than the last, so the loop always
terminates instead of oscillating around the cap.
"""
token_ids = self.encoding.encode(text)
keep = self.max_embedding_tokens
while len(token_ids) > self.max_embedding_tokens:
keep = min(self.max_embedding_tokens, len(token_ids) - 1)
keep = min(keep, len(token_ids) - 1)
if keep < 1:
return "", 0
text = self.encoding.decode(token_ids[:keep])
token_ids = self.encoding.encode(text)
# Guarantee forward progress: if this slice still re-encodes over
# the cap, the next one must be strictly smaller.
keep -= 1
return text, len(token_ids)
async def simple_batch_embed(

View File

@ -301,7 +301,14 @@ async def _sync_documents(
EmbeddingCallPurpose.VECTOR_SYNC.value,
parent_category="reconciliation",
):
new_embeddings = await embedding_client.simple_batch_embed(contents)
# Document content is stored whole, so an observation can be
# over the per-item embedding cap (the write path truncates the
# embedded prefix, not the stored text). Truncate here too:
# raising would drop every other document in this batch on
# every reconciler pass.
new_embeddings = await embedding_client.simple_batch_embed(
contents, on_oversize="truncate"
)
if len(new_embeddings) != len(docs_needing_embed):
logger.warning(

View File

@ -551,7 +551,9 @@ class TestReEmbedding:
# Mock embedding client to track batch calls
batch_call_count = 0
async def track_batch_embed(contents: list[str]) -> list[list[float]]:
async def track_batch_embed(
contents: list[str], **_kwargs: object
) -> list[list[float]]:
nonlocal batch_call_count
batch_call_count += 1
return [[1.0] * 1536 for _ in contents]