Workflow evaluations

import OpenAI, { toFile } from "openai";
import {
  WorkflowBatchTask,
  defineWorkflowEval,
  type EvalScorerArgs,
  type WorkflowEvalStore,
} from "braintrust";

export const client = new OpenAI();

// Data your submit callback returns and the SDK saves between calls.
// The SDK passes it back to the polling and collection callbacks.
// Here it holds the OpenAI batch ID. You choose the fields. Use a type
// alias rather than an interface, because TypeScript can't confirm that an
// interface is JSON-serializable.
export type Submission = { batchId: string };

// One request in a batch. customId is the ID the SDK gives each item in a
// batch. It goes to OpenAI as custom_id, so results can be matched to items.
export type BatchRequest = { customId: string; prompt: string };

// One entry per request, as read from the batch's output and error files.
export type BatchText =
  | { customId: string; text: string }
  | { customId: string; error: unknown };

// The fields this example reads from one line of a batch's output or error
// file.
type BatchResultLine = {
  custom_id: string;
  error?: unknown;
  response?: { status_code: number; body: OpenAI.Responses.Response };
};

// Checks that a parsed line has the custom_id that matches it to an item.
function isBatchResultLine(value: unknown): value is BatchResultLine {
  return (
    typeof value === "object" &&
    value !== null &&
    "custom_id" in value &&
    typeof value.custom_id === "string"
  );
}

// Sends the prompts to OpenAI as one batch.
export async function submitBatch(
  requests: BatchRequest[],
): Promise<Submission> {
  // The Batch API reads its requests from an uploaded JSONL file.
  const lines = requests.map(({ customId, prompt }) =>
    JSON.stringify({
      custom_id: customId,
      method: "POST",
      url: "/v1/responses",
      body: { model: "gpt-5-mini", input: prompt },
    }),
  );
  const batchFile = await client.files.create({
    file: await toFile(Buffer.from(lines.join("\n") + "\n"), "requests.jsonl"),
    purpose: "batch",
  });
  const batch = await client.batches.create({
    input_file_id: batchFile.id,
    endpoint: "/v1/responses",
    completion_window: "24h",
  });
  // The SDK saves this object and passes it to the polling and collect
  // callbacks.
  return { batchId: batch.id };
}

// The polling callback. Called during each poll() until the batch reports
// complete. It reports complete once OpenAI stops working on the batch,
// including when the batch fails or expires, so collect can record the
// requests that have no result as failed and the run can finish.
export async function pollBatch({ batchId }: Submission) {
  const batch = await client.batches.retrieve(batchId);
  switch (batch.status) {
    case "completed":
    case "failed":
    case "expired":
    case "cancelled":
      return { status: "complete" } as const;
    default:
      return { status: "pending" } as const;
  }
}

// Reads each request's answer, or its error, from the batch's output and
// error files.
export async function* collectText({
  batchId,
}: Submission): AsyncGenerator<BatchText> {
  const batch = await client.batches.retrieve(batchId);
  for (const fileId of [batch.output_file_id, batch.error_file_id]) {
    if (!fileId) {
      continue;
    }
    const content = await client.files.content(fileId);
    for (const line of (await content.text()).split("\n")) {
      if (!line.trim()) {
        continue;
      }
      const result: unknown = JSON.parse(line);
      // Stops on a line that can't be matched to an item. The SDK calls
      // collect again on the next poll() or webhook delivery.
      if (!isBatchResultLine(result)) {
        throw new Error(`Unexpected line in batch ${batchId} results: ${line}`);
      }
      if (result.error || result.response?.status_code !== 200) {
        yield {
          customId: result.custom_id,
          error: result.error ?? result.response?.body,
        };
        continue;
      }
      const text = result.response.body.output
        .flatMap((item) => (item.type === "message" ? item.content : []))
        .flatMap((part) => (part.type === "output_text" ? [part.text] : []))
        .join("");
      yield { customId: result.custom_id, text };
    }
  }
}

// A plain scorer function, like the ones you pass to Eval(). The SDK runs
// it on each case's collected output.
function exactMatch({
  output,
  expected,
}: EvalScorerArgs<string, string, string>) {
  return output.trim().toLowerCase() === expected.toLowerCase() ? 1 : 0;
}

// The store and completion mode are parameters so the Redis and webhook
// examples later on this page can reuse this function.
export function createEvaluation(
  store: WorkflowEvalStore,
  completionMode: "poll" | "webhook",
) {
  // The type arguments are the input, output, and expected output types.
  // TypeScript uses them to infer the parameter types of the callbacks below.
  return defineWorkflowEval<string, string, string>("capital-cities-eval", {
    store,
    // Each case needs a stable, unique id so the run can resume.
    data: [
      { id: "france", input: "Capital of France?", expected: "Paris" },
      { id: "japan", input: "Capital of Japan?", expected: "Tokyo" },
    ],
    task: new WorkflowBatchTask({
      // Puts up to 100 cases in each OpenAI batch, so this evaluation's two
      // cases go in one batch. Add trialCount to run each case more than once.
      batching: { maxSize: 100 },
      // Called once for each batch, with one item for each case in it.
      submit: (items) =>
        submitBatch(
          items.map(({ customId, input }) => ({
            customId,
            prompt: `${input} Reply with only the city name.`,
          })),
        ),
      // How the SDK learns that a batch finished: by polling, or from a
      // webhook that the SDK matches to the batch ID getExternalId returns.
      completion:
        completionMode === "poll"
          ? { mode: "poll", poll: pollBatch }
          : { mode: "webhook", getExternalId: ({ batchId }) => batchId },
      // Called after the batch completes. Returns one entry for each item,
      // matched by customId. The SDK records an item without an entry as
      // failed.
      collect: async function* (submission) {
        for await (const entry of collectText(submission)) {
          yield "error" in entry
            ? entry
            : { customId: entry.customId, result: { output: entry.text } };
        }
      },
    }),
    // Plain scorer functions, WorkflowScorer, and WorkflowBatchScorer
    // instances can go here.
    scores: [exactMatch],
  });
}

import { setTimeout as sleep } from "node:timers/promises";
import { WorkflowEvalMemoryStore } from "braintrust";
import { createEvaluation } from "./evaluation";

async function main() {
  // The memory store works in this example because the whole run
  // happens in one process.
  const evaluation = createEvaluation(new WorkflowEvalMemoryStore(), "poll");
  // Submits the cases in one batch and returns right away with status
  // "waiting".
  let result = await evaluation.start();
  console.log(result.runId, result.status);
  // In an application, schedule poll() instead of sleeping in one process.
  while (result.status === "waiting") {
    await sleep(60_000);
    result = await evaluation.poll({ runId: result.runId });
    // pending counts submitted batches that haven't been collected.
    console.log(result.status, result.pending);
  }
  console.log(result.summary);
}

main();