fix(kb): batch Voyage requests by token budget, not just request count
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) <noreply@anthropic.com>
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user