Skip to content

Commit ce4a259

Browse files
sunlit-dengvincentkoc
sunlit-deng
and
vincentkoc
authored
fix(openai): bound embedding batch file downloads (#98554)
* fix(openai): bound embedding batch file downloads * fix(openai): bound batch output records --------- Co-authored-by: Vincent Koc <[email protected]>
1 parent 8abd5d4 commit ce4a259

2 files changed

Lines changed: 395 additions & 16 deletions

File tree

‎extensions/openai/embedding-batch.test.ts‎

Lines changed: 256 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
// Openai tests cover embedding batch plugin behavior.
2+
import { createServer } from "node:http";
23
import { describe, expect, it, vi } from "vitest";
34
import { parseOpenAiBatchOutput, runOpenAiEmbeddingBatches } from "./embedding-batch.js";
45

@@ -54,6 +55,33 @@ function parseStringBody(init: RequestInit | undefined): unknown {
5455
return JSON.parse(init.body) as unknown;
5556
}
5657

58+
async function listenLoopbackServer(server: ReturnType<typeof createServer>): Promise<number> {
59+
return await new Promise((resolve, reject) => {
60+
server.once("error", reject);
61+
server.listen(0, "127.0.0.1", () => {
62+
server.off("error", reject);
63+
const address = server.address();
64+
if (!address || typeof address === "string") {
65+
reject(new Error("expected loopback TCP address"));
66+
return;
67+
}
68+
resolve(address.port);
69+
});
70+
});
71+
}
72+
73+
async function closeServer(server: ReturnType<typeof createServer>): Promise<void> {
74+
await new Promise<void>((resolve, reject) => {
75+
server.close((err) => {
76+
if (err) {
77+
reject(err);
78+
return;
79+
}
80+
resolve();
81+
});
82+
});
83+
}
84+
5785
describe("OpenAI embedding batch output", () => {
5886
it("wraps malformed JSONL output", () => {
5987
expect(() => parseOpenAiBatchOutput('{"custom_id":"ok"}\n{not json')).toThrow(
@@ -330,6 +358,234 @@ describe("OpenAI embedding batch output", () => {
330358
expect(readCount).toBeLessThan(chunkCount);
331359
});
332360

361+
it("streams valid batch output files larger than the provider text cap", async () => {
362+
const outputLineCount = 18;
363+
const padding = "x".repeat(1024 * 1024);
364+
const requests: Parameters<typeof runOpenAiEmbeddingBatches>[0]["requests"] = Array.from(
365+
{ length: outputLineCount },
366+
(_, index) => ({
367+
custom_id: String(index),
368+
method: "POST" as const,
369+
url: "/v1/embeddings",
370+
body: { model: "text-embedding-3-small", input: `payload-${index}` },
371+
}),
372+
);
373+
let outputLinesSent = 0;
374+
const outputResponse = new Response(
375+
new ReadableStream<Uint8Array>({
376+
pull(controller) {
377+
if (outputLinesSent >= outputLineCount) {
378+
controller.close();
379+
return;
380+
}
381+
const customId = String(outputLinesSent);
382+
controller.enqueue(
383+
jsonlEncoder.encode(
384+
`${JSON.stringify({
385+
custom_id: customId,
386+
response: {
387+
status_code: 200,
388+
body: { data: [{ embedding: [outputLinesSent + 1] }] },
389+
},
390+
padding,
391+
})}\n`,
392+
),
393+
);
394+
outputLinesSent += 1;
395+
},
396+
}),
397+
{ status: 200, headers: { "Content-Type": "application/jsonl" } },
398+
);
399+
const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
400+
const url = fetchInputUrl(input);
401+
if (url.endsWith("/files") && init?.method === "POST") {
402+
return jsonResponse({ id: "file-0" });
403+
}
404+
if (url.endsWith("/batches") && init?.method === "POST") {
405+
return jsonResponse({ id: "batch-0", status: "completed", output_file_id: "output-0" });
406+
}
407+
if (url.endsWith("/files/output-0/content")) {
408+
return outputResponse;
409+
}
410+
return new Response("unexpected request", { status: 500 });
411+
});
412+
413+
const byCustomId = await runOpenAiEmbeddingBatches({
414+
openAi: {
415+
baseUrl: "https://openai-compatible.example/v1",
416+
headers: { Authorization: "Bearer test" },
417+
model: "text-embedding-3-small",
418+
fetchImpl,
419+
},
420+
agentId: "main",
421+
requests,
422+
wait: true,
423+
concurrency: 1,
424+
pollIntervalMs: 1000,
425+
timeoutMs: 60_000,
426+
});
427+
428+
expect(outputLinesSent).toBe(outputLineCount);
429+
expect([...byCustomId.entries()]).toEqual(
430+
requests.map((request, index) => [request.custom_id, [index + 1]]),
431+
);
432+
});
433+
434+
it("stops reading batch output after all requested custom IDs are accounted for", async () => {
435+
const outputLineCount = 1024;
436+
let outputLinesSent = 0;
437+
let canceled = false;
438+
const outputResponse = new Response(
439+
new ReadableStream<Uint8Array>({
440+
pull(controller) {
441+
if (outputLinesSent >= outputLineCount) {
442+
controller.close();
443+
return;
444+
}
445+
const line =
446+
outputLinesSent === 0
447+
? {
448+
custom_id: "0",
449+
response: {
450+
status_code: 200,
451+
body: { data: [{ embedding: [1] }] },
452+
},
453+
}
454+
: {
455+
custom_id: `extra-${outputLinesSent}`,
456+
response: {
457+
status_code: 200,
458+
body: { data: [{ embedding: [outputLinesSent] }] },
459+
},
460+
};
461+
controller.enqueue(jsonlEncoder.encode(`${JSON.stringify(line)}\n`));
462+
outputLinesSent += 1;
463+
},
464+
cancel() {
465+
canceled = true;
466+
},
467+
}),
468+
{ status: 200, headers: { "Content-Type": "application/jsonl" } },
469+
);
470+
const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
471+
const url = fetchInputUrl(input);
472+
if (url.endsWith("/files") && init?.method === "POST") {
473+
return jsonResponse({ id: "file-0" });
474+
}
475+
if (url.endsWith("/batches") && init?.method === "POST") {
476+
return jsonResponse({ id: "batch-0", status: "completed", output_file_id: "output-0" });
477+
}
478+
if (url.endsWith("/files/output-0/content")) {
479+
return outputResponse;
480+
}
481+
return new Response("unexpected request", { status: 500 });
482+
});
483+
484+
const byCustomId = await runOpenAiEmbeddingBatches({
485+
openAi: {
486+
baseUrl: "https://openai-compatible.example/v1",
487+
headers: { Authorization: "Bearer test" },
488+
model: "text-embedding-3-small",
489+
fetchImpl,
490+
},
491+
agentId: "main",
492+
requests: [
493+
{
494+
custom_id: "0",
495+
method: "POST",
496+
url: "/v1/embeddings",
497+
body: { model: "text-embedding-3-small", input: "payload" },
498+
},
499+
],
500+
wait: true,
501+
concurrency: 1,
502+
pollIntervalMs: 1000,
503+
timeoutMs: 60_000,
504+
});
505+
506+
expect([...byCustomId.entries()]).toEqual([["0", [1]]]);
507+
expect(canceled).toBe(true);
508+
expect(outputLinesSent).toBeLessThan(outputLineCount);
509+
});
510+
511+
it("bounds batch output file content without buffering the whole response", async () => {
512+
const outputChunkCount = 1024;
513+
let outputChunksSent = 0;
514+
const server = createServer((req, res) => {
515+
const url = req.url ?? "";
516+
if (url === "/v1/files") {
517+
res.writeHead(200, { "Content-Type": "application/json" });
518+
res.end(JSON.stringify({ id: "file-0" }));
519+
return;
520+
}
521+
if (url === "/v1/batches") {
522+
res.writeHead(200, { "Content-Type": "application/json" });
523+
res.end(JSON.stringify({ id: "batch-0", status: "completed", output_file_id: "output-0" }));
524+
return;
525+
}
526+
if (url === "/v1/files/output-0/content") {
527+
res.writeHead(200, { "Content-Type": "application/jsonl" });
528+
const chunkSize = 1024 * 1024;
529+
const writeNext = () => {
530+
if (outputChunksSent >= outputChunkCount) {
531+
res.end();
532+
return;
533+
}
534+
outputChunksSent += 1;
535+
if (res.write(Buffer.alloc(chunkSize))) {
536+
setImmediate(writeNext);
537+
} else {
538+
res.once("drain", writeNext);
539+
}
540+
};
541+
writeNext();
542+
return;
543+
}
544+
res.writeHead(500);
545+
res.end("unexpected request");
546+
});
547+
548+
const port = await listenLoopbackServer(server);
549+
const realFetch = globalThis.fetch.bind(globalThis);
550+
const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
551+
const originalUrl = new URL(fetchInputUrl(input));
552+
const loopbackUrl = new URL(
553+
`${originalUrl.pathname}${originalUrl.search}`,
554+
`http://127.0.0.1:${port}`,
555+
);
556+
return await realFetch(loopbackUrl, init);
557+
});
558+
559+
try {
560+
await expect(
561+
runOpenAiEmbeddingBatches({
562+
openAi: {
563+
baseUrl: "https://openai-compatible.example/v1",
564+
headers: { Authorization: "Bearer test" },
565+
model: "text-embedding-3-small",
566+
fetchImpl,
567+
},
568+
agentId: "main",
569+
requests: [
570+
{
571+
custom_id: "0",
572+
method: "POST",
573+
url: "/v1/embeddings",
574+
body: { model: "text-embedding-3-small", input: "payload" },
575+
},
576+
],
577+
wait: true,
578+
concurrency: 1,
579+
pollIntervalMs: 1000,
580+
timeoutMs: 60_000,
581+
}),
582+
).rejects.toThrow(/openai\.batch-file-content/);
583+
} finally {
584+
await closeServer(server);
585+
}
586+
expect(outputChunksSent).toBeLessThan(outputChunkCount);
587+
});
588+
333589
it("bounds batch resource error bodies without using response.text()", async () => {
334590
const tracked = cancelTrackedResponse(`${"batch status unavailable ".repeat(1024)}tail`, {
335591
status: 400,

0 commit comments

Comments
 (0)