From 0249bbf98dd028989409a3a82c29f8ae25da1239 Mon Sep 17 00:00:00 2001 From: Dmitri Khokhlov Date: Sun, 16 Aug 2026 21:29:38 -0700 Subject: [PATCH 1/3] feat(embeddings): batch ollama embeds via /api/embed Send multiple chunks per POST /api/embed request (input: string[] -> embeddings: number[][]) to amortize N per-chunk HTTP round-trips into one. ~10x per-chunk speedup measured against a remote ollama host (236 ms/chunk -> ~22 ms/chunk for 32 inputs). - OllamaEmbeddingProvider: new embedMany (batched /api/embed) and embedOneByOne (shared legacy per-text /api/embeddings path). embedBatch: 0 -> empty, 1 -> legacy (no /api/embed probe), >1 -> batched with fallback to per-text on context-length / 404 / malformed-batch errors so one bad vector fails only its chunk. Malformed/non-JSON/null 200 bodies engage the fallback. - getDynamicBatchOptions: ollama defaults maxBatchItems 16, maxBatchTokens 65536; new embedding.batch.{maxBatchItems,maxBatchTokens} config knobs override them. Number.isFinite guard rejects NaN/Infinity. - Recovery path (index() retry loop and retryFailedBatches) re-embeds previously-failed chunks one per request (ollama-scoped) so a permanently-failing chunk is isolated instead of re-poisoning its co-batched healthy chunks. - docs/configuration.md: new "Batching Ollama embeddings" section and a resource-pressure note (ollama concurrency is fixed at 5; ~80 chunks in flight at the default 16). Backward compatible: old ollama without /api/embed (404) and context-length overflow fall back to the unchanged per-text path. --- docs/configuration.md | 46 ++++++ src/config/schema.ts | 39 +++++ src/embeddings/providers/ollama.ts | 131 ++++++++++++++++- src/indexer/index.ts | 50 +++++-- tests/config.test.ts | 41 ++++++ tests/custom-provider.test.ts | 24 ++-- tests/indexer-failed-batches.test.ts | 74 ++++++++++ tests/ollama-embed.test.ts | 204 +++++++++++++++++++++++++++ 8 files changed, 578 insertions(+), 31 deletions(-) create mode 100644 tests/ollama-embed.test.ts diff --git a/docs/configuration.md b/docs/configuration.md index 9989bcb1..9fb1ecaf 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -93,6 +93,52 @@ Keep the default when the additional local model download and query latency are not worthwhile for your project. See the [local model comparison](benchmarks/2026-08-12-local-embedding-model-comparison.md) for the methodology and measured results. +#### Batching Ollama embeddings + +The indexer sends multiple chunks to Ollama in one request. This decreases the +number of HTTP requests and accelerates indexing against a remote Ollama host. +The `/api/embed` endpoint accepts an array of texts and returns one vector per +text. The indexer uses this endpoint for batches of two or more chunks. A +single-chunk batch uses the legacy `/api/embeddings` endpoint. + +If the batched endpoint is not available, the indexer falls back to the legacy +per-chunk endpoint. If one chunk exceeds the model context length, the indexer +truncates and retries that chunk by itself. If the batch response is malformed, +the indexer retries each chunk by itself so one bad vector fails only its chunk. + +Control the batch size with the `embedding.batch` section: + +```json +{ + "embedding": { + "batch": { + "maxBatchItems": 32, + "maxBatchTokens": 65536 + } + } +} +``` + +| Option | Ollama default | Purpose | +|---|---:|---| +| `maxBatchItems` | `16` | Maximum number of chunks in one request | +| `maxBatchTokens` | `65536` | Maximum total estimated tokens in one request | + +Ollama encodes each text independently. The model context length applies to each +text and not to the batch total. Set `maxBatchTokens` to bound the request size and +the processing time. Set `maxBatchItems` to bound the number of texts. Both values +are optional and must be at least 1. When you omit a value, the indexer uses the +Ollama default. These knobs apply to any provider when you set them; the Ollama +defaults apply only to the Ollama provider. + +The indexer runs up to five Ollama requests at the same time. Each request carries +up to `maxBatchItems` chunks, so the worst case is five times `maxBatchItems` chunks +in flight (80 chunks at the default 16). A remote or memory-limited Ollama host can +run out of memory or exceed the 120-second request timeout when this number is too +high. Lower `maxBatchItems` for a small or remote host. The Ollama concurrency is +fixed at five and is not configurable; the custom provider exposes concurrency +through `customProvider.concurrency`. + ### OpenAI and Google Set the provider and corresponding environment credentials: diff --git a/src/config/schema.ts b/src/config/schema.ts index e09d52f8..072d9c2d 100644 --- a/src/config/schema.ts +++ b/src/config/schema.ts @@ -146,11 +146,29 @@ export interface CustomProviderConfig { max_batch_size?: number; } +export interface EmbeddingBatchConfig { + /** Max texts per embedding request. Default: provider-specific (ollama 16). */ + maxBatchItems?: number; + /** Max total input tokens per embedding request. This is a request size/time guard, + * not a model context limit: ollama encodes each input independently, so the per-batch + * token sum is not bounded by the model context length. Default: provider-specific (ollama 65536). + * Raise together with maxBatchItems to pack more texts per request; lower it if a single + * request approaches the request timeout. */ + maxBatchTokens?: number; +} + +export interface EmbeddingConfig { + /** Embedding request batching options. Currently applied to the ollama provider. */ + batch?: EmbeddingBatchConfig; +} + export interface CodebaseIndexConfig { embeddingProvider: EmbeddingProvider | 'custom' | 'auto'; embeddingModel?: EmbeddingModelName; /** Configuration for custom OpenAI-compatible embedding providers (required when embeddingProvider is 'custom') */ customProvider?: CustomProviderConfig; + /** Embedding request shape options (e.g. batch sizes). Currently applied to the ollama provider. */ + embedding?: EmbeddingConfig; scope: IndexScope; indexing?: Partial; search?: Partial; @@ -177,6 +195,7 @@ export type ParsedCodebaseIndexConfig = CodebaseIndexConfig & { reranker?: RerankerConfig; knowledgeBases: string[]; additionalInclude: string[]; + embedding: EmbeddingConfig; }; export function parseConfig(raw: unknown): ParsedCodebaseIndexConfig { @@ -376,10 +395,30 @@ export function parseConfig(raw: unknown): ParsedCodebaseIndexConfig { }; } + const rawEmbedding = (input.embedding && typeof input.embedding === "object" ? input.embedding : {}) as Record; + const rawEmbeddingBatch = (rawEmbedding.batch && typeof rawEmbedding.batch === "object" ? rawEmbedding.batch : null) as Record | null; + const embeddingMaxBatchItems = typeof rawEmbeddingBatch?.maxBatchItems === "number" + && Number.isFinite(rawEmbeddingBatch.maxBatchItems) + ? Math.max(1, Math.floor(rawEmbeddingBatch.maxBatchItems)) + : undefined; + const embeddingMaxBatchTokens = typeof rawEmbeddingBatch?.maxBatchTokens === "number" + && Number.isFinite(rawEmbeddingBatch.maxBatchTokens) + ? Math.max(1, Math.floor(rawEmbeddingBatch.maxBatchTokens)) + : undefined; + const embedding: EmbeddingConfig = (embeddingMaxBatchItems !== undefined || embeddingMaxBatchTokens !== undefined) + ? { + batch: { + ...(embeddingMaxBatchItems !== undefined ? { maxBatchItems: embeddingMaxBatchItems } : {}), + ...(embeddingMaxBatchTokens !== undefined ? { maxBatchTokens: embeddingMaxBatchTokens } : {}), + }, + } + : {}; + return { embeddingProvider, embeddingModel, customProvider, + embedding, scope: isValidScope(scopeValue) ? scopeValue : "project", include: includeValue ?? DEFAULT_INCLUDE, exclude: excludeValue ?? DEFAULT_EXCLUDE, diff --git a/src/embeddings/providers/ollama.ts b/src/embeddings/providers/ollama.ts index 81c614f3..15c099c7 100644 --- a/src/embeddings/providers/ollama.ts +++ b/src/embeddings/providers/ollama.ts @@ -33,6 +33,22 @@ export class OllamaEmbeddingProvider extends BaseEmbeddingProvider(); @@ -155,9 +171,87 @@ export class OllamaEmbeddingProvider extends BaseEmbeddingProvider { - const results: Array<{ embedding: number[]; tokensUsed: number }> = []; + // Embeds many texts in one POST /api/embed request (input: string[]). Ollama + // encodes each input independently, so the model context length applies per input + // (the upstream splitter already bounds each input), not over the batch. This + // amortizes N HTTP round-trips into one. + private async embedMany(texts: string[]): Promise { + const controller = new AbortController(); + const timeout = setTimeout( + () => controller.abort(), + OllamaEmbeddingProvider.REQUEST_TIMEOUT_MS, + ); + let response: Response; + try { + response = await fetch(`${this.credentials.baseUrl}/api/embed`, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + model: this.modelInfo.model, + input: texts, + truncate: false, + }), + signal: controller.signal, + }); + } catch (error) { + if (error instanceof Error && error.name === "AbortError") { + throw new Error( + `Ollama embedding request timed out after ${OllamaEmbeddingProvider.REQUEST_TIMEOUT_MS}ms`, + ); + } + throw error; + } finally { + clearTimeout(timeout); + } + if (!response.ok) { + const error = (await response.text()).slice(0, 500); + if (response.status === 404) { + throw new Error(`Ollama /api/embed not available: ${response.status} - ${error}`); + } + throw new Error(`Ollama embedding API error: ${response.status} - ${error}`); + } + + let parsed: unknown; + try { + parsed = await response.json(); + } catch { + // Invalid JSON (e.g. an empty or truncated 200 body) -> treat as a malformed + // batch so embedBatch falls back to the per-text path instead of propagating + // a parse error that skips the fallback. + throw new Error( + `Ollama returned an invalid embedding batch; expected ${texts.length} vectors of ${this.modelInfo.dimensions} finite dimensions`, + ); + } + const data = (parsed && typeof parsed === "object" ? parsed : {}) as { embeddings?: unknown }; + if ( + !Array.isArray(data.embeddings) + || data.embeddings.length !== texts.length + || data.embeddings.some( + (value) => + !Array.isArray(value) + || value.length !== this.modelInfo.dimensions + || value.some((v) => typeof v !== "number" || !Number.isFinite(v)), + ) + ) { + throw new Error( + `Ollama returned an invalid embedding batch; expected ${texts.length} vectors of ${this.modelInfo.dimensions} finite dimensions`, + ); + } + + return { + embeddings: data.embeddings, + totalTokensUsed: texts.reduce((sum, text) => sum + this.estimateTokens(text), 0), + }; + } + + // Per-text /api/embeddings path shared by the single-text fast path and the + // batch fallback. Uses the legacy endpoint one text at a time, so each text gets + // its own truncation safety net and a vector validated on its own. + private async embedOneByOne(texts: string[]): Promise { + const results: Array<{ embedding: number[]; tokensUsed: number }> = []; for (const text of texts) { results.push(await this.embedSingleWithFallback(text)); } @@ -167,4 +261,37 @@ export class OllamaEmbeddingProvider extends BaseEmbeddingProvider sum + r.tokensUsed, 0), }; } + + public async embedBatch(texts: string[]): Promise { + if (texts.length === 0) { + return { embeddings: [], totalTokensUsed: 0 }; + } + + // A single-text batch gets no batching benefit; send it to the legacy per-text + // path directly so its truncation/timeout/error behavior stays unchanged and no + // /api/embed probe runs (this also keeps old ollama installs on the legacy path + // for the common single-chunk case). + if (texts.length === 1) { + return this.embedOneByOne(texts); + } + + try { + return await this.embedMany(texts); + } catch (error) { + // Fall back to the per-text /api/embeddings path when the batched endpoint is + // unavailable (old ollama, 404), a single input overflowed the model context + // length (per-text truncation net), or the batch response was malformed + // (per-chunk isolation: one bad vector fails only its own chunk instead of + // the whole batch). Other errors propagate for the indexer's pRetry to handle. + if ( + !this.isContextLengthError(error) + && !this.isBatchEndpointUnavailableError(error) + && !this.isBatchValidationError(error) + ) { + throw error; + } + + return this.embedOneByOne(texts); + } + } } diff --git a/src/indexer/index.ts b/src/indexer/index.ts index 1f165bef..6eaa0ccd 100644 --- a/src/indexer/index.ts +++ b/src/indexer/index.ts @@ -6,7 +6,7 @@ import { promisify } from "util"; import PQueue from "p-queue"; import pRetry from "p-retry"; -import { ParsedCodebaseIndexConfig, type RerankerConfig } from "../config/schema.js"; +import { type EmbeddingBatchConfig, ParsedCodebaseIndexConfig, type RerankerConfig } from "../config/schema.js"; import { detectEmbeddingProvider, ConfiguredProviderInfo, tryDetectProvider, createCustomProviderInfo } from "../embeddings/detector.js"; import { createEmbeddingProvider, @@ -305,15 +305,27 @@ function getSafeEmbeddingChunkTokenLimit(provider: ConfiguredProviderInfo): numb return Math.min(2000, maxChunkTokens); } -function getDynamicBatchOptions(provider: ConfiguredProviderInfo): { maxBatchTokens?: number; maxBatchItems?: number } { - if (provider.provider === "ollama") { - return { - maxBatchTokens: provider.modelInfo.maxTokens, - maxBatchItems: 1, - }; - } - - return {}; +// Ollama default batch caps. maxBatchItems is the count limiter; maxBatchTokens is a +// request size/time guard (ollama encodes each input independently, so the per-batch +// token sum is not bounded by the model context). 65536 allows the default 16 items +// (each input is split to <= ~1536 tokens) and is overridable via embedding.batch. +// 16 (not 32) bounds the in-flight workload: ollama concurrency is fixed at 5, so the +// worst case is 5 concurrent batches * 16 inputs (~80 chunks) rather than 160. +const DEFAULT_OLLAMA_MAX_BATCH_ITEMS = 16; +const DEFAULT_OLLAMA_MAX_BATCH_TOKENS = 65_536; + +function getDynamicBatchOptions( + provider: ConfiguredProviderInfo, + embeddingBatch?: EmbeddingBatchConfig, +): { maxBatchTokens?: number; maxBatchItems?: number } { + const base = provider.provider === "ollama" + ? { maxBatchTokens: DEFAULT_OLLAMA_MAX_BATCH_TOKENS, maxBatchItems: DEFAULT_OLLAMA_MAX_BATCH_ITEMS } + : {}; + return { + ...base, + ...(typeof embeddingBatch?.maxBatchTokens === "number" ? { maxBatchTokens: embeddingBatch.maxBatchTokens } : {}), + ...(typeof embeddingBatch?.maxBatchItems === "number" ? { maxBatchItems: embeddingBatch.maxBatchItems } : {}), + }; } function isSqliteCorruptionError(error: unknown): boolean { @@ -2522,6 +2534,9 @@ export class Indexer { forceReembed: boolean; reuseCachedEmbeddings: boolean; incrementRepeatedFailures: boolean; + // When true (recovery path), embed previously-failed chunks one per request + // so a permanently-failing chunk is isolated instead of failing its batch. + forceSingleItemBatches?: boolean; onSucceeded?: (chunks: PendingChunk[]) => void; onProgress?: (progress: Readonly) => void; }, @@ -2580,10 +2595,15 @@ export class Indexer { const embeddingPartsByChunk = new Map>(); const completedVectorsByChunkId = new Map(); const completedChunkIds = new Set(); - const requestBatches = createPendingEmbeddingRequestBatches( - chunksNeedingEmbedding, - getDynamicBatchOptions(options.configuredProviderInfo), - ); + const batchOptions = getDynamicBatchOptions(options.configuredProviderInfo, this.config.embedding?.batch); + // On the recovery path, embed previously-failed chunks one per request so a + // permanently-failing chunk is isolated instead of failing its whole batch. + // Scoped to ollama (whose default batch size groups chunks); other providers + // keep their existing recovery batching. + if (options.forceSingleItemBatches && options.configuredProviderInfo.provider === "ollama") { + batchOptions.maxBatchItems = 1; + } + const requestBatches = createPendingEmbeddingRequestBatches(chunksNeedingEmbedding, batchOptions); let fatalError: unknown; for (const requestBatch of requestBatches) { @@ -4662,6 +4682,7 @@ export class Indexer { forceReembed: forceScopedReembed, reuseCachedEmbeddings: true, incrementRepeatedFailures: true, + forceSingleItemBatches: true, onSucceeded: (succeededChunks) => { database.addChunksToBranchBatch( this.getBranchCatalogKey(), @@ -5898,6 +5919,7 @@ export class Indexer { forceReembed: false, reuseCachedEmbeddings: false, incrementRepeatedFailures: false, + forceSingleItemBatches: true, onSucceeded: (succeededChunks) => { database.addChunksToBranchBatch( this.getBranchCatalogKey(), diff --git a/tests/config.test.ts b/tests/config.test.ts index fb243947..3496e5cd 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -441,6 +441,47 @@ describe("config schema", () => { }); }); + describe("embedding.batch config", () => { + it("defaults to an empty embedding config when no batch knobs are set", () => { + const config = parseConfig({}); + expect(config.embedding).toEqual({}); + }); + + it("parses maxBatchItems and maxBatchTokens", () => { + const config = parseConfig({ + embedding: { batch: { maxBatchItems: 16, maxBatchTokens: 4096 } }, + }); + expect(config.embedding.batch).toEqual({ maxBatchItems: 16, maxBatchTokens: 4096 }); + }); + + it("accepts a single knob and omits the other", () => { + const config = parseConfig({ embedding: { batch: { maxBatchItems: 8 } } }); + expect(config.embedding.batch).toEqual({ maxBatchItems: 8 }); + expect(config.embedding.batch?.maxBatchTokens).toBeUndefined(); + }); + + it("clamps non-positive and fractional values to a minimum of 1", () => { + const config = parseConfig({ + embedding: { batch: { maxBatchItems: 0.5, maxBatchTokens: -10 } }, + }); + expect(config.embedding.batch).toEqual({ maxBatchItems: 1, maxBatchTokens: 1 }); + }); + + it("ignores non-number batch knobs", () => { + const config = parseConfig({ + embedding: { batch: { maxBatchItems: "32", maxBatchTokens: true } }, + } as unknown as Parameters[0]); + expect(config.embedding).toEqual({}); + }); + + it("ignores Infinity and NaN so they cannot defeat batch splitting", () => { + const config = parseConfig({ + embedding: { batch: { maxBatchItems: Infinity, maxBatchTokens: NaN } }, + } as unknown as Parameters[0]); + expect(config.embedding).toEqual({}); + }); + }); + describe("custom provider config", () => { it("should parse valid custom provider config", () => { const config = parseConfig({ diff --git a/tests/custom-provider.test.ts b/tests/custom-provider.test.ts index d7f30964..9a14d7d4 100644 --- a/tests/custom-provider.test.ts +++ b/tests/custom-provider.test.ts @@ -543,29 +543,23 @@ describe("OllamaEmbeddingProvider", () => { expect(result.embeddings).toHaveLength(1); }); - it("processes ollama embedBatch requests sequentially", async () => { - let activeRequests = 0; - let maxActiveRequests = 0; - + it("batches multiple ollama embedBatch texts into a single /api/embed request", async () => { + let calls = 0; fetchSpy.mockImplementation(async (_url, init) => { - const body = JSON.parse(String(init?.body ?? "{}")) as { truncate?: boolean }; + calls += 1; + const body = JSON.parse(String(init?.body ?? "{}")) as { input?: string[]; truncate?: boolean }; expect(body.truncate).toBe(false); - - activeRequests += 1; - maxActiveRequests = Math.max(maxActiveRequests, activeRequests); - - await new Promise((resolve) => setTimeout(resolve, 5)); - - activeRequests -= 1; - return new Response(JSON.stringify({ embedding: new Array(768).fill(0.1) }), { status: 200 }); + expect(body.input).toEqual(["first", "second", "third"]); + return new Response(JSON.stringify({ + embeddings: [0.1, 0.2, 0.3].map((v) => new Array(768).fill(v)), + }), { status: 200 }); }); const provider = createOllamaProvider(); const result = await provider.embedBatch(["first", "second", "third"]); expect(result.embeddings).toHaveLength(3); - expect(maxActiveRequests).toBe(1); - expect(fetchSpy).toHaveBeenCalledTimes(3); + expect(calls).toBe(1); }); it("rethrows non-context ollama errors", async () => { diff --git a/tests/indexer-failed-batches.test.ts b/tests/indexer-failed-batches.test.ts index b535f494..026feb96 100644 --- a/tests/indexer-failed-batches.test.ts +++ b/tests/indexer-failed-batches.test.ts @@ -201,6 +201,27 @@ describe("indexer failed batch recovery", () => { } function createOllamaIndexer(): Indexer { + const config = parseConfig({ + embeddingProvider: "ollama", + embeddingModel: "nomic-embed-text", + // These tests verify per-text recovery machinery (truncation, splitting, + // attemptCount, same-run failure isolation), which requires one chunk per + // batch. The batched /api/embed fast path is covered in ollama-embed.test.ts. + embedding: { batch: { maxBatchItems: 1 } }, + indexing: { + watchFiles: false, + retries: 0, + retryDelayMs: 1, + }, + }); + + return _indexers[_indexers.push(new Indexer(tempDir, config, "opencode")) - 1]; + } + + // Ollama indexer with the DEFAULT batch size (no maxBatchItems pin). Used by + // the recovery-isolation regression test, which needs the batched main run to + // fail a whole batch and the recovery run to re-embed one chunk per request. + function createBatchedOllamaIndexer(): Indexer { const config = parseConfig({ embeddingProvider: "ollama", embeddingModel: "nomic-embed-text", @@ -866,6 +887,59 @@ describe("indexer failed batch recovery", () => { expect(persistedBatches[0]?.error).toContain("persistent split failure"); }); + it("isolates a permanently-failing chunk from healthy chunks on the recovery run (batched ollama)", async () => { + // A poison chunk deterministically fails with a NON-fallback error (HTTP 500, + // not context-length). With the default 16-item batch, run 1 fails the whole + // batch that contains the poison. The recovery run re-embeds one chunk per + // request, so the poison fails alone and the healthy chunks that shared its + // batch index. Without the recovery-path override this regresses: recovery + // re-batches the failed chunks together and the poison fails them all again. + fetchSpy.mockImplementation(async (url, init) => { + if (String(url).endsWith("/api/tags")) { + return new Response(JSON.stringify({ models: [{ name: "nomic-embed-text" }] }), { status: 200 }); + } + const body = JSON.parse(String(init?.body ?? "{}")) as { input?: string[]; prompt?: string }; + const isPoison = (text: string) => text.includes("POISON"); + if (String(url).endsWith("/api/embed")) { + const texts = Array.isArray(body.input) ? body.input : []; + if (texts.some(isPoison)) { + return new Response(JSON.stringify({ error: "persistent failure" }), { status: 500 }); + } + return new Response(JSON.stringify({ + embeddings: texts.map(() => Array.from({ length: 768 }, () => 0.1)), + }), { status: 200 }); + } + if (String(url).endsWith("/api/embeddings")) { + if (isPoison(body.prompt ?? "")) { + return new Response(JSON.stringify({ error: "persistent failure" }), { status: 500 }); + } + return new Response(JSON.stringify({ embedding: Array.from({ length: 768 }, () => 0.1) }), { status: 200 }); + } + throw new Error(`Unexpected request: ${String(url)}`); + }); + + const lines: string[] = []; + for (let i = 0; i < 30; i++) { + lines.push(`export const healthy${i} = 'healthy${i}';`); + } + lines.splice(15, 0, "export const POISON = 'POISON';"); + fs.writeFileSync(sourceFile, lines.join("\n"), "utf-8"); + + const indexer = createBatchedOllamaIndexer(); + const first = await indexer.index(); + // Run 1: the batch containing POISON fails wholesale (non-fallback 500). + expect(first.failedChunks).toBeGreaterThan(0); + + // Recovery run: failed chunks are re-embedded one per request, so only the + // poison chunk stays failed and its healthy co-batch chunks get indexed. + const recovered = await indexer.index(); + expect(recovered.indexedChunks).toBeGreaterThan(0); + expect(recovered.failedChunks).toBe(1); + + const status = await indexer.getStatus(); + expect(status.failedBatchesCount).toBe(1); + }); + it("persists failed batches when storage fails after pooling embeddings", async () => { const addBatchSpy = vi.spyOn(VectorStore.prototype, "addBatch").mockImplementation(() => { throw new Error("vector store write failed"); diff --git a/tests/ollama-embed.test.ts b/tests/ollama-embed.test.ts new file mode 100644 index 00000000..ae115115 --- /dev/null +++ b/tests/ollama-embed.test.ts @@ -0,0 +1,204 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { OllamaEmbeddingProvider } from "../src/embeddings/providers/ollama.js"; + +// Small dimensions keep the mocked vectors tiny; the provider validates every +// vector against modelInfo.dimensions, so this still exercises the real path. +const modelInfo = { + provider: "ollama" as const, + model: "nomic-embed-text", + dimensions: 3, + maxTokens: 2048, + costPer1MTokens: 0, +}; + +function makeProvider() { + return new OllamaEmbeddingProvider( + { provider: "ollama", baseUrl: "http://localhost:11434" }, + modelInfo, + ); +} + +afterEach(() => { + vi.restoreAllMocks(); +}); + +describe("OllamaEmbeddingProvider.embedBatch", () => { + it("sends all texts to /api/embed in one request and parses embeddings in order", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + const body = JSON.parse(String(init?.body)); + expect(body.input).toEqual(["aaa", "bbb"]); + expect(body.truncate).toBe(false); + expect(body.model).toBe("nomic-embed-text"); + return new Response(JSON.stringify({ + embeddings: [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], + })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const result = await makeProvider().embedBatch(["aaa", "bbb"]); + + expect(result.embeddings).toEqual([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]); + // estimateTokens = ceil(len / 4): "aaa" -> 1, "bbb" -> 1 + expect(result.totalTokensUsed).toBe(2); + expect(fetchSpy).toHaveBeenCalledTimes(1); + }); + + it("returns an empty result for an empty batch without calling ollama", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async () => new Response("")); + + const result = await makeProvider().embedBatch([]); + + expect(result).toEqual({ embeddings: [], totalTokensUsed: 0 }); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("uses the legacy /api/embeddings path for a single text and never probes /api/embed", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => { + const url = String(input); + if (url.endsWith("/api/embeddings")) { + const body = JSON.parse(String(init?.body)); + expect(body.prompt).toBe("aaa"); + expect(body.truncate).toBe(false); + return new Response(JSON.stringify({ embedding: [0.1, 0.2, 0.3] })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const result = await makeProvider().embedBatch(["aaa"]); + + expect(result.embeddings).toEqual([[0.1, 0.2, 0.3]]); + expect(result.totalTokensUsed).toBe(1); + expect(fetchSpy).toHaveBeenCalledTimes(1); + }); + + it("falls back to per-text /api/embeddings when /api/embed returns 404 (old ollama)", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + return new Response("not found", { status: 404 }); + } + if (url.endsWith("/api/embeddings")) { + return new Response(JSON.stringify({ embedding: [0.1, 0.2, 0.3] })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const result = await makeProvider().embedBatch(["aaa", "bbb"]); + + expect(result.embeddings).toEqual([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]); + // 1 batched attempt (404) + 2 per-text fallback requests + expect(fetchSpy).toHaveBeenCalledTimes(3); + }); + + it("falls back to per-text truncation when /api/embed signals a context-length error", async () => { + const longText = "x".repeat(10_000); // exceeds maxTokens * 4 = 8192 chars + vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + return new Response(JSON.stringify({ error: "input length exceeds the context length" }), { status: 400 }); + } + if (url.endsWith("/api/embeddings")) { + const body = JSON.parse(String(init?.body)); + if (body.prompt.length > 8192) { + return new Response(JSON.stringify({ error: "context length exceeded" }), { status: 400 }); + } + return new Response(JSON.stringify({ embedding: [0.1, 0.2, 0.3] })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const result = await makeProvider().embedBatch([longText, "aaa"]); + + expect(result.embeddings).toEqual([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]); + }); + + it("falls back to per-text /api/embeddings when /api/embed returns a malformed batch (one bad vector fails only its own chunk)", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + // wrong vector count: 1 vector for 2 inputs + return new Response(JSON.stringify({ embeddings: [[0.1, 0.2, 0.3]] })); + } + if (url.endsWith("/api/embeddings")) { + return new Response(JSON.stringify({ embedding: [0.1, 0.2, 0.3] })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const result = await makeProvider().embedBatch(["aaa", "bbb"]); + + expect(result.embeddings).toEqual([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]); + // 1 batched attempt (rejected) + 2 per-text fallback requests + expect(fetchSpy).toHaveBeenCalledTimes(3); + }); + + it("propagates the per-text error when the fallback path itself fails (no retry loop)", async () => { + vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + // malformed batch triggers the fallback + return new Response(JSON.stringify({ embeddings: [[0.1, 0.2, 0.3]] })); + } + if (url.endsWith("/api/embeddings")) { + // per-text vector is also invalid (wrong dimensions) + return new Response(JSON.stringify({ embedding: [0.1, 0.2] })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + await expect(makeProvider().embedBatch(["aaa", "bbb"])).rejects.toThrow("invalid embedding"); + }); + + it("propagates non-404, non-context-length errors from /api/embed (lets pRetry handle them)", async () => { + vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + return new Response(JSON.stringify({ error: "server error" }), { status: 500 }); + } + throw new Error(`Unexpected request: ${url}`); + }); + + await expect(makeProvider().embedBatch(["aaa", "bbb"])).rejects.toThrow("Ollama embedding API error: 500"); + }); + + it("falls back to per-text /api/embeddings when /api/embed returns a non-JSON 200 body", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + return new Response("not json", { status: 200 }); + } + if (url.endsWith("/api/embeddings")) { + return new Response(JSON.stringify({ embedding: [0.1, 0.2, 0.3] })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const result = await makeProvider().embedBatch(["aaa", "bbb"]); + + expect(result.embeddings).toEqual([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]); + // 1 batched attempt (malformed) + 2 per-text fallback requests + expect(fetchSpy).toHaveBeenCalledTimes(3); + }); + + it("falls back to per-text /api/embeddings when /api/embed returns a null 200 body", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + return new Response("null", { status: 200 }); + } + if (url.endsWith("/api/embeddings")) { + return new Response(JSON.stringify({ embedding: [0.1, 0.2, 0.3] })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const result = await makeProvider().embedBatch(["aaa", "bbb"]); + + expect(result.embeddings).toEqual([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]); + expect(fetchSpy).toHaveBeenCalledTimes(3); + }); +}); \ No newline at end of file From 157ac3686e37ed17242a0dc3e4a7a618f1cdc2a9 Mon Sep 17 00:00:00 2001 From: Dmitri Khokhlov Date: Sun, 16 Aug 2026 23:30:17 -0700 Subject: [PATCH 2/3] fix(embeddings): scope embedding.batch.* to the ollama provider getDynamicBatchOptions applied embedding.batch.{maxBatchItems,maxBatchTokens} to every provider via an unconditional spread outside the provider === "ollama" ternary. The config schema and docs state these options are ollama-only, so a non-ollama provider (openai/google/custom) configured with embedding.batch.maxBatchItems: 1 was split into one-item embedding requests, changing its existing behavior. Return an empty options object immediately for non-ollama providers, and move the embeddingBatch override spread inside the ollama branch. Export the helper so the provider gate can be unit-tested directly. Add a regression test asserting that non-ollama providers ignore embedding.batch.* (an aggressive maxBatchItems: 1 must not split their requests), while ollama still honors overrides and the documented defaults. --- src/indexer/index.ts | 12 ++++++++---- tests/embedding-batches.test.ts | 23 +++++++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/src/indexer/index.ts b/src/indexer/index.ts index 6eaa0ccd..91990016 100644 --- a/src/indexer/index.ts +++ b/src/indexer/index.ts @@ -314,13 +314,17 @@ function getSafeEmbeddingChunkTokenLimit(provider: ConfiguredProviderInfo): numb const DEFAULT_OLLAMA_MAX_BATCH_ITEMS = 16; const DEFAULT_OLLAMA_MAX_BATCH_TOKENS = 65_536; -function getDynamicBatchOptions( +export function getDynamicBatchOptions( provider: ConfiguredProviderInfo, embeddingBatch?: EmbeddingBatchConfig, ): { maxBatchTokens?: number; maxBatchItems?: number } { - const base = provider.provider === "ollama" - ? { maxBatchTokens: DEFAULT_OLLAMA_MAX_BATCH_TOKENS, maxBatchItems: DEFAULT_OLLAMA_MAX_BATCH_ITEMS } - : {}; + // embedding.batch.* is documented as ollama-only. Non-ollama providers keep + // their existing (unbatched-by-this-layer) behavior, so return an empty + // options object regardless of any user-supplied batch config. + if (provider.provider !== "ollama") { + return {}; + } + const base = { maxBatchTokens: DEFAULT_OLLAMA_MAX_BATCH_TOKENS, maxBatchItems: DEFAULT_OLLAMA_MAX_BATCH_ITEMS }; return { ...base, ...(typeof embeddingBatch?.maxBatchTokens === "number" ? { maxBatchTokens: embeddingBatch.maxBatchTokens } : {}), diff --git a/tests/embedding-batches.test.ts b/tests/embedding-batches.test.ts index 9f662265..49a5d2fe 100644 --- a/tests/embedding-batches.test.ts +++ b/tests/embedding-batches.test.ts @@ -14,6 +14,8 @@ import { type FailedBatch, type PendingChunk, } from "../src/indexer/embedding-batches.js"; +import { getDynamicBatchOptions } from "../src/indexer/index.js"; +import type { ConfiguredProviderInfo } from "../src/embeddings/detector.js"; function metadata(filePath = "src/example.ts"): ChunkMetadata { return { @@ -142,4 +144,25 @@ describe("embedding batch helpers", () => { expect(getPendingChunkFilePath(pendingChunk("path"))).toBe("src/path.ts"); expect(getPendingChunkFilePath({ metadata: {} })).toBeNull(); }); + + it("applies embedding.batch.* only to the ollama provider", () => { + // getDynamicBatchOptions reads only provider.provider; the remaining fields + // are irrelevant to this gate, so a minimal cast is sufficient. + const ollama = { provider: "ollama", credentials: {}, modelInfo: { maxTokens: 8192 } } as unknown as ConfiguredProviderInfo; + const openai = { provider: "openai", credentials: {}, modelInfo: { maxTokens: 8192 } } as unknown as ConfiguredProviderInfo; + const google = { provider: "google", credentials: {}, modelInfo: { maxTokens: 8192 } } as unknown as ConfiguredProviderInfo; + const custom = { provider: "custom", credentials: {}, modelInfo: { maxTokens: 8192 } } as unknown as ConfiguredProviderInfo; + + // Ollama with no overrides gets the documented defaults. + expect(getDynamicBatchOptions(ollama)).toEqual({ maxBatchTokens: 65_536, maxBatchItems: 16 }); + // Ollama honors user overrides. + expect(getDynamicBatchOptions(ollama, { maxBatchItems: 1, maxBatchTokens: 100 })).toEqual({ maxBatchTokens: 100, maxBatchItems: 1 }); + + // Non-ollama providers ignore embedding.batch.* entirely: an aggressive + // maxBatchItems: 1 must NOT split their requests into one-item batches. + const aggressiveBatch = { maxBatchItems: 1, maxBatchTokens: 1 }; + expect(getDynamicBatchOptions(openai, aggressiveBatch)).toEqual({}); + expect(getDynamicBatchOptions(google, aggressiveBatch)).toEqual({}); + expect(getDynamicBatchOptions(custom, aggressiveBatch)).toEqual({}); + }); }); From 7dc5d21ab988e4d5a728a576b7534cbb4da97d46 Mon Sep 17 00:00:00 2001 From: Dmitri Khokhlov Date: Mon, 17 Aug 2026 00:41:20 -0700 Subject: [PATCH 3/3] fix(embeddings): address review follow-ups for ollama batched embeds Follow-up to 427d06f and 121d836 incorporating findings from an independent review (claude + codex) of PR #300. No happy-path behavior change; corrects docs/comments and hardens edge paths. - ollama.ts: cache a 404 from /api/embed per provider instance so a legacy ollama install is probed once, not once per multi-text batch. Soften the fallback comments to match the actual behavior: a malformed batch re-embeds each text cleanly, but a text that hard-fails per-text throws and fails the whole request batch. In-run per-text isolation is not provided; a poison text is isolated only on the recovery run, which re-embeds one text per request. - index.ts: guard the getDynamicBatchOptions embedding.batch overrides with Number.isFinite so a programmatic NaN/Infinity cannot poison createDynamicBatches. Correct the in-flight gloss (~80 texts, not chunks). - docs/configuration.md: state that embedding.batch.* is ollama-only (OpenAI, Google, and custom providers ignore it and keep their existing request behavior); describe batched items as embedding texts/parts rather than chunks; document the failed-batch recovery semantics so a one-shot indexer knows healthy co-batch chunks recover on the next index() run. - tests: add a regression test for the 404 cache (one /api/embed hit across two consecutive batches) and for non-finite override rejection; add a trailing newline to ollama-embed.test.ts. Verified: affected suites green (192 tests); typecheck, lint, and build:ts clean. The three watcher-snapshot-reconciler failures are pre-existing on the clean PR head and unrelated to this change. --- docs/configuration.md | 36 ++++++++++++++++++------------ src/embeddings/providers/ollama.ts | 35 ++++++++++++++++++++--------- src/indexer/index.ts | 6 ++--- tests/embedding-batches.test.ts | 13 +++++++++++ tests/ollama-embed.test.ts | 29 +++++++++++++++++++++++- 5 files changed, 91 insertions(+), 28 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 9fb1ecaf..5d4c4254 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -95,16 +95,23 @@ for the methodology and measured results. #### Batching Ollama embeddings -The indexer sends multiple chunks to Ollama in one request. This decreases the -number of HTTP requests and accelerates indexing against a remote Ollama host. +The indexer sends multiple embedding texts to Ollama in one request. This decreases +the number of HTTP requests and accelerates indexing against a remote Ollama host. The `/api/embed` endpoint accepts an array of texts and returns one vector per -text. The indexer uses this endpoint for batches of two or more chunks. A -single-chunk batch uses the legacy `/api/embeddings` endpoint. +text. The indexer uses this endpoint for batches of two or more texts. A single-text +batch uses the legacy `/api/embeddings` endpoint. A chunk that the splitter divides +into multiple parts sends one text per part, so a batch can carry several parts of +the same chunk. If the batched endpoint is not available, the indexer falls back to the legacy -per-chunk endpoint. If one chunk exceeds the model context length, the indexer -truncates and retries that chunk by itself. If the batch response is malformed, -the indexer retries each chunk by itself so one bad vector fails only its chunk. +per-text endpoint and remembers the result, so later batches skip the probe. If one +text exceeds the model context length, the indexer truncates and retries that text +by itself. If the batch response is malformed, the indexer retries each text by +itself so a bad batch response re-embeds each text cleanly. This is not in-run +per-text isolation: if a text then hard-fails per-text, the whole request batch +fails and is marked failed. The chunks in that batch recover on the next `index()` +run, where the recovery path re-embeds one text per request so a persistently-failing +text is isolated from healthy texts. Control the batch size with the `embedding.batch` section: @@ -121,19 +128,20 @@ Control the batch size with the `embedding.batch` section: | Option | Ollama default | Purpose | |---|---:|---| -| `maxBatchItems` | `16` | Maximum number of chunks in one request | +| `maxBatchItems` | `16` | Maximum number of embedding texts in one request | | `maxBatchTokens` | `65536` | Maximum total estimated tokens in one request | Ollama encodes each text independently. The model context length applies to each text and not to the batch total. Set `maxBatchTokens` to bound the request size and -the processing time. Set `maxBatchItems` to bound the number of texts. Both values -are optional and must be at least 1. When you omit a value, the indexer uses the -Ollama default. These knobs apply to any provider when you set them; the Ollama -defaults apply only to the Ollama provider. +the processing time. Set `maxBatchItems` to bound the number of texts (a chunk split +into multiple parts counts as one text per part). Both values are optional and must +be at least 1. When you omit a value, the indexer uses the Ollama default. These +knobs apply only to the Ollama provider; OpenAI, Google, and custom providers ignore +them and keep their existing request behavior. The indexer runs up to five Ollama requests at the same time. Each request carries -up to `maxBatchItems` chunks, so the worst case is five times `maxBatchItems` chunks -in flight (80 chunks at the default 16). A remote or memory-limited Ollama host can +up to `maxBatchItems` texts, so the worst case is five times `maxBatchItems` texts +in flight (80 texts at the default 16). A remote or memory-limited Ollama host can run out of memory or exceed the 120-second request timeout when this number is too high. Lower `maxBatchItems` for a small or remote host. The Ollama concurrency is fixed at five and is not configurable; the custom provider exposes concurrency diff --git a/src/embeddings/providers/ollama.ts b/src/embeddings/providers/ollama.ts index 15c099c7..900fec68 100644 --- a/src/embeddings/providers/ollama.ts +++ b/src/embeddings/providers/ollama.ts @@ -7,6 +7,11 @@ export class OllamaEmbeddingProvider extends BaseEmbeddingProvider { const results: Array<{ embedding: number[]; tokensUsed: number }> = []; for (const text of texts) { @@ -269,9 +277,9 @@ export class OllamaEmbeddingProvider extends BaseEmbeddingProvider { expect(getDynamicBatchOptions(google, aggressiveBatch)).toEqual({}); expect(getDynamicBatchOptions(custom, aggressiveBatch)).toEqual({}); }); + + it("ignores non-finite embedding.batch.* values and falls back to the ollama defaults", () => { + const ollama = { provider: "ollama", credentials: {}, modelInfo: { maxTokens: 8192 } } as unknown as ConfiguredProviderInfo; + // NaN, Infinity, and -Infinity are typeof "number" but must not poison the batch + // options; the documented defaults apply instead. + expect(getDynamicBatchOptions(ollama, { maxBatchItems: Number.NaN, maxBatchTokens: Number.NaN })) + .toEqual({ maxBatchTokens: 65_536, maxBatchItems: 16 }); + expect(getDynamicBatchOptions(ollama, { maxBatchItems: Number.POSITIVE_INFINITY, maxBatchTokens: Number.NEGATIVE_INFINITY })) + .toEqual({ maxBatchTokens: 65_536, maxBatchItems: 16 }); + // A finite override still wins; a non-finite sibling still falls back per-field. + expect(getDynamicBatchOptions(ollama, { maxBatchItems: 4, maxBatchTokens: Number.NaN })) + .toEqual({ maxBatchTokens: 65_536, maxBatchItems: 4 }); + }); }); diff --git a/tests/ollama-embed.test.ts b/tests/ollama-embed.test.ts index ae115115..e87b15f9 100644 --- a/tests/ollama-embed.test.ts +++ b/tests/ollama-embed.test.ts @@ -94,6 +94,33 @@ describe("OllamaEmbeddingProvider.embedBatch", () => { expect(fetchSpy).toHaveBeenCalledTimes(3); }); + it("caches the 404 result so later multi-text batches skip the /api/embed probe", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const url = String(input); + if (url.endsWith("/api/embed")) { + return new Response("not found", { status: 404 }); + } + if (url.endsWith("/api/embeddings")) { + return new Response(JSON.stringify({ embedding: [0.1, 0.2, 0.3] })); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const provider = makeProvider(); + + // First multi-text batch: 1 batched attempt (404) + 2 per-text fallback = 3 calls. + await provider.embedBatch(["aaa", "bbb"]); + // Second multi-text batch: the 404 is cached, so it goes straight to the legacy + // path (2 per-text requests) and does NOT probe /api/embed again. + const result = await provider.embedBatch(["ccc", "ddd"]); + + expect(result.embeddings).toEqual([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]); + // 3 (first batch) + 2 (second batch, legacy only) = 5 total; only 1 hit /api/embed. + expect(fetchSpy).toHaveBeenCalledTimes(5); + const embedCalls = fetchSpy.mock.calls.filter((call) => String(call[0]).endsWith("/api/embed")); + expect(embedCalls).toHaveLength(1); + }); + it("falls back to per-text truncation when /api/embed signals a context-length error", async () => { const longText = "x".repeat(10_000); // exceeds maxTokens * 4 = 8192 chars vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => { @@ -201,4 +228,4 @@ describe("OllamaEmbeddingProvider.embedBatch", () => { expect(result.embeddings).toEqual([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]); expect(fetchSpy).toHaveBeenCalledTimes(3); }); -}); \ No newline at end of file +});