diff --git a/.env.example b/.env.example index ea90402971..de474e33d1 100644 --- a/.env.example +++ b/.env.example @@ -301,6 +301,12 @@ REDIS_URL=redis://redis:6379 # REQUIRED for the self-host review # QDRANT_DIM=768 # vector dimension of the collection (768 = nomic-embed-text:latest; # # 1024 = bge-m3/mxbai-embed-large). Must match AI_EMBED_MODEL; # # recreate the Qdrant collection when changing this after startup. +# AI_EMBED_BATCH=96 # items per RAG embed-provider call (#4327). Defaults to 96 +# # (a conservative bound sized for Workers AI's 100-item cap). +# # Tune upward for GPU-accelerated self-host Ollama throughput -- +# # benchmarking on an RTX A5000 found 96 already near-optimal +# # (~34ms/chunk vs ~70ms/chunk at 32), so leave unset unless you +# # have hardware-specific data suggesting otherwise. # GITTENSORY_REPORTING_SOURCE_DATABASE_URL= # optional Postgres reporting reader URL. Defaults to DATABASE_URL. # GITTENSORY_BACKUP_SOURCE_DATABASE_URL= # optional Postgres backup reader URL. Defaults to DATABASE_URL. # MIGRATIONS_DIR=/app/migrations diff --git a/src/env.d.ts b/src/env.d.ts index 3ebb3f8711..854a6586a2 100644 --- a/src/env.d.ts +++ b/src/env.d.ts @@ -17,6 +17,10 @@ declare global { VECTORIZE?: Vectorize; /** Self-host RAG vector width. Must match the configured embedding model and vector backend. */ QDRANT_DIM?: string; + /** Self-host RAG embed batch size (items per embed-provider call). Defaults to the shipped + * Workers-AI-safe constant (96) when unset — this override exists for self-host operators tuning + * throughput on their own hardware (e.g. GPU-accelerated Ollama), not to change the hosted default. */ + AI_EMBED_BATCH?: string; /** Optional self-host review audit + visual-capture blob store. The Node runtime injects a filesystem-backed * store when REVIEW_AUDIT_DIR is set, or an S3-compatible-bucket-backed store (an operator's own Cloudflare * R2 bucket, or any other S3-compatible provider) when REVIEW_AUDIT_S3_BUCKET + _ENDPOINT + diff --git a/src/review/adapters.ts b/src/review/adapters.ts index 7ac655a1a2..fad8dc5a43 100644 --- a/src/review/adapters.ts +++ b/src/review/adapters.ts @@ -13,7 +13,7 @@ // fail-safe on a missing vector/inference adapter ("no vector index → no RAG", "no AI → no context"), so the // modules NEVER throw — they degrade to no-context. Storage (D1 `DB`) is always present (the Worker cannot run // without it); its wrapper is a thin pass-through with the prepare→bind→all/first/run + batch surface RAG uses. -import { ragDimensionsFromEnv, type InferenceAdapter, type RagInfra, type StorageAdapter, type VectorAdapter } from "./rag"; +import { ragDimensionsFromEnv, ragEmbedBatchFromEnv, type InferenceAdapter, type RagInfra, type StorageAdapter, type VectorAdapter } from "./rag"; // ── Storage (D1 → StorageAdapter). Always present. A thin pass-through over `env.DB` — structurally the // prepare→bind→{all,first,run} + batch surface the ported modules use. Byte-faithful to reviewbot's @@ -70,6 +70,7 @@ export function reviewInferenceAdapter(ai: Ai): InferenceAdapter { export function createReviewAdapters(env: Env): RagInfra { const infra: RagInfra = { storage: reviewStorageAdapter(env) }; if (env.QDRANT_DIM !== undefined) infra.embeddingDimensions = ragDimensionsFromEnv(env.QDRANT_DIM); + if (env.AI_EMBED_BATCH !== undefined) infra.embedBatch = ragEmbedBatchFromEnv(env.AI_EMBED_BATCH); if (env.VECTORIZE) infra.vector = reviewVectorAdapter(env.VECTORIZE); // Embeddings use the DEDICATED embed provider (env.AI_EMBED) when configured — keeping the review chat chain // frontier-only — and fall back to env.AI otherwise (byte-identical to before). diff --git a/src/review/rag.ts b/src/review/rag.ts index e85560813a..0832f62091 100644 --- a/src/review/rag.ts +++ b/src/review/rag.ts @@ -70,6 +70,8 @@ export interface RagInfra { vector?: VectorAdapter; inference?: InferenceAdapter; embeddingDimensions?: number; + /** Items per embed-provider call. Defaults to `EMBED_BATCH` when unset — see `ragEmbedBatchFromEnv`. */ + embedBatch?: number; } export type RagRetrievalMetrics = { @@ -119,6 +121,11 @@ export function ragDimensionsFromEnv(value: string | undefined): number { return Number.isFinite(dim) && dim > 0 ? Math.floor(dim) : RAG_DIMENSIONS; } +export function ragEmbedBatchFromEnv(value: string | undefined): number { + const batch = Number(value); + return Number.isFinite(batch) && batch > 0 ? Math.floor(batch) : EMBED_BATCH; +} + // ── Filtering: index CODE, not content/data corpora (the primary free-tier cost guard) ─────────── const SKIP_DIR_RE = /(^|\/)(node_modules|dist|build|out|coverage|vendor|\.git|\.next|\.nuxt|\.svelte-kit|\.turbo|\.cache|target|\.gradle|_build|\.venv|venv|__pycache__|\.mypy_cache|\.pytest_cache|\.ruff_cache|\.tox|\.terraform|content|data|fixtures|__snapshots__|__fixtures__|testdata|generated|public)\//i; @@ -284,12 +291,17 @@ export async function countRepoChunks(storage: StorageAdapter, project: string, } // ── Embedding (fail-safe: null on any failure) ──────────────────────────────────────────────────── -export async function embedTexts(inference: InferenceAdapter | undefined, texts: string[], expectedDimensions = RAG_DIMENSIONS): Promise { +export async function embedTexts( + inference: InferenceAdapter | undefined, + texts: string[], + expectedDimensions = RAG_DIMENSIONS, + batchSize = EMBED_BATCH, +): Promise { if (!inference || texts.length === 0) return null; try { const out: number[][] = []; - for (let i = 0; i < texts.length; i += EMBED_BATCH) { - const batch = texts.slice(i, i + EMBED_BATCH); + for (let i = 0; i < texts.length; i += batchSize) { + const batch = texts.slice(i, i + batchSize); const res = (await inference.run(EMBED_MODEL, { text: batch })) as { data?: number[][] } | null; const data = res?.data; // Validate COUNT and DIMENSION: a self-host embedding endpoint can return a structurally-valid response @@ -316,7 +328,7 @@ export async function upsertChunks(infra: RagInfra, project: string, repo: strin const { storage: db, vector: vec, inference } = infra; if (!vec || !inference || chunks.length === 0) return 0; const namespace = ragNamespace(project, repo); - const vectors = await embedTexts(inference, chunks.map((c) => c.text), infra.embeddingDimensions ?? RAG_DIMENSIONS); + const vectors = await embedTexts(inference, chunks.map((c) => c.text), infra.embeddingDimensions ?? RAG_DIMENSIONS, infra.embedBatch ?? EMBED_BATCH); if (!vectors) return 0; try { await vec.upsert( @@ -420,7 +432,7 @@ export async function retrieveContextWithMetrics( // entirely — no point spending an inference call (and vector query budget) on an empty namespace. (#audit cost) if (!(await hasIndexedChunks(storage, opts.project, opts.repo, Date.now()))) return emptyRagRetrievalResult(configuredMinScore); try { - const embedded = await embedTexts(inference, [opts.queryText.slice(0, 16000)], infra.embeddingDimensions ?? RAG_DIMENSIONS); + const embedded = await embedTexts(inference, [opts.queryText.slice(0, 16000)], infra.embeddingDimensions ?? RAG_DIMENSIONS, infra.embedBatch ?? EMBED_BATCH); const vec = embedded?.[0]; if (!vec) return emptyRagRetrievalResult(configuredMinScore); const res = await vectorAdapter.query(vec, { diff --git a/test/unit/rag.test.ts b/test/unit/rag.test.ts index 3114cef0d3..7b6e87ffe3 100644 --- a/test/unit/rag.test.ts +++ b/test/unit/rag.test.ts @@ -16,6 +16,7 @@ import { type RagInfra, RAG_DIMENSIONS, ragDimensionsFromEnv, + ragEmbedBatchFromEnv, ragNamespace, readChunkTexts, retrieveContext, @@ -157,6 +158,21 @@ describe("ragDimensionsFromEnv", () => { }); }); +describe("ragEmbedBatchFromEnv", () => { + it("uses a positive integer batch size from configuration", () => { + expect(ragEmbedBatchFromEnv("32")).toBe(32); + expect(ragEmbedBatchFromEnv("256.9")).toBe(256); + }); + + it("falls back to the shipped EMBED_BATCH default (96) for unset, invalid, or non-positive values", () => { + expect(ragEmbedBatchFromEnv(undefined)).toBe(96); + expect(ragEmbedBatchFromEnv("")).toBe(96); + expect(ragEmbedBatchFromEnv("not-a-number")).toBe(96); + expect(ragEmbedBatchFromEnv("0")).toBe(96); + expect(ragEmbedBatchFromEnv("-5")).toBe(96); + }); +}); + describe("rag: per-file chunking", () => { it("emits one chunk for a small file", () => { const chunks = chunkFile("src/a.ts", "export const x = 1;\n"); @@ -344,6 +360,27 @@ describe("rag: fail-safe (never throws; degrades to no context)", () => { expect(out).not.toContain("src/changed.ts"); // the file under review is excluded → only RELATED code surfaces }); + it("threads a configured infra.embedBatch into the query-embed call too (#4327)", async () => { + const matches = [{ id: "src/a.ts::0", score: 0.9, metadata: { path: "src/a.ts" } }]; + const vector = { query: async () => ({ matches }) } as unknown as VectorAdapter; + let queryEmbedBatch = 0; + const inference: InferenceAdapter = { + run: async (_model, options) => { + queryEmbedBatch = (options as { text: string[] }).text.length; + return { data: [Array(1024).fill(0.1)] }; + }, + }; + const infra: RagInfra = { + storage: storageStub({ count: 1, rows: [{ id: "src/a.ts::0", text: "export const x = 1;" }] }), + vector, + inference, + embedBatch: 8, // arbitrary non-default value; only one query text is ever embedded, so this proves plumbing, not chunking + }; + const out = await retrieveContext(infra, { project: "p", repo: "o/r", queryText: "refactor the auth token verification and add coverage" }); + expect(out).toContain("src/a.ts"); + expect(queryEmbedBatch).toBe(1); // a single query string, regardless of the configured batch size + }); + it("retrieveContextWithMetrics reports candidates, injected chars, and unique retrieved paths", async () => { const matches = [ { id: "src/a.ts::0", score: 0.9, metadata: { path: "src/a.ts" } }, @@ -454,6 +491,27 @@ describe("rag: upsertChunks (embed + vector upsert + chunk-text store)", () => { expect((upserted[0]?.[0]?.values ?? []).length).toBe(768); }); + it("threads a configured infra.embedBatch into the embed call (self-host GPU tuning, #4327)", async () => { + const upserted: VectorUpsert[][] = []; + const vector = { upsert: async (v: VectorUpsert[]) => { upserted.push(v); } } as unknown as VectorAdapter; + const storage = { + prepare: () => ({ bind: () => ({ run: async () => undefined }) as unknown as BoundStatement }), + batch: async () => undefined, + } as unknown as StorageAdapter; + const batchSizes: number[] = []; + const inference: InferenceAdapter = { + run: async (_model, options) => { + const batch = (options as { text: string[] }).text; + batchSizes.push(batch.length); + return { data: batch.map(() => Array(1024).fill(0.1)) }; + }, + }; + const manyChunks: RagChunk[] = Array.from({ length: 10 }, (_, i) => ({ id: `ns|src/a.ts::${i}`, path: "src/a.ts", chunkIndex: i, kind: "code", text: `chunk ${i}` })); + const n = await upsertChunks({ storage, vector, inference, embedBatch: 4 }, "gittensory", "o/r", manyChunks); + expect(n).toBe(10); + expect(batchSizes).toEqual([4, 4, 2]); // proves infra.embedBatch (4) drove the batching, not the 96 default + }); + it("returns 0 with no vector / no inference / empty chunks (the fail-safe guard)", async () => { const vector = { upsert: async () => undefined } as unknown as VectorAdapter; const storage = storageStub(); @@ -675,6 +733,21 @@ describe("rag: embedTexts validation branches", () => { expect(calls).toEqual([96, 54]); // proves the batching loop ran twice }); + it("honors a configured batchSize override instead of the EMBED_BATCH=96 default (self-host GPU tuning)", async () => { + const calls: number[] = []; + const inference: InferenceAdapter = { + run: async (_model, options) => { + const batch = (options as { text: string[] }).text; + calls.push(batch.length); + return { data: batch.map(() => Array(1024).fill(0.1)) }; + }, + }; + const texts = Array.from({ length: 150 }, (_, i) => `t${i}`); + const out = await embedTexts(inference, texts, RAG_DIMENSIONS, 50); + expect(out?.length).toBe(150); + expect(calls).toEqual([50, 50, 50]); // three batches of 50, not the default 96/54 split + }); + it("fails the WHOLE embed when a LATER batch is malformed (early-return mid-loop)", async () => { let call = 0; const inference: InferenceAdapter = { diff --git a/test/unit/review-adapters.test.ts b/test/unit/review-adapters.test.ts index 8e92e664f3..112552562b 100644 --- a/test/unit/review-adapters.test.ts +++ b/test/unit/review-adapters.test.ts @@ -68,6 +68,12 @@ describe("createReviewAdapters: bundle assembly + graceful degradation", () => { expect(infra.embeddingDimensions).toBe(768); }); + it("carries the configured embed batch size into the infra bundle (self-host GPU tuning, #4327)", () => { + const { DB } = dbStub(); + const infra = createReviewAdapters({ DB, VECTORIZE: vectorizeStub(), AI: aiStub(), AI_EMBED_BATCH: "32" } as unknown as Env); + expect(infra.embedBatch).toBe(32); + }); + it("prefers the dedicated AI_EMBED provider for inference, keeping the review chain frontier-only", async () => { const { DB } = dbStub(); const reviewAi = { run: vi.fn(async () => ({ response: "review text" })) }; // would NOT return embed data @@ -112,6 +118,7 @@ describe("createReviewAdapters: bundle assembly + graceful degradation", () => { expect("vector" in infra).toBe(false); expect("inference" in infra).toBe(false); expect("embeddingDimensions" in infra).toBe(false); + expect("embedBatch" in infra).toBe(false); }); });