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]