From 5b55813585a26dc2f73e11c8e1282753010d6111 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Wed, 12 Aug 2026 18:03:10 -0400 Subject: [PATCH] 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 --- src/embedding_client.py | 14 ++++++++++---- src/reconciler/sync_vectors.py | 9 ++++++++- tests/deriver/test_vector_reconciliation.py | 4 +++- 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/src/embedding_client.py b/src/embedding_client.py index 70f6b34a..1656338d 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -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( diff --git a/src/reconciler/sync_vectors.py b/src/reconciler/sync_vectors.py index 68e9ac59..8f510f47 100644 --- a/src/reconciler/sync_vectors.py +++ b/src/reconciler/sync_vectors.py @@ -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( diff --git a/tests/deriver/test_vector_reconciliation.py b/tests/deriver/test_vector_reconciliation.py index 5e2e4712..b22d71a8 100644 --- a/tests/deriver/test_vector_reconciliation.py +++ b/tests/deriver/test_vector_reconciliation.py @@ -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]