diff --git a/docs/configuration.md b/docs/configuration.md index 9989bcb..5d4c425 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -93,6 +93,60 @@ 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 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 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-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: + +```json +{ + "embedding": { + "batch": { + "maxBatchItems": 32, + "maxBatchTokens": 65536 + } + } +} +``` + +| Option | Ollama default | Purpose | +|---|---:|---| +| `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 (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` 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 +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 e09d52f..072d9c2 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 81c614f..900fec6 100644 --- a/src/embeddings/providers/ollama.ts +++ b/src/embeddings/providers/ollama.ts @@ -7,6 +7,11 @@ export class OllamaEmbeddingProvider extends BaseEmbeddingProvider(); @@ -155,9 +177,89 @@ 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. A text that + // hard-fails per-text throws here and fails the whole request batch; the recovery + // run re-embeds one text per request to isolate it. + 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 +269,44 @@ 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. Once /api/embed is known unavailable, multi-text batches + // also skip the probe (see batchEndpointUnavailable). + if (texts.length === 1 || this.batchEndpointUnavailable) { + 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 (so a + // bad batch response re-embeds each text cleanly). This is not in-run per-text + // isolation: a text that hard-fails per-text throws and fails the whole request + // batch, and is isolated only on the recovery run, which re-embeds one text per + // request. Other errors propagate for the indexer's pRetry to handle. + if (this.isBatchEndpointUnavailableError(error)) { + // Old ollama without /api/embed: cache the result so later batches skip the probe. + this.batchEndpointUnavailable = true; + return this.embedOneByOne(texts); + } + + if ( + !this.isContextLengthError(error) + && !this.isBatchValidationError(error) + ) { + throw error; + } + + return this.embedOneByOne(texts); + } + } } diff --git a/src/indexer/index.ts b/src/indexer/index.ts index 1f165be..a49c355 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,31 @@ 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, - }; +// 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 texts) rather than 160. +const DEFAULT_OLLAMA_MAX_BATCH_ITEMS = 16; +const DEFAULT_OLLAMA_MAX_BATCH_TOKENS = 65_536; + +export function getDynamicBatchOptions( + provider: ConfiguredProviderInfo, + embeddingBatch?: EmbeddingBatchConfig, +): { maxBatchTokens?: number; maxBatchItems?: number } { + // 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 {}; } - - return {}; + const base = { maxBatchTokens: DEFAULT_OLLAMA_MAX_BATCH_TOKENS, maxBatchItems: DEFAULT_OLLAMA_MAX_BATCH_ITEMS }; + return { + ...base, + ...(typeof embeddingBatch?.maxBatchTokens === "number" && Number.isFinite(embeddingBatch.maxBatchTokens) ? { maxBatchTokens: embeddingBatch.maxBatchTokens } : {}), + ...(typeof embeddingBatch?.maxBatchItems === "number" && Number.isFinite(embeddingBatch.maxBatchItems) ? { maxBatchItems: embeddingBatch.maxBatchItems } : {}), + }; } function isSqliteCorruptionError(error: unknown): boolean { @@ -2522,6 +2538,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 +2599,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 +4686,7 @@ export class Indexer { forceReembed: forceScopedReembed, reuseCachedEmbeddings: true, incrementRepeatedFailures: true, + forceSingleItemBatches: true, onSucceeded: (succeededChunks) => { database.addChunksToBranchBatch( this.getBranchCatalogKey(), @@ -5898,6 +5923,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 fb24394..3496e5c 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 d7f3096..9a14d7d 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/embedding-batches.test.ts b/tests/embedding-batches.test.ts index 9f66226..b4be656 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,38 @@ 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({}); + }); + + 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/indexer-failed-batches.test.ts b/tests/indexer-failed-batches.test.ts index b535f49..026feb9 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 0000000..e87b15f --- /dev/null +++ b/tests/ollama-embed.test.ts @@ -0,0 +1,231 @@ +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("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) => { + 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); + }); +});