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:
parent
144ea2a07b
commit
5b55813585
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Reference in New Issue