fix(js): Support Uint8Array and Buffer JSONL sources in OpenAI batch instrumentation

.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

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 =

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",

js/src/instrumentation/plugins/openai-batch-instrumentation.ts

@@ -234,8 +234,17 @@ async function* jsonlRecords(
   onIssue: (error: Error) => void = () => {},
 ): AsyncGenerator<unknown> {
   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;
       }

js/src/openai-batch-types.ts

@@ -1,5 +1,6 @@
 export type OpenAIBatchJSONL =
   | string
+  | ArrayBuffer
   | Iterable<unknown>
   | AsyncIterable<unknown>;