diff --git a/.changeset/fix-openai-batch-binary-jsonl.md b/.changeset/fix-openai-batch-binary-jsonl.md new file mode 100644 index 000000000..4e0fc6c9e --- /dev/null +++ b/.changeset/fix-openai-batch-binary-jsonl.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(js): Support Uint8Array and Buffer JSONL sources in OpenAI batch instrumentation diff --git a/e2e/scenarios/openai-instrumentation/assertions.ts b/e2e/scenarios/openai-instrumentation/assertions.ts index af9f516fe..601f37fca 100644 --- a/e2e/scenarios/openai-instrumentation/assertions.ts +++ b/e2e/scenarios/openai-instrumentation/assertions.ts @@ -864,6 +864,43 @@ export function defineOpenAIInstrumentationAssertions(options: { } }); + test( + "accepts binary JSONL sources for OpenAI Batch traces", + testConfig, + () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan( + events, + "openai-batch-binary-jsonl-operation", + ); + const task = findOpenAISpan(events, operation?.span.id, [ + "openai.batch", + ]); + const children = findOpenAISpans(events, task?.span.id, [ + "Chat Completion", + ]); + + expect(operation).toBeDefined(); + expect(operation?.row.metadata).toMatchObject({ + operation: "batch-binary-jsonl", + }); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(task?.row.metadata).toMatchObject({ + endpoint: "/v1/chat/completions", + input_file_id: "file_binary_batch_e2e_fixture", + provider: "openai", + }); + expect(task?.span.parentIds).toEqual([operation?.span.id ?? ""]); + expect(spanInstrumentationName(task)).toBe("openai"); + expect(children).toHaveLength(EXPECTED_BATCH_OUTPUTS.size); + for (const child of children) { + expect(child.span.parentIds).toEqual([task?.span.id ?? ""]); + expect(spanInstrumentationName(child)).toBe("openai"); + } + validateChatBatchSpans(children); + }, + ); + const scenarioDir = path.dirname(fileURLToPath(options.testFileUrl)); const cassetteMode = process.env.BRAINTRUST_E2E_CASSETTE_MODE; const cassetteEngaged = diff --git a/e2e/scenarios/openai-instrumentation/scenario.impl.mjs b/e2e/scenarios/openai-instrumentation/scenario.impl.mjs index 634d50a06..7c0e2e14c 100644 --- a/e2e/scenarios/openai-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/openai-instrumentation/scenario.impl.mjs @@ -819,6 +819,84 @@ export async function runOpenAIInstrumentationScenario(options) { } }); + await runOperation( + "openai-batch-binary-jsonl-operation", + "batch-binary-jsonl", + async () => { + const batchItems = [ + { + customId: "batch_binary_alpha", + prompt: "Reply with exactly ALPHA.", + response: "ALPHA", + }, + { + customId: "batch_binary_bravo", + prompt: "Reply with exactly BRAVO.", + response: "BRAVO", + }, + { + customId: "batch_binary_charlie", + prompt: "Reply with exactly CHARLIE.", + response: "CHARLIE", + }, + ]; + const input = batchItems + .map((item) => + JSON.stringify({ + custom_id: item.customId, + method: "POST", + url: "/v1/chat/completions", + body: { + model: OPENAI_MODEL, + messages: [{ role: "user", content: item.prompt }], + }, + }), + ) + .join("\n"); + const output = [batchItems[1], batchItems[0]] + .map((item, index) => + JSON.stringify({ + custom_id: item.customId, + response: { + status_code: 200, + body: { + choices: [ + { + index: 0, + finish_reason: "stop", + message: { + role: "assistant", + content: item.response, + }, + }, + ], + usage: { + prompt_tokens: 8 + index, + completion_tokens: 1, + total_tokens: 9 + index, + }, + }, + }, + }), + ) + .join("\n"); + const error = JSON.stringify({ + custom_id: batchItems[2].customId, + error: { + code: "fixture_error", + message: "Batch fixture request failed", + }, + }); + + await completeOpenAIBatchTrace({ + inputFileId: "file_binary_batch_e2e_fixture", + inputFileContent: new TextEncoder().encode(input), + outputFileContent: Promise.resolve(Buffer.from(output)), + errorFileContent: await new Blob([error]).arrayBuffer(), + }); + }, + ); + await runOperation( "openai-embedding-batch-operation", "embedding-batch", diff --git a/js/src/instrumentation/plugins/openai-batch-instrumentation.ts b/js/src/instrumentation/plugins/openai-batch-instrumentation.ts index 636de2d27..378d60865 100644 --- a/js/src/instrumentation/plugins/openai-batch-instrumentation.ts +++ b/js/src/instrumentation/plugins/openai-batch-instrumentation.ts @@ -234,8 +234,17 @@ async function* jsonlRecords( onIssue: (error: Error) => void = () => {}, ): AsyncGenerator { const resolvedFile = await file; - if (typeof resolvedFile === "string") { - for (const line of resolvedFile.split("\n")) { + if ( + typeof resolvedFile === "string" || + resolvedFile instanceof Uint8Array || + resolvedFile instanceof ArrayBuffer || + (typeof Buffer !== "undefined" && Buffer.isBuffer(resolvedFile)) + ) { + const text = + typeof resolvedFile === "string" + ? resolvedFile + : new TextDecoder().decode(resolvedFile); + for (const line of text.split("\n")) { if (!line.trim()) { continue; } diff --git a/js/src/openai-batch-types.ts b/js/src/openai-batch-types.ts index b6052d575..eebd92789 100644 --- a/js/src/openai-batch-types.ts +++ b/js/src/openai-batch-types.ts @@ -1,5 +1,6 @@ export type OpenAIBatchJSONL = | string + | ArrayBuffer | Iterable | AsyncIterable;