diff --git a/docs/0-requirements.ja.md b/docs/0-requirements.ja.md index 789defe..6c67f3b 100644 --- a/docs/0-requirements.ja.md +++ b/docs/0-requirements.ja.md @@ -248,8 +248,8 @@ Responsibilities: - embed 経路では D1 FTS5 upsert 失敗は Vectorize upsert を無効化しない。保存した bodyHash が次回の試行を駆動し、次の reindex で sparse 側が reconcile される - metadata のみの経路（body は変わらず state / labels / milestone / assignees が変わった場合）では、mirror 書き込みの失敗を best-effort 扱いに**しない**。差分検出の基準を進めずに保持し、次の poll / webhook 配信で再試行させる。基準は IssueStore の record そのものなので、失敗した mirror を追い越して基準を進めると取り残しが恒久化する — state だけの変更は、embed 経路が待っている body 変更を二度と連れてこない（issue #209） - この経路の dense / sparse mirror は互いに独立して書く。vector が欠けている行（issue #210）でも sparse 側の state は更新される -- commit diff は 1 commit 分の file リストを batch embed（Workers AI の `text: string[]` 対応を利用）し、1 回の Vectorize upsert で N vector を書き込む -- batch size は `MAX_EMBEDDING_BATCH_SIZE`（既定 20）で上限。これを超える commit は複数 batch call に分割する +- commit diff は 1 commit 分の file リストを batch embed（Workers AI の `text: string[]` 対応を利用）し、batch ごとに 1 回の Vectorize upsert で N vector を書き込む +- batch の切れ目は file 件数ではなく推定 token 予算（`MAX_EMBEDDING_BATCH_TOKENS`）で決める。この pipeline が embed する内容の「1 token あたりの文字数」は桁で変わる — ASCII のソースで約 3、CJK の散文で約 1 — ため、件数が call を縛れるのは「どの patch も安い側である」と仮定した場合だけである。その仮定を破る commit は chunk 全体を失敗させ、vector が載らなかった commit は diff watermark が留まる対象（issue #178）なので、その commit を飛ばすのではなく surface がそこで恒久的に停止していた。しかも毎 cron 同じ commit で同じ結果になる決定論的な失敗で、同じログに同居する一過性の subrequest 超過とはそこで性質が分かれる（issue #236）。天井は bge-m3 の documented な per-input 上限 8192 token ではなく、1 call の input 全体を合算した endpoint 側の上限である。batch は合算され、拒否応答はその合計値を名指しする（`3030: Max context reached 85920 tokens but model supports only 60000`）。この値は非公開なので `WORKERS_AI_BATCH_CONTEXT_LIMIT` としてそのエラーから記録し、予算はその半分に置く。半分にしているのは効く方向の誤差が 1 つだけだからで、`estimateEmbeddingTokens` は近似であり、真値より小さく出た推定こそが call を天井の向こうへ送って stall を再現する。同時に天井から離しすぎない — batch が 1 つ増えるごとに subrequest を 2 つ消費し、この worker は既に invocation 予算を超過しているため。単独で予算を超える input はそのまま単独で送る — さらに削るのは truncate 軸（`MAX_EMBEDDING_INPUT_CHARS`）の仕事であり、落とせばその file が索引から消えるため **索引欠落の修復.** watermark の修正は漏れを止めるだけで、既に空いた穴は埋まらない — 取り残された項目が再 fetch されるのは `updated_at` が動いたときだけで、閉じた履歴はもう動かない。`POST /admin/backfill-issue-index?repo=owner/repo`（installation guide 参照）が欠落そのものを走査する。repository の issue 番号空間は密かつ有界なので、`search_docs` に issue / PR 行が無い番号がそのまま欠落集合であり、数値 cursor が「どこまで走査したか」を厳密に表せる。同じ集合を時刻 cursor で辿ると、欠陥が突いた順序をそのまま持ち込むことになる。GitHub 側に既に無い番号（削除・transfer 済み）は 404 を返すので、retry せず計上のみ。取り込みは body-hash 判定を強制的に飛ばす: 候補はいずれも retrieval surface が欠けていると分かっている項目であり、hash が一致していると（embed 成功後に FTS5 mirror が失敗した行がこの状態になる）そのまま恒久的に skip されてしまうため。state 修復と違いこちらは embed を伴い、1 候補ごとに取り込み fan-out の全額がかかるので、呼び出し側が batch を跨いで sweep を進める。`dry_run=true` は予算を使わずに欠落量だけを測る。sweep は poller の watermark と同じ不変条件に従う: **取り込みそこねた最初の候補を cursor が追い越さない。** これにより 1 call の上限を大きく取りすぎたときの代償は 1 回分の無駄な呼び出しであって取りこぼしではなく、結果が「上限値が subrequest 予算をどれだけ正確に写しているか」に依存しなくなる（issue #216）。GitHub 側に既に無い番号は例外 — 誰が何回試しても取り込めないので、そこで止めると retry の境界にならず sweep が停止する。トレードオフは poller と同じで、恒久的に失敗する候補があると sweep は止まる。ただしそれは応答に現れる（`nextCursor` が渡した `cursor` と同じ値で返る）し、この endpoint は cron ではなく人間 / AI が駆動するので、詰まった番号を手動の `cursor` 指定で越えられる。 diff --git a/docs/0-requirements.md b/docs/0-requirements.md index ce21821..9d7dcb4 100644 --- a/docs/0-requirements.md +++ b/docs/0-requirements.md @@ -253,8 +253,8 @@ Responsibilities: - D1 FTS5 upsert failures do not invalidate a successful Vectorize upsert on the embed path; the stored bodyHash drives the next attempt and the next reindex reconciles the sparse side - on the metadata-only path (state / labels / milestone / assignees changed, body did not) a failed mirror write is **not** best-effort: the diff baseline is held so the next poll or webhook delivery retries. The baseline is the IssueStore record itself, so advancing it past a failed mirror makes the miss permanent — a state-only change never brings the body change the embed path waits for (issue #209) - the dense and sparse mirrors on that path are written independently: a row with no vector (issue #210) still gets its sparse state updated -- for commit diffs: batch-embed a commit's file list in a single Workers AI call (`text: string[]`) and upsert the resulting N vectors in one `VECTORIZE.upsert` call -- batch size is capped by `MAX_EMBEDDING_BATCH_SIZE` (default 20); commits exceeding it are split across multiple batch calls +- for commit diffs: batch-embed a commit's file list through Workers AI (`text: string[]`) and upsert each batch's N vectors in one `VECTORIZE.upsert` call +- batches are cut on an estimated token budget (`MAX_EMBEDDING_BATCH_TOKENS`), not on a file count. Characters per token vary by an order of magnitude across what this pipeline embeds — roughly 3 for ASCII source, roughly 1 for CJK prose — so N inputs bound a call only where every patch is assumed to be the cheap kind. A commit that broke the assumption failed its whole chunk, and a commit whose vectors never landed is one the diff watermark holds on (issue #178), so the surface stalled there permanently instead of passing it by — deterministically, on the same commit every cron tick, which is what separated it from the transient subrequest overruns sharing the log (issue #236). The ceiling is the endpoint's aggregate context across a call's inputs, not bge-m3's documented 8192-token per-input maximum: the batch is summed, and the rejection names the sum (`3030: Max context reached 85920 tokens but model supports only 60000`). It is unpublished, so `WORKERS_AI_BATCH_CONTEXT_LIMIT` records it from that error, and the budget is half of it. The halving covers the one error direction that matters — `estimateEmbeddingTokens` approximates, and an estimate landing under the true count is what puts a call over the ceiling and reproduces the stall — while staying near enough the ceiling to keep the batch count down, since every extra batch spends two subrequests on an invocation budget this worker already overruns. An input whose own estimate exceeds the budget is sent alone — cutting it down further is the truncation axis (`MAX_EMBEDDING_INPUT_CHARS`), and dropping it would lose the file from the index **Missing-entry repair.** The watermark fix stops the leak but does not fill the hole: a stranded item is only re-fetched when its `updated_at` moves, and closed history never moves again. `POST /admin/backfill-issue-index?repo=owner/repo` (see the installation guide) walks the gap directly — the repository's issue-number space is dense and bounded, so the numbers with no `search_docs` issue / PR row are exactly the missing set, and a numeric cursor states how far the sweep has reached. A timestamp cursor over the same set would reintroduce the ordering the defect exploited. Numbers GitHub no longer has (deleted or transferred) answer 404 and are counted rather than retried. The ingest is forced past the body-hash check: every candidate is known to be missing a retrieval surface, and a matching hash — which an embed whose FTS5 mirror failed leaves behind — would otherwise skip it permanently. Unlike the state repair this one embeds, so every candidate carries the full ingest fan-out and the caller drives the sweep one batch at a time; `dry_run=true` measures the gap without spending it. The sweep obeys the same invariant as the poller's watermark: **the cursor never advances past the first candidate a call failed to ingest**, so a per-call limit set too high costs a wasted call rather than a missed item, and the result does not depend on how accurately that limit models the subrequest budget (issue #216). A number GitHub no longer has is exempt — nothing will ever ingest it, so holding there would stall the sweep instead of bounding a retry. The tradeoff is the poller's: a candidate that fails on every attempt stops the sweep. Here that is visible rather than silent (`nextCursor` comes back equal to the `cursor` passed in), and because a human or an AI drives this endpoint rather than cron, stepping over the blocking number is a matter of passing the next `cursor` by hand. diff --git a/src/pipeline/embed-diff.test.ts b/src/pipeline/embed-diff.test.ts new file mode 100644 index 0000000..ef9b7df --- /dev/null +++ b/src/pipeline/embed-diff.test.ts @@ -0,0 +1,152 @@ +import { describe, it, expect, vi } from "vitest"; +import type { Env } from "../types.js"; +import { processAndUpsertCommitDiff, type GitHubCommitDetail } from "./embed-diff.js"; +import { estimateEmbeddingTokens, MAX_EMBEDDING_BATCH_TOKENS } from "./embedding.js"; +import { diffVectorId } from "./vector-id.js"; + +const REPO = "acme/widgets"; + +function mkCommit(patches: string[]): GitHubCommitDetail { + return { + sha: "c0ffee", + commit: { + message: "a commit message", + author: { name: "author", date: "2026-01-01T00:00:00Z" }, + }, + author: { login: "author" }, + files: patches.map((patch, i) => ({ + filename: `src/file-${i}.ts`, + status: "modified", + patch, + sha: `blob${i}`, + })), + }; +} + +interface EnvStubOptions { + /** Batch indices (in call order) whose embed call should throw. */ + embedFailsOnCall?: number[]; +} + +function mkEnv(options: EnvStubOptions = {}) { + /** Inputs handed to each Workers AI call, in call order. */ + const aiCalls: string[][] = []; + const upsertedIds: string[][] = []; + + const env = { + GITHUB_TOKEN: "t", + AI: { + run: vi.fn(async (_model: string, input: { text: string[] }) => { + const callIndex = aiCalls.length; + aiCalls.push(input.text); + if (options.embedFailsOnCall?.includes(callIndex)) { + throw new Error("workers ai rejected the batch"); + } + return { data: input.text.map(() => [0.1, 0.2]) }; + }), + }, + VECTORIZE: { + upsert: vi.fn(async (vectors: Array<{ id: string }>) => { + upsertedIds.push(vectors.map((v) => v.id)); + }), + }, + DB_FTS: { + prepare: () => ({ + bind: () => ({ run: async () => ({}) }), + }), + }, + } as unknown as Env; + + return { env, aiCalls, upsertedIds }; +} + +function mkStore() { + return { + fetch: vi.fn(async () => new Response("{}", { status: 200 })), + } as unknown as DurableObjectStub; +} + +describe("embed-diff: the batch axis is the token budget, not the file count", () => { + it("still sends an ordinary commit as one call", async () => { + // 30 files well past the retired count cap of 20, each a small patch. + const { env, aiCalls } = mkEnv(); + const commit = mkCommit(Array.from({ length: 30 }, (_, i) => `@@ -1 +1 @@\n+line ${i}`)); + + const result = await processAndUpsertCommitDiff(env, mkStore(), REPO, commit); + + expect(result).toEqual({ embedded: 30, skipped: 0, failed: 0, batches: 1 }); + expect(aiCalls).toHaveLength(1); + expect(aiCalls[0]).toHaveLength(30); + }); + + it("splits a commit of large patches that a file count would have kept in one call", async () => { + // Twenty maximal patches — exactly one chunk under the retired count cap, and + // the shape the endpoint rejected in production ("Max context reached 85920 + // tokens but model supports only 60000"), taking every file in the chunk down + // with the call. + const { env, aiCalls, upsertedIds } = mkEnv(); + const commit = mkCommit(Array.from({ length: 20 }, () => "あ".repeat(9000))); + + const result = await processAndUpsertCommitDiff(env, mkStore(), REPO, commit); + + expect(result.embedded).toBe(20); + expect(result.failed).toBe(0); + expect(aiCalls.length).toBeGreaterThan(1); + expect(result.batches).toBe(aiCalls.length); + + // Every call carrying more than one input stays inside the budget. + for (const call of aiCalls) { + if (call.length > 1) { + const total = call.reduce((sum, text) => sum + estimateEmbeddingTokens(text), 0); + expect(total).toBeLessThanOrEqual(MAX_EMBEDDING_BATCH_TOKENS); + } + } + + // No file is dropped or duplicated on the way through the split. + expect(aiCalls.flat()).toHaveLength(20); + expect(new Set(upsertedIds.flat()).size).toBe(20); + }); + + it("keeps each embed call paired with its own slice of files", async () => { + // Vectors are matched to files by position, so a split has to cut the inputs + // and the files on the same boundary. Reading each call's file paths back out + // of its inputs and rebuilding the vector IDs from them catches a slice that + // drifted — the failure mode that would file every patch under a neighbour. + const { env, aiCalls, upsertedIds } = mkEnv(); + const commit = mkCommit(Array.from({ length: 6 }, () => "あ".repeat(9000))); + + await processAndUpsertCommitDiff(env, mkStore(), REPO, commit); + + expect(upsertedIds).toHaveLength(aiCalls.length); + for (let batch = 0; batch < aiCalls.length; batch++) { + // Input format is "{message}\n\n{path}\n\n{patch}". + const paths = aiCalls[batch].map((text) => text.split("\n\n")[1]); + const expected = await Promise.all( + paths.map((path) => diffVectorId(REPO, commit.sha, path)), + ); + expect(upsertedIds[batch]).toEqual(expected); + } + }); + + it("loses only the failing batch when one embed call is rejected", async () => { + const { env, aiCalls, upsertedIds } = mkEnv({ embedFailsOnCall: [0] }); + const commit = mkCommit(Array.from({ length: 6 }, () => "あ".repeat(9000))); + + const result = await processAndUpsertCommitDiff(env, mkStore(), REPO, commit); + + expect(result.failed).toBe(aiCalls[0].length); + expect(result.embedded).toBe(6 - aiCalls[0].length); + expect(upsertedIds.flat()).toHaveLength(result.embedded); + }); + + it("skips files with no patch and reports them separately", async () => { + const { env, aiCalls } = mkEnv(); + const commit = mkCommit(["@@ -1 +1 @@\n+one"]); + commit.files!.push({ filename: "assets/logo.png", status: "modified", sha: "blobX" }); + + const result = await processAndUpsertCommitDiff(env, mkStore(), REPO, commit); + + expect(result).toEqual({ embedded: 1, skipped: 1, failed: 0, batches: 1 }); + expect(aiCalls[0]).toHaveLength(1); + }); +}); diff --git a/src/pipeline/embed-diff.ts b/src/pipeline/embed-diff.ts index dc3bc3c..0b3037a 100644 --- a/src/pipeline/embed-diff.ts +++ b/src/pipeline/embed-diff.ts @@ -12,7 +12,7 @@ import { upsertFtsRow } from "../fts.js"; import { prepareDiffEmbeddingInput } from "./hash.js"; import { generateEmbeddingBatch, - MAX_EMBEDDING_BATCH_SIZE, + planEmbeddingBatches, } from "./embedding.js"; import { diffVectorId } from "./vector-id.js"; @@ -104,7 +104,8 @@ function normaliseFileStatus(status: string): DiffFileStatus { * Flow: * 1. Filter `files[]` to those with a textual `patch` (binary / oversized files are skipped). * 2. Build embedding inputs = commit message + file path + patch, truncated. - * 3. Batch-embed inputs via Workers AI (chunked by MAX_EMBEDDING_BATCH_SIZE). + * 3. Batch-embed inputs via Workers AI (chunked by `planEmbeddingBatches`, which + * splits on an estimated token budget rather than a file count). * 4. Upsert all vectors into Vectorize in the same chunks. * 5. Record DiffRecord rows into the Durable Object store for each indexed file. * @@ -150,21 +151,26 @@ export async function processAndUpsertCommitDiff( let failed = 0; let batches = 0; - // Chunk to respect Workers AI / Vectorize batch limits. - for (let offset = 0; offset < indexable.length; offset += MAX_EMBEDDING_BATCH_SIZE) { - const chunk = indexable.slice(offset, offset + MAX_EMBEDDING_BATCH_SIZE); - batches++; + const allInputs = indexable.map((f) => + prepareDiffEmbeddingInput(commitMessage, f.filename, f.patch), + ); - const inputs = chunk.map((f) => - prepareDiffEmbeddingInput(commitMessage, f.filename, f.patch), - ); + // Chunk on the estimated token total of the inputs, not on how many files the + // commit touched. A file count bounds a call only if every patch is assumed to + // be small, and a commit that breaks that assumption used to fail the whole + // chunk — which the poller reads as an uningested commit and holds the diff + // watermark on, so the surface stalls there rather than skipping past it (#236). + for (const { start, end } of planEmbeddingBatches(allInputs)) { + const chunk = indexable.slice(start, end); + const inputs = allInputs.slice(start, end); + batches++; let embeddings: number[][]; try { embeddings = await generateEmbeddingBatch(env.AI, inputs); } catch (err) { console.error( - `Failed to batch-embed diffs for ${repo}@${commitSha} chunk offset ${offset}:`, + `Failed to batch-embed diffs for ${repo}@${commitSha} chunk offset ${start}:`, err instanceof Error ? err.message : String(err), ); failed += chunk.length; @@ -212,7 +218,7 @@ export async function processAndUpsertCommitDiff( await env.VECTORIZE.upsert(vectors); } catch (err) { console.error( - `Failed to upsert diff vectors for ${repo}@${commitSha} chunk offset ${offset}:`, + `Failed to upsert diff vectors for ${repo}@${commitSha} chunk offset ${start}:`, err instanceof Error ? err.message : String(err), ); failed += chunk.length; diff --git a/src/pipeline/embedding.test.ts b/src/pipeline/embedding.test.ts new file mode 100644 index 0000000..afce34d --- /dev/null +++ b/src/pipeline/embedding.test.ts @@ -0,0 +1,132 @@ +import { describe, it, expect } from "vitest"; +import { + estimateEmbeddingTokens, + planEmbeddingBatches, + MAX_EMBEDDING_BATCH_TOKENS, + MAX_EMBEDDING_INPUT_CHARS, + WORKERS_AI_BATCH_CONTEXT_LIMIT, +} from "./embedding.js"; + +/** Sum of the per-input estimates over one planned range. */ +function rangeTokens(inputs: string[], range: { start: number; end: number }): number { + return inputs + .slice(range.start, range.end) + .reduce((total, text) => total + estimateEmbeddingTokens(text), 0); +} + +describe("estimateEmbeddingTokens", () => { + it("counts ASCII at the assumed characters-per-token ratio", () => { + expect(estimateEmbeddingTokens("x".repeat(300))).toBe(100); + }); + + it("rounds a partial token up rather than down", () => { + expect(estimateEmbeddingTokens("x")).toBe(1); + expect(estimateEmbeddingTokens("xxxx")).toBe(2); + }); + + it("counts non-ASCII one token per code unit — the CJK worst case", () => { + expect(estimateEmbeddingTokens("あ".repeat(100))).toBe(100); + }); + + it("separates the two ranges inside one mixed input", () => { + // 300 ASCII (100 tokens) + 50 CJK (50 tokens). + expect(estimateEmbeddingTokens("x".repeat(300) + "あ".repeat(50))).toBe(150); + }); + + it("is empty for empty input", () => { + expect(estimateEmbeddingTokens("")).toBe(0); + }); + + it("separates content of equal character length by token cost", () => { + // The whole point of the axis: same chars, different budget consumption. + const ascii = "x".repeat(MAX_EMBEDDING_INPUT_CHARS); + const cjk = "あ".repeat(MAX_EMBEDDING_INPUT_CHARS); + expect(estimateEmbeddingTokens(cjk)).toBeGreaterThan(estimateEmbeddingTokens(ascii)); + }); +}); + +describe("MAX_EMBEDDING_BATCH_TOKENS", () => { + it("keeps a margin against the endpoint's aggregate ceiling", () => { + // Asserted rather than left to the comment at the constant, because the + // pressure on this number runs one way: every batch costs two subrequests on + // an axis this worker already overruns, so the temptation is to walk the + // budget up toward the ceiling. The estimator approximates, and an estimate + // that lands under the true count is what reproduces the stall this batching + // exists to prevent. A failing assertion here is not a verdict that the new + // value is wrong — it says the estimator now has to earn the thinner margin. + expect(MAX_EMBEDDING_BATCH_TOKENS).toBeLessThanOrEqual( + WORKERS_AI_BATCH_CONTEXT_LIMIT / 2, + ); + }); + + it("stays above the per-input maximum a truncated input can reach", () => { + // The floor on the same number: below this, a maximal input could not be sent + // even alone, and the planner would be handing the endpoint a call it rejects. + expect(MAX_EMBEDDING_BATCH_TOKENS).toBeGreaterThanOrEqual(MAX_EMBEDDING_INPUT_CHARS); + }); +}); + +describe("planEmbeddingBatches", () => { + it("returns no range for no inputs", () => { + expect(planEmbeddingBatches([])).toEqual([]); + }); + + it("keeps inputs that fit the budget in a single call", () => { + const inputs = Array.from({ length: 50 }, () => "x".repeat(300)); + expect(planEmbeddingBatches(inputs)).toEqual([{ start: 0, end: 50 }]); + }); + + it("splits when the running total would overrun the budget", () => { + // 3 tokens each against a budget of 6: two per call, then the remainder. + const inputs = Array.from({ length: 5 }, () => "x".repeat(9)); + expect(planEmbeddingBatches(inputs, 6)).toEqual([ + { start: 0, end: 2 }, + { start: 2, end: 4 }, + { start: 4, end: 5 }, + ]); + }); + + it("covers every input exactly once, in order", () => { + const inputs = Array.from({ length: 37 }, (_, i) => "x".repeat(i * 30 + 1)); + const ranges = planEmbeddingBatches(inputs, 200); + + expect(ranges[0].start).toBe(0); + expect(ranges[ranges.length - 1].end).toBe(inputs.length); + for (let i = 1; i < ranges.length; i++) { + expect(ranges[i].start).toBe(ranges[i - 1].end); + } + expect(ranges.every((r) => r.end > r.start)).toBe(true); + }); + + it("holds every multi-input call inside the budget", () => { + const inputs = Array.from({ length: 40 }, (_, i) => + i % 3 === 0 ? "あ".repeat(400) : "x".repeat(900), + ); + const ranges = planEmbeddingBatches(inputs, 1000); + + for (const range of ranges) { + if (range.end - range.start > 1) { + expect(rangeTokens(inputs, range)).toBeLessThanOrEqual(1000); + } + } + }); + + it("gives an over-budget input a call of its own instead of stalling", () => { + // The guard that makes this pass is what keeps the planner from closing an + // empty range and looping on the same index forever. + const inputs = ["x".repeat(30), "あ".repeat(5000), "x".repeat(30)]; + expect(planEmbeddingBatches(inputs, 100)).toEqual([ + { start: 0, end: 1 }, + { start: 1, end: 2 }, + { start: 2, end: 3 }, + ]); + }); + + it("admits any single truncated input under the default budget", () => { + // Truncation caps an input at MAX_EMBEDDING_INPUT_CHARS; at the CJK worst case + // that is one token per character, and the default budget must still take it. + const maximal = "あ".repeat(MAX_EMBEDDING_INPUT_CHARS); + expect(estimateEmbeddingTokens(maximal)).toBeLessThanOrEqual(MAX_EMBEDDING_BATCH_TOKENS); + expect(planEmbeddingBatches([maximal])).toEqual([{ start: 0, end: 1 }]); + }); +}); diff --git a/src/pipeline/embedding.ts b/src/pipeline/embedding.ts index 4e87936..0092c88 100644 --- a/src/pipeline/embedding.ts +++ b/src/pipeline/embedding.ts @@ -2,26 +2,121 @@ * Workers AI embedding wrappers and batching constants. * * Wraps the BGE-M3 model behind `generateEmbedding` (single text) and - * `generateEmbeddingBatch` (multi-text). Callers must chunk inputs by - * `MAX_EMBEDDING_BATCH_SIZE`; the helper does not split internally. + * `generateEmbeddingBatch` (multi-text). Callers must chunk inputs with + * `planEmbeddingBatches`; the helper does not split internally. */ /** Maximum characters for embedding input (BGE-M3 context limit ~8192 tokens, conservative char limit) */ export const MAX_EMBEDDING_INPUT_CHARS = 8000; /** - * Maximum number of inputs per Workers AI batch embed call. - * Cloudflare Workers AI does not publish a hard cap on batched embedding inputs, - * so we split large commits into multiple calls. 20 × 8000 chars ≈ 160k chars per - * call keeps payload size comfortably inside observed request limits. + * Aggregate context the Workers AI endpoint accepts across all inputs of one + * batched embed call. Not the per-input maximum bge-m3 documents (8192) — the + * batch is summed, and the endpoint reports the sum it rejected: + * + * 3030: Max context reached 85920 tokens but model supports only 60000 + * + * Unpublished, so this is read off the error rather than a docs page. Kept named + * because the budget below is a margin against it, and a margin whose reference + * is inlined reads as an arbitrary number the next time someone retunes it. */ -export const MAX_EMBEDDING_BATCH_SIZE = 20; +export const WORKERS_AI_BATCH_CONTEXT_LIMIT = 60000; /** - * Maximum number of vectors per single Vectorize.upsert call. - * We mirror MAX_EMBEDDING_BATCH_SIZE so each embed batch maps 1:1 onto one upsert. + * Token budget for the inputs of one batched Workers AI embed call. + * + * Half the ceiling above. The halving is sized to the one error direction that + * matters: `estimateEmbeddingTokens` approximates, and an estimate that comes in + * *under* the true count is what puts a call over the ceiling — which fails the + * whole chunk, and a commit whose vectors never landed is one the diff watermark + * holds on, so the surface stalls there rather than passing it by. Punctuation- + * dense payloads (lockfile hashes, minified sources) are where the ASCII ratio + * below runs optimistic, and 2x covers that class with room left. + * + * The margin is not free and is not larger than it needs to be. Every extra batch + * costs two subrequests (the AI call and its `VECTORIZE.upsert`) against an + * invocation budget this worker already overruns, so a budget far below the + * ceiling buys no safety and spends a neighbouring axis that is genuinely tight. + * + * A count cap cannot express any of this. Characters per token vary by an order of + * magnitude across the content this pipeline embeds — roughly 3 for ASCII source, + * roughly 1 for CJK prose — so N inputs bound the request only when every input + * is assumed to be the cheap kind. */ -export const MAX_VECTORIZE_UPSERT_BATCH_SIZE = 20; +export const MAX_EMBEDDING_BATCH_TOKENS = WORKERS_AI_BATCH_CONTEXT_LIMIT / 2; + +/** + * Characters per token assumed for the ASCII range. bge-m3 tokenizes with an + * XLM-RoBERTa SentencePiece vocabulary, where English prose runs near 4 and + * punctuation-dense source code runs nearer 3. The lower figure is used because + * the diff surface is source code and an underestimate is what overruns a call. + */ +const ASCII_CHARS_PER_TOKEN = 3; + +/** + * Estimate the token cost of one embedding input. + * + * Deliberately an estimate: the tokenizer is not available inside the Worker, and + * the value is only ever compared against a budget that is itself conservative. + * Non-ASCII code units are counted one token each (the CJK worst case), ASCII at + * `ASCII_CHARS_PER_TOKEN`. Counting UTF-16 code units rather than code points + * makes a surrogate pair cost two, which errs toward the safe side. + */ +export function estimateEmbeddingTokens(text: string): number { + let wide = 0; + for (let i = 0; i < text.length; i++) { + if (text.charCodeAt(i) > 127) wide++; + } + const ascii = text.length - wide; + return wide + Math.ceil(ascii / ASCII_CHARS_PER_TOKEN); +} + +/** Half-open `[start, end)` index range over a caller's input array. */ +export interface EmbeddingBatchRange { + start: number; + end: number; +} + +/** + * Split embedding inputs into batches whose estimated token totals stay within + * `budgetTokens`. + * + * Index ranges are returned rather than the strings themselves so the caller can + * slice its own parallel arrays (files, metadata) by the same boundaries — the + * position-for-position correspondence between inputs and returned vectors is + * what the upsert depends on. + * + * Contract: + * - order is preserved, ranges are contiguous, and every input falls in exactly one + * - no returned range is empty + * - an input whose own estimate already exceeds the budget occupies a range of + * one. Cutting it down further belongs to the truncation axis + * (`MAX_EMBEDDING_INPUT_CHARS`), and dropping it would lose a file from the index. + */ +export function planEmbeddingBatches( + inputs: string[], + budgetTokens: number = MAX_EMBEDDING_BATCH_TOKENS, +): EmbeddingBatchRange[] { + const ranges: EmbeddingBatchRange[] = []; + let start = 0; + let total = 0; + + for (let i = 0; i < inputs.length; i++) { + const cost = estimateEmbeddingTokens(inputs[i]); + // Close the open range before an input that would overrun the budget. The + // `i > start` guard is what keeps an oversized input in a batch of its own + // instead of closing an empty range and looping on it forever. + if (i > start && total + cost > budgetTokens) { + ranges.push({ start, end: i }); + start = i; + total = 0; + } + total += cost; + } + + if (inputs.length > start) ranges.push({ start, end: inputs.length }); + return ranges; +} /** * Generate embedding for a text input using Workers AI BGE-M3. @@ -47,9 +142,9 @@ export async function generateEmbedding( * Generate embeddings for multiple text inputs in one batched Workers AI call. * Input order is preserved in the returned array. * - * Workers AI does not publish a hard limit on the number of inputs per call, - * so callers must chunk by MAX_EMBEDDING_BATCH_SIZE before invoking this - * function. Throws if the returned vector count does not match the input count. + * Workers AI does not publish a hard limit on the size of one embed call, so + * callers must chunk with `planEmbeddingBatches` before invoking this function. + * Throws if the returned vector count does not match the input count. */ export async function generateEmbeddingBatch( ai: Ai,