Workflow evaluations

import OpenAI, { toFile } from "openai";
import {
  WorkflowTask,
  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 };

// Sends one prompt to OpenAI as a single-request batch. context holds the
// SDK's run ID and submission ID.
export async function submitRequest(
  prompt: string,
  context: { runId: string; submissionId: string },
): Promise<Submission> {
  const request = {
    // Ties the batch request to this SDK submission.
    custom_id: context.submissionId,
    method: "POST",
    url: "/v1/responses",
    body: { model: "gpt-5-mini", input: prompt },
  };
  // The Batch API reads its requests from an uploaded JSONL file.
  const batchFile = await client.files.create({
    file: await toFile(
      Buffer.from(JSON.stringify(request) + "\n"),
      "requests.jsonl",
    ),
    purpose: "batch",
  });
  const batch = await client.batches.create({
    input_file_id: batchFile.id,
    endpoint: "/v1/responses",
    completion_window: "24h",
    // Lets a webhook handler find the run.
    metadata: { braintrust_run_id: context.runId },
  });
  // 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.
export async function pollSubmission({ batchId }: Submission) {
  const batch = await client.batches.retrieve(batchId);
  switch (batch.status) {
    case "completed":
      return { status: "complete" } as const;
    case "failed":
    case "expired":
    case "cancelled": {
      const errors = JSON.stringify(batch.errors);
      const error = new Error(`Batch ${batch.id}: ${batch.status}, ${errors}`);
      return { status: "failed", error } as const;
    }
    default:
      return { status: "pending" } as const;
  }
}

// Reads the model's answer from the batch output file. The file has one
// JSON line because each batch holds one request.
export async function collectText({ batchId }: Submission): Promise<string> {
  const batch = await client.batches.retrieve(batchId);
  if (batch.error_file_id) {
    const errors = await client.files.content(batch.error_file_id);
    throw new Error(await errors.text());
  }
  if (!batch.output_file_id) {
    throw new Error(`Batch ${batch.id} has no output file`);
  }
  const content = await client.files.content(batch.output_file_id);
  const result = JSON.parse(await content.text());
  const response: OpenAI.Responses.Response = result.response.body;
  return response.output
    .flatMap((item) => (item.type === "message" ? item.content : []))
    .flatMap((part) => (part.type === "output_text" ? [part.text] : []))
    .join("");
}

// 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 WorkflowTask({
      // Called once for each case in data, or once per trial if you set
      // trialCount.
      submit: ({ input }, context) =>
        submitRequest(`${input} Reply with only the city name.`, context),
      // 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: pollSubmission }
          : { mode: "webhook", getExternalId: ({ batchId }) => batchId },
      // Called after the batch completes.
      collect: async (submission) => ({
        output: await collectText(submission),
      }),
    }),
    // Plain scorer functions and WorkflowScorer 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 one batch per case 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 requests that haven't been collected.
    console.log(result.status, result.pending);
  }
  console.log(result.summary);
}

main();