From 600150a2cb75dabcafd4222a31f2edcab48f1d95 Mon Sep 17 00:00:00 2001 From: Jake Date: Wed, 28 Jan 2026 11:26:33 +1300 Subject: [PATCH] fix(memory): optimize voyage batch memory usage with streaming and deduplicate code Cherry-picked from PR #2519. Fixed lint error: changed this.runWithConcurrency to use imported runWithConcurrency function after extraction to internal.ts --- src/memory/batch-voyage.test.ts | 170 +++++++++++++++++++++++++++ src/memory/batch-voyage.ts | 143 ++++++++++------------ src/memory/embeddings-voyage.test.ts | 100 ++++++++++++++++ src/memory/internal.ts | 30 +++++ src/memory/manager.ts | 32 +---- 5 files changed, 362 insertions(+), 113 deletions(-) create mode 100644 src/memory/batch-voyage.test.ts create mode 100644 src/memory/embeddings-voyage.test.ts diff --git a/src/memory/batch-voyage.test.ts b/src/memory/batch-voyage.test.ts new file mode 100644 index 000000000..e0e757f19 --- /dev/null +++ b/src/memory/batch-voyage.test.ts @@ -0,0 +1,170 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { ReadableStream } from "node:stream/web"; +import type { VoyageBatchOutputLine, VoyageBatchRequest } from "./batch-voyage.js"; +import type { VoyageEmbeddingClient } from "./embeddings-voyage.js"; + +// Mock internal.js if needed, but runWithConcurrency is simple enough to keep real. +// We DO need to mock retryAsync to avoid actual delays/retries logic complicating tests +vi.mock("../infra/retry.js", () => ({ + retryAsync: async (fn: () => Promise) => fn(), +})); + +describe("runVoyageEmbeddingBatches", () => { + afterEach(() => { + vi.resetAllMocks(); + vi.unstubAllGlobals(); + }); + + const mockClient: VoyageEmbeddingClient = { + baseUrl: "https://api.voyageai.com/v1", + headers: { Authorization: "Bearer test-key" }, + model: "voyage-4-large", + }; + + const mockRequests: VoyageBatchRequest[] = [ + { custom_id: "req-1", body: { input: "text1" } }, + { custom_id: "req-2", body: { input: "text2" } }, + ]; + + it("successfully submits batch, waits, and streams results", async () => { + const fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); + + // Sequence of fetch calls: + // 1. Upload file + fetchMock.mockResolvedValueOnce({ + ok: true, + json: async () => ({ id: "file-123" }), + }); + + // 2. Create batch + fetchMock.mockResolvedValueOnce({ + ok: true, + json: async () => ({ id: "batch-abc", status: "pending" }), + }); + + // 3. Poll status (pending) - Optional depending on wait loop, let's say it finishes immediately for this test + // Actually the code does: initial check (if completed) -> wait loop. + // If create returns "pending", it enters waitForVoyageBatch. + // waitForVoyageBatch fetches status. + + // 3. Poll status (completed) + fetchMock.mockResolvedValueOnce({ + ok: true, + json: async () => ({ + id: "batch-abc", + status: "completed", + output_file_id: "file-out-999", + }), + }); + + // 4. Download content (Streaming) + const outputLines: VoyageBatchOutputLine[] = [ + { + custom_id: "req-1", + response: { status_code: 200, body: { data: [{ embedding: [0.1, 0.1] }] } }, + }, + { + custom_id: "req-2", + response: { status_code: 200, body: { data: [{ embedding: [0.2, 0.2] }] } }, + }, + ]; + + // Create a stream that emits the NDJSON lines + const stream = new ReadableStream({ + start(controller) { + const text = outputLines.map((l) => JSON.stringify(l)).join("\n"); + controller.enqueue(new TextEncoder().encode(text)); + controller.close(); + }, + }); + + fetchMock.mockResolvedValueOnce({ + ok: true, + body: stream, + }); + + const { runVoyageEmbeddingBatches } = await import("./batch-voyage.js"); + + const results = await runVoyageEmbeddingBatches({ + client: mockClient, + agentId: "agent-1", + requests: mockRequests, + wait: true, + pollIntervalMs: 1, // fast poll + timeoutMs: 1000, + concurrency: 1, + }); + + expect(results.size).toBe(2); + expect(results.get("req-1")).toEqual([0.1, 0.1]); + expect(results.get("req-2")).toEqual([0.2, 0.2]); + + // Verify calls + expect(fetchMock).toHaveBeenCalledTimes(4); + + // Verify File Upload + expect(fetchMock.mock.calls[0][0]).toContain("/files"); + const uploadBody = fetchMock.mock.calls[0][1].body as FormData; + expect(uploadBody).toBeInstanceOf(FormData); + expect(uploadBody.get("purpose")).toBe("batch"); + + // Verify Batch Create + expect(fetchMock.mock.calls[1][0]).toContain("/batches"); + const createBody = JSON.parse(fetchMock.mock.calls[1][1].body); + expect(createBody.input_file_id).toBe("file-123"); + expect(createBody.completion_window).toBe("12h"); + + // Verify Content Fetch + expect(fetchMock.mock.calls[3][0]).toContain("/files/file-out-999/content"); + }); + + it("handles empty lines and stream chunks correctly", async () => { + const fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); + + // 1. Upload + fetchMock.mockResolvedValueOnce({ ok: true, json: async () => ({ id: "f1" }) }); + // 2. Create (completed immediately) + fetchMock.mockResolvedValueOnce({ + ok: true, + json: async () => ({ id: "b1", status: "completed", output_file_id: "out1" }), + }); + // 3. Download Content (Streaming with chunks and newlines) + const stream = new ReadableStream({ + start(controller) { + const line1 = JSON.stringify({ + custom_id: "req-1", + response: { body: { data: [{ embedding: [1] }] } }, + }); + const line2 = JSON.stringify({ + custom_id: "req-2", + response: { body: { data: [{ embedding: [2] }] } }, + }); + + // Split across chunks + controller.enqueue(new TextEncoder().encode(line1 + "\n")); + controller.enqueue(new TextEncoder().encode("\n")); // empty line + controller.enqueue(new TextEncoder().encode(line2)); // no newline at EOF + controller.close(); + }, + }); + + fetchMock.mockResolvedValueOnce({ ok: true, body: stream }); + + const { runVoyageEmbeddingBatches } = await import("./batch-voyage.js"); + + const results = await runVoyageEmbeddingBatches({ + client: mockClient, + agentId: "a1", + requests: mockRequests, + wait: true, + pollIntervalMs: 1, + timeoutMs: 1000, + concurrency: 1, + }); + + expect(results.get("req-1")).toEqual([1]); + expect(results.get("req-2")).toEqual([2]); + }); +}); diff --git a/src/memory/batch-voyage.ts b/src/memory/batch-voyage.ts index a1f101e0c..5e882738c 100644 --- a/src/memory/batch-voyage.ts +++ b/src/memory/batch-voyage.ts @@ -1,6 +1,9 @@ +import { createInterface } from "node:readline"; +import { Readable } from "node:stream"; + import { retryAsync } from "../infra/retry.js"; import type { VoyageEmbeddingClient } from "./embeddings-voyage.js"; -import { hashText } from "./internal.js"; +import { hashText, runWithConcurrency } from "./internal.js"; /** * Voyage Batch API Input Line format. @@ -153,40 +156,26 @@ async function fetchVoyageBatchStatus(params: { return (await res.json()) as VoyageBatchStatus; } -async function fetchVoyageFileContent(params: { - client: VoyageEmbeddingClient; - fileId: string; -}): Promise { - const baseUrl = getVoyageBaseUrl(params.client); - const res = await fetch(`${baseUrl}/files/${params.fileId}/content`, { - headers: getVoyageHeaders(params.client, { json: true }), - }); - if (!res.ok) { - const text = await res.text(); - throw new Error(`voyage batch file content failed: ${res.status} ${text}`); - } - return await res.text(); -} - -function parseVoyageBatchOutput(text: string): VoyageBatchOutputLine[] { - if (!text.trim()) return []; - return text - .split("\n") - .map((line) => line.trim()) - .filter(Boolean) - .map((line) => JSON.parse(line) as VoyageBatchOutputLine); -} - async function readVoyageBatchError(params: { client: VoyageEmbeddingClient; errorFileId: string; }): Promise { try { - const content = await fetchVoyageFileContent({ - client: params.client, - fileId: params.errorFileId, + const baseUrl = getVoyageBaseUrl(params.client); + const res = await fetch(`${baseUrl}/files/${params.errorFileId}/content`, { + headers: getVoyageHeaders(params.client, { json: true }), }); - const lines = parseVoyageBatchOutput(content); + if (!res.ok) { + const text = await res.text(); + throw new Error(`voyage batch error file content failed: ${res.status} ${text}`); + } + const text = await res.text(); + if (!text.trim()) return undefined; + const lines = text + .split("\n") + .map((line) => line.trim()) + .filter(Boolean) + .map((line) => JSON.parse(line) as VoyageBatchOutputLine); const first = lines.find((line) => line.error?.message || line.response?.body?.error); const message = first?.error?.message ?? @@ -247,33 +236,6 @@ async function waitForVoyageBatch(params: { } } -async function runWithConcurrency(tasks: Array<() => Promise>, limit: number): Promise { - if (tasks.length === 0) return []; - const resolvedLimit = Math.max(1, Math.min(limit, tasks.length)); - const results: T[] = Array.from({ length: tasks.length }); - let next = 0; - let firstError: unknown = null; - - const workers = Array.from({ length: resolvedLimit }, async () => { - while (true) { - if (firstError) return; - const index = next; - next += 1; - if (index >= tasks.length) return; - try { - results[index] = await tasks[index](); - } catch (err) { - firstError = err; - return; - } - } - }); - - await Promise.allSettled(workers); - if (firstError) throw firstError; - return results; -} - export async function runVoyageEmbeddingBatches(params: { client: VoyageEmbeddingClient; agentId: string; @@ -331,39 +293,52 @@ export async function runVoyageEmbeddingBatches(params: { throw new Error(`voyage batch ${batchInfo.id} completed without output file`); } - const content = await fetchVoyageFileContent({ - client: params.client, - fileId: completed.outputFileId, + const baseUrl = getVoyageBaseUrl(params.client); + const contentRes = await fetch(`${baseUrl}/files/${completed.outputFileId}/content`, { + headers: getVoyageHeaders(params.client, { json: true }), }); - const outputLines = parseVoyageBatchOutput(content); + if (!contentRes.ok) { + const text = await contentRes.text(); + throw new Error(`voyage batch file content failed: ${contentRes.status} ${text}`); + } + const errors: string[] = []; const remaining = new Set(group.map((request) => request.custom_id)); - for (const line of outputLines) { - const customId = line.custom_id; - if (!customId) continue; - remaining.delete(customId); - if (line.error?.message) { - errors.push(`${customId}: ${line.error.message}`); - continue; + if (contentRes.body) { + const reader = createInterface({ + input: Readable.fromWeb(contentRes.body as any), + terminal: false, + }); + + for await (const rawLine of reader) { + if (!rawLine.trim()) continue; + const line = JSON.parse(rawLine) as VoyageBatchOutputLine; + const customId = line.custom_id; + if (!customId) continue; + remaining.delete(customId); + if (line.error?.message) { + errors.push(`${customId}: ${line.error.message}`); + continue; + } + const response = line.response; + const statusCode = response?.status_code ?? 0; + if (statusCode >= 400) { + const message = + response?.body?.error?.message ?? + (typeof response?.body === "string" ? response.body : undefined) ?? + "unknown error"; + errors.push(`${customId}: ${message}`); + continue; + } + const data = response?.body?.data ?? []; + const embedding = data[0]?.embedding ?? []; + if (embedding.length === 0) { + errors.push(`${customId}: empty embedding`); + continue; + } + byCustomId.set(customId, embedding); } - const response = line.response; - const statusCode = response?.status_code ?? 0; - if (statusCode >= 400) { - const message = - response?.body?.error?.message ?? - (typeof response?.body === "string" ? response.body : undefined) ?? - "unknown error"; - errors.push(`${customId}: ${message}`); - continue; - } - const data = response?.body?.data ?? []; - const embedding = data[0]?.embedding ?? []; - if (embedding.length === 0) { - errors.push(`${customId}: empty embedding`); - continue; - } - byCustomId.set(customId, embedding); } if (errors.length > 0) { diff --git a/src/memory/embeddings-voyage.test.ts b/src/memory/embeddings-voyage.test.ts new file mode 100644 index 000000000..d9cc1d541 --- /dev/null +++ b/src/memory/embeddings-voyage.test.ts @@ -0,0 +1,100 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../agents/model-auth.js", () => ({ + resolveApiKeyForProvider: vi.fn(), + requireApiKey: (auth: { apiKey?: string; mode?: string }, provider: string) => { + if (auth?.apiKey) return auth.apiKey; + throw new Error(`No API key resolved for provider "${provider}" (auth mode: ${auth?.mode}).`); + }, +})); + +const createFetchMock = () => + vi.fn(async () => ({ + ok: true, + status: 200, + json: async () => ({ data: [{ embedding: [0.1, 0.2, 0.3] }] }), + })) as unknown as typeof fetch; + +describe("voyage embedding provider", () => { + afterEach(() => { + vi.resetAllMocks(); + vi.resetModules(); + vi.unstubAllGlobals(); + }); + + it("configures client with correct defaults and headers", async () => { + const fetchMock = createFetchMock(); + vi.stubGlobal("fetch", fetchMock); + + const { createVoyageEmbeddingProvider } = await import("./embeddings-voyage.js"); + const authModule = await import("../agents/model-auth.js"); + + vi.mocked(authModule.resolveApiKeyForProvider).mockResolvedValue({ + apiKey: "voyage-key-123", + mode: "api-key", + source: "test", + }); + + const result = await createVoyageEmbeddingProvider({ + config: {} as never, + provider: "voyage", + model: "voyage-4-large", + fallback: "none", + }); + + await result.provider.embedQuery("test query"); + + expect(authModule.resolveApiKeyForProvider).toHaveBeenCalledWith( + expect.objectContaining({ provider: "voyage" }), + ); + + const [url, init] = fetchMock.mock.calls[0] ?? []; + expect(url).toBe("https://api.voyageai.com/v1/embeddings"); + + const headers = (init?.headers ?? {}) as Record; + expect(headers.Authorization).toBe("Bearer voyage-key-123"); + expect(headers["Content-Type"]).toBe("application/json"); + + const body = JSON.parse(init?.body as string); + expect(body).toEqual({ + model: "voyage-4-large", + input: ["test query"], + }); + }); + + it("respects remote overrides for baseUrl and apiKey", async () => { + const fetchMock = createFetchMock(); + vi.stubGlobal("fetch", fetchMock); + + const { createVoyageEmbeddingProvider } = await import("./embeddings-voyage.js"); + + const result = await createVoyageEmbeddingProvider({ + config: {} as never, + provider: "voyage", + model: "voyage-4-lite", + fallback: "none", + remote: { + baseUrl: "https://proxy.example.com", + apiKey: "remote-override-key", + headers: { "X-Custom": "123" }, + }, + }); + + await result.provider.embedQuery("test"); + + const [url, init] = fetchMock.mock.calls[0] ?? []; + expect(url).toBe("https://proxy.example.com/embeddings"); + + const headers = (init?.headers ?? {}) as Record; + expect(headers.Authorization).toBe("Bearer remote-override-key"); + expect(headers["X-Custom"]).toBe("123"); + }); + + it("normalizes model names", async () => { + const { normalizeVoyageModel } = await import("./embeddings-voyage.js"); + expect(normalizeVoyageModel("voyage/voyage-large-2")).toBe("voyage-large-2"); + expect(normalizeVoyageModel("voyage-4-large")).toBe("voyage-4-large"); + expect(normalizeVoyageModel(" voyage-lite ")).toBe("voyage-lite"); + expect(normalizeVoyageModel("")).toBe("voyage-4-large"); // Default + }); +}); diff --git a/src/memory/internal.ts b/src/memory/internal.ts index b68570c35..0a548a27d 100644 --- a/src/memory/internal.ts +++ b/src/memory/internal.ts @@ -203,3 +203,33 @@ export function cosineSimilarity(a: number[], b: number[]): number { if (normA === 0 || normB === 0) return 0; return dot / (Math.sqrt(normA) * Math.sqrt(normB)); } + +export async function runWithConcurrency( + tasks: Array<() => Promise>, + limit: number, +): Promise { + if (tasks.length === 0) return []; + const resolvedLimit = Math.max(1, Math.min(limit, tasks.length)); + const results: T[] = Array.from({ length: tasks.length }); + let next = 0; + let firstError: unknown = null; + + const workers = Array.from({ length: resolvedLimit }, async () => { + while (true) { + if (firstError) return; + const index = next; + next += 1; + if (index >= tasks.length) return; + try { + results[index] = await tasks[index](); + } catch (err) { + firstError = err; + return; + } + } + }); + + await Promise.allSettled(workers); + if (firstError) throw firstError; + return results; +} diff --git a/src/memory/manager.ts b/src/memory/manager.ts index a592bb1f3..40217ecc5 100644 --- a/src/memory/manager.ts +++ b/src/memory/manager.ts @@ -42,6 +42,7 @@ import { type MemoryFileEntry, normalizeRelPath, parseEmbedding, + runWithConcurrency, } from "./internal.js"; import { bm25RankToScore, buildFtsQuery, mergeHybridResults } from "./hybrid.js"; import { searchKeyword, searchVector } from "./manager-search.js"; @@ -1023,7 +1024,7 @@ export class MemoryIndexManager { }); } }); - await this.runWithConcurrency(tasks, this.getIndexConcurrency()); + await runWithConcurrency(tasks, this.getIndexConcurrency()); const staleRows = this.db .prepare(`SELECT path FROM files WHERE source = ?`) @@ -1118,7 +1119,7 @@ export class MemoryIndexManager { }); } }); - await this.runWithConcurrency(tasks, this.getIndexConcurrency()); + await runWithConcurrency(tasks, this.getIndexConcurrency()); const staleRows = this.db .prepare(`SELECT path FROM files WHERE source = ?`) @@ -2021,33 +2022,6 @@ export class MemoryIndexManager { } } - private async runWithConcurrency(tasks: Array<() => Promise>, limit: number): Promise { - if (tasks.length === 0) return []; - const resolvedLimit = Math.max(1, Math.min(limit, tasks.length)); - const results: T[] = Array.from({ length: tasks.length }); - let next = 0; - let firstError: unknown = null; - - const workers = Array.from({ length: resolvedLimit }, async () => { - while (true) { - if (firstError) return; - const index = next; - next += 1; - if (index >= tasks.length) return; - try { - results[index] = await tasks[index](); - } catch (err) { - firstError = err; - return; - } - } - }); - - await Promise.allSettled(workers); - if (firstError) throw firstError; - return results; - } - private async withBatchFailureLock(fn: () => Promise): Promise { let release: () => void; const wait = this.batchFailureLock;