From 3948c387b5d001664a0f359cfab848ef54087f2a Mon Sep 17 00:00:00 2001 From: Chaim Date: Tue, 21 Apr 2026 15:32:12 +0000 Subject: [PATCH] fix(kb): batch Voyage requests by token budget, not just request count MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous batching sent up to 128 inputs per request but didn't cap total tokens. On large Hebrew statutes (ספר הליקויים, ספר המבחנים) a single request exceeded Voyage's 120K-token-per-batch hard cap and failed with 400. - voyage.embed now pre-computes an estimated token count per input (Hebrew is ~0.55 tokens/char in voyage-multilingual-2) and emits multiple requests so each stays under a 90K-token safety margin. - chunker lowers the per-chunk ceiling to 1500 tokens and updates the chars→tokens estimate to match the observed tokenizer behavior. Refs Task Master #2 Co-Authored-By: Claude Opus 4.7 (1M context) --- api/services/kb/chunker.py | 8 +++++--- api/services/kb/voyage.py | 31 ++++++++++++++++++++++++------- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/api/services/kb/chunker.py b/api/services/kb/chunker.py index 398c8a8..cd3631f 100644 --- a/api/services/kb/chunker.py +++ b/api/services/kb/chunker.py @@ -11,11 +11,13 @@ from dataclasses import dataclass from typing import Literal # Approx tokens per char for Hebrew/mixed text (voyage-multilingual-2 tokenizer). -_TOKENS_PER_CHAR = 0.35 +# Voyage tends to split Hebrew into 1.5-2 tokens per word; 0.55 tokens/char +# is a conservative estimate that matches what the API actually reports. +_TOKENS_PER_CHAR = 0.55 # Hard safety ceiling per chunk. Voyage accepts up to 32K tokens per item, -# but a single batch is capped at 120K tokens. Keep chunks well below both. -_MAX_TOKENS_PER_CHUNK = 2000 +# and a batch is capped at 120K tokens. Keep chunks small so batches fit. +_MAX_TOKENS_PER_CHUNK = 1500 # A "section" in Israeli statutes: "סעיף 12", "תקנה 5", "5א.", "12(ב)" _SECTION_RE = re.compile( diff --git a/api/services/kb/voyage.py b/api/services/kb/voyage.py index e5dd3db..c492e86 100644 --- a/api/services/kb/voyage.py +++ b/api/services/kb/voyage.py @@ -12,6 +12,12 @@ logger = logging.getLogger("shira.kb.voyage") _VOYAGE_URL = "https://api.voyageai.com/v1/embeddings" _MAX_BATCH = 128 +# Voyage caps each request at 120K tokens total. Keep a safety margin so our +# char-based estimate doesn't push over when the real tokenizer disagrees. +_MAX_BATCH_TOKENS = 90_000 +# Rough chars→tokens for Hebrew/mixed text, used only to pre-batch requests. +# Voyage's multilingual tokenizer splits Hebrew more aggressively than ASCII. +_TOKENS_PER_CHAR_EST = 0.55 class VoyageError(RuntimeError): @@ -32,10 +38,23 @@ async def embed( model = model or os.environ.get("VOYAGE_MODEL", "voyage-multilingual-2") + # Pre-batch by estimated token budget so we never exceed Voyage's 120K cap. + batches: list[list[str]] = [] + current: list[str] = [] + current_tokens = 0 + for t in texts: + est = max(1, int(len(t) * _TOKENS_PER_CHAR_EST)) + if current and (current_tokens + est > _MAX_BATCH_TOKENS or len(current) >= _MAX_BATCH): + batches.append(current) + current, current_tokens = [], 0 + current.append(t) + current_tokens += est + if current: + batches.append(current) + out: list[list[float]] = [] - async with httpx.AsyncClient(timeout=60) as http: - for start in range(0, len(texts), _MAX_BATCH): - batch = texts[start : start + _MAX_BATCH] + async with httpx.AsyncClient(timeout=120) as http: + for bi, batch in enumerate(batches): resp = await http.post( _VOYAGE_URL, headers={"Authorization": f"Bearer {api_key}"}, @@ -47,9 +66,7 @@ async def embed( for item in sorted(data["data"], key=lambda d: d["index"]): out.append(item["embedding"]) logger.info( - "[kb.voyage] embedded batch=%d/%d type=%s", - min(start + _MAX_BATCH, len(texts)), - len(texts), - input_type, + "[kb.voyage] batch %d/%d size=%d type=%s", + bi + 1, len(batches), len(batch), input_type, ) return out