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
This commit is contained in:
Jake
Keith the Silly Goose
parent
0face480eb
commit
600150a2cb
170
src/memory/batch-voyage.test.ts
Normal file
170
src/memory/batch-voyage.test.ts
Normal file
@ -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 <T>(fn: () => Promise<T>) => 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]);
|
||||
});
|
||||
});
|
||||
@ -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<string> {
|
||||
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<string | undefined> {
|
||||
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<T>(tasks: Array<() => Promise<T>>, limit: number): Promise<T[]> {
|
||||
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) {
|
||||
|
||||
100
src/memory/embeddings-voyage.test.ts
Normal file
100
src/memory/embeddings-voyage.test.ts
Normal file
@ -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<string, string>;
|
||||
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<string, string>;
|
||||
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
|
||||
});
|
||||
});
|
||||
@ -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<T>(
|
||||
tasks: Array<() => Promise<T>>,
|
||||
limit: number,
|
||||
): Promise<T[]> {
|
||||
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;
|
||||
}
|
||||
|
||||
@ -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<T>(tasks: Array<() => Promise<T>>, limit: number): Promise<T[]> {
|
||||
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<T>(fn: () => Promise<T>): Promise<T> {
|
||||
let release: () => void;
|
||||
const wait = this.batchFailureLock;
|
||||
|
||||
Loading…
Reference in New Issue
Block a user