wip feat: Add batch evals api

.changeset/durable-batches-evaluate.md

@@ -0,0 +1,5 @@
+---
+"braintrust": minor
+---
+
+feat: Add experimental batch/durable evals API

e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts

@@ -245,7 +245,16 @@ describe.sequential("HarnessAgent instrumentation variants", () => {
             );
             const harnessSpans = findAllSpans(events, "harness");
             expect(harnessSpans).toHaveLength(4);
-            const bashSpans = findAllSpans(events, "bash");
+            // The harness may issue additional bash calls while coordinating a
+            // suspended turn. Assert only the two commands requested from the
+            // agent; coordination calls are not part of this contract.
+            const bashSpans = findAllSpans(events, "bash").filter((span) => {
+              const input = String(span.input);
+              return (
+                input.includes("printf GENERATE_OK") ||
+                input.includes("printf STREAM_OK")
+              );
+            });
             expect(bashSpans).toHaveLength(2);
             for (const bashSpan of bashSpans) {
               expect(bashSpan.span.type).toBe("tool");

e2e/scenarios/durable-eval-webhook/scenario.test.ts

@@ -0,0 +1,84 @@
+import { expect, test } from "vitest";
+import {
+  prepareScenarioDir,
+  resolveScenarioDir,
+  withScenarioHarness,
+} from "../../helpers/scenario-harness";
+import { findAllSpans } from "../../helpers/trace-selectors";
+
+const scenarioDir = await prepareScenarioDir({
+  scenarioDir: resolveScenarioDir(import.meta.url),
+});
+
+test("durable eval collects task and scorer webhook sub-batches", async () => {
+  await withScenarioHarness(
+    async ({ events, runScenarioDir, testRunEvents }) => {
+      await runScenarioDir({ scenarioDir });
+
+      const evalSpans = findAllSpans(testRunEvents(), "eval");
+      const webhookSpans = evalSpans.filter(
+        (event) => event.metadata?.kind === "webhook",
+      );
+      expect(webhookSpans).toHaveLength(3);
+      expect(webhookSpans.map((event) => event.output).sort()).toEqual([
+        2, 4, 6,
+      ]);
+      expect(
+        webhookSpans
+          .map((event) => event.scores)
+          .sort((left, right) =>
+            JSON.stringify(left).localeCompare(JSON.stringify(right)),
+          ),
+      ).toEqual([
+        { batch_exact: 1, exact: 1 },
+        { batch_exact: 1, exact: 1 },
+        { batch_exact: 1, exact: 1 },
+      ]);
+      expect(webhookSpans.map((event) => event.metadata?.durable_eval)).toEqual(
+        [
+          expect.objectContaining({ run_id: expect.any(String) }),
+          expect.objectContaining({ run_id: expect.any(String) }),
+          expect.objectContaining({ run_id: expect.any(String) }),
+        ],
+      );
+
+      const taskSpans = findAllSpans(events(), "task");
+      expect(taskSpans).toHaveLength(3);
+      expect(taskSpans.map((event) => event.output).sort()).toEqual([2, 4, 6]);
+
+      const exactScoreSpans = findAllSpans(events(), "exact");
+      expect(exactScoreSpans).toHaveLength(3);
+      expect(exactScoreSpans.map((event) => event.scores)).toEqual([
+        { exact: 1 },
+        { exact: 1 },
+        { exact: 1 },
+      ]);
+      expect(exactScoreSpans.map((event) => event.metadata?.method)).toEqual([
+        "shared-eval-runtime",
+        "shared-eval-runtime",
+        "shared-eval-runtime",
+      ]);
+
+      const batchScoreSpans = findAllSpans(events(), "batch_exact");
+      expect(batchScoreSpans).toHaveLength(3);
+      expect(batchScoreSpans.map((event) => event.scores)).toEqual([
+        { batch_exact: 1 },
+        { batch_exact: 1 },
+        { batch_exact: 1 },
+      ]);
+      expect(batchScoreSpans.map((event) => event.metadata?.method)).toEqual([
+        "batch-provider",
+        "batch-provider",
+        "batch-provider",
+      ]);
+
+      const classifierSpans = findAllSpans(events(), "quality");
+      expect(classifierSpans).toHaveLength(3);
+      expect(webhookSpans.map((event) => event.row.classifications)).toEqual([
+        { quality: [{ id: "pass", label: "Pass" }] },
+        { quality: [{ id: "pass", label: "Pass" }] },
+        { quality: [{ id: "pass", label: "Pass" }] },
+      ]);
+    },
+  );
+});

e2e/scenarios/durable-eval-webhook/scenario.ts

@@ -0,0 +1,162 @@
+import {
+  BatchScorer,
+  BatchTask,
+  defineDurableEval,
+  DurableEvalMemoryStore,
+} from "braintrust";
+import {
+  getTestRunId,
+  runMain,
+  scopedName,
+} from "../../helpers/scenario-runtime";
+
+async function main() {
+  const testRunId = getTestRunId();
+  const store = new DurableEvalMemoryStore();
+  const jobs = new Map<string, unknown[]>();
+  const webhookCompletion = {
+    mode: "webhook" as const,
+    externalId: (handle: { id: string }) => handle.id,
+  };
+  const task = BatchTask<
+    number,
+    number,
+    number,
+    { testRunId: string; kind: string },
+    Record<string, never>
+  >({
+    workflow(workflow) {
+      const generated = workflow.batch("generate", {
+        batchSize: 2,
+        input: (item) => item.input,
+        async submit(items) {
+          const id = `generate-${jobs.size + 1}`;
+          jobs.set(id, items);
+          return { id };
+        },
+        completion: webhookCompletion,
+        async collect(handle) {
+          const items = (jobs.get(handle.id) ?? []) as Array<{
+            id: string;
+            input: number;
+          }>;
+          return items.map((item) => ({
+            id: item.id,
+            output: item.input * 2,
+          }));
+        },
+      });
+      return workflow.batch("finalize", {
+        needs: { generated },
+        input: (_item, { generated }) => generated,
+        batchSize: 2,
+        async submit(items) {
+          const id = `finalize-${jobs.size + 1}`;
+          jobs.set(id, items);
+          return { id };
+        },
+        completion: webhookCompletion,
+        async collect(handle) {
+          const items = (jobs.get(handle.id) ?? []) as Array<{
+            id: string;
+            input: number;
+          }>;
+          return items.map((item) => ({
+            id: item.id,
+            output: item.input,
+          }));
+        },
+      });
+    },
+  });
+  const scorer = BatchScorer<
+    number,
+    number,
+    number,
+    { testRunId: string; kind: string },
+    { id: string }
+  >({
+    name: "batch_exact",
+    batchSize: 2,
+    async submit(items) {
+      const id = `score-${jobs.size + 1}`;
+      jobs.set(id, items);
+      return { id };
+    },
+    completion: webhookCompletion,
+    async collect(handle) {
+      const items = (jobs.get(handle.id) ?? []) as Array<{
+        id: string;
+        output: number;
+        expected: number;
+      }>;
+      return items.map((item) => ({
+        id: item.id,
+        score: {
+          name: "batch_exact",
+          score: item.output === item.expected ? 1 : 0,
+          metadata: { method: "batch-provider" },
+        },
+      }));
+    },
+  });
+  const definition = defineDurableEval(
+    scopedName("e2e-durable-eval-webhook-project", testRunId),
+    {
+      store,
+      experimentName: scopedName(
+        "e2e-durable-eval-webhook-experiment",
+        testRunId,
+      ),
+      data: [1, 2, 3].map((input) => ({
+        id: `case-${input}`,
+        input,
+        expected: input * 2,
+        metadata: { testRunId, kind: "webhook" },
+      })),
+      task,
+      scores: [
+        function exact({ output, expected }) {
+          return {
+            name: "exact",
+            score: output === expected ? 1 : 0,
+            metadata: { method: "shared-eval-runtime" },
+          };
+        },
+        scorer,
+      ],
+      classifiers: [
+        function quality({ output, expected }) {
+          return {
+            name: "quality",
+            id: output === expected ? "pass" : "fail",
+            label: output === expected ? "Pass" : "Fail",
+          };
+        },
+      ],
+    },
+  );
+
+  const waiting = await definition.start();
+  if (waiting.status !== "waiting" || jobs.size !== 2) {
+    throw new Error("Durable eval did not pause with two webhook batches");
+  }
+
+  let completed = false;
+  const completedJobs = new Set<string>();
+  while (completedJobs.size < jobs.size || !completed) {
+    const externalId = [...jobs.keys()].find((id) => !completedJobs.has(id));
+    if (!externalId) throw new Error("Durable eval stopped before completion");
+    completedJobs.add(externalId);
+    const processed = await definition.processBatchResult({
+      runId: waiting.runId,
+      externalId,
+    });
+    completed = processed.status === "completed";
+  }
+  if ([...jobs.keys()].filter((id) => id.startsWith("score-")).length !== 2) {
+    throw new Error("Batch scorer did not split three cases into two batches");
+  }
+}
+
+runMain(main);