Skip to content
5 changes: 5 additions & 0 deletions .changeset/durable-batches-evaluate.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"braintrust": minor
---

feat: Add experimental batch/durable evals API
11 changes: 10 additions & 1 deletion e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
84 changes: 84 additions & 0 deletions e2e/scenarios/durable-eval-webhook/scenario.test.ts
Original file line number Diff line number Diff line change
@@ -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" }] },
]);
},
);
});
162 changes: 162 additions & 0 deletions e2e/scenarios/durable-eval-webhook/scenario.ts
Original file line number Diff line number Diff line change
@@ -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);
Loading
Loading