diff --git a/.changeset/durable-batches-evaluate.md b/.changeset/durable-batches-evaluate.md new file mode 100644 index 000000000..6067dc216 --- /dev/null +++ b/.changeset/durable-batches-evaluate.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Add experimental batch/durable evals API diff --git a/e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts b/e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts index 87054af1f..c934f6dc9 100644 --- a/e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts +++ b/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"); diff --git a/e2e/scenarios/durable-eval-webhook/scenario.test.ts b/e2e/scenarios/durable-eval-webhook/scenario.test.ts new file mode 100644 index 000000000..4a0b3eae9 --- /dev/null +++ b/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" }] }, + ]); + }, + ); +}); diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts new file mode 100644 index 000000000..a7645a678 --- /dev/null +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -0,0 +1,136 @@ +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(); + const webhookCompletion = { + mode: "webhook" as const, + externalId: (handle: { id: string }) => handle.id, + }; + const task = BatchTask< + number, + number, + number, + { testRunId: string; kind: string }, + Record + >({ + batchSize: 2, + async submit(items) { + const id = `task-${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, + })); + }, + }); + 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(); + 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); diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts new file mode 100644 index 000000000..b13d52c09 --- /dev/null +++ b/js/src/durable-eval.test.ts @@ -0,0 +1,645 @@ +import { describe, expect, test, vi } from "vitest"; +import { configureNode } from "./node/config"; +import { + BatchScorer, + BatchTask, + defineDurableEval, + DurableEvalMemoryStore, + DurableEvalRedisStore, + type DurableBatchScorerItem, + type DurableBatchTaskItem, + type DurableEvalStore, +} from "./durable-eval"; + +configureNode(); + +describe("durable eval stores", () => { + test("memory store copies values on read and write", async () => { + const store = new DurableEvalMemoryStore(); + const value = new Uint8Array([1, 2, 3]); + + await store.write("run", value); + value[0] = 9; + + const firstRead = await store.read("run"); + expect(firstRead).toEqual(new Uint8Array([1, 2, 3])); + firstRead![1] = 9; + expect(await store.read("run")).toEqual(new Uint8Array([1, 2, 3])); + expect(await store.read("missing")).toBeUndefined(); + + const [first, second] = await Promise.all([ + store.getOrSet("claim", new Uint8Array([1])), + store.getOrSet("claim", new Uint8Array([2])), + ]); + expect([first.created, second.created]).toEqual([true, false]); + expect(first.value).toEqual(new Uint8Array([1])); + expect(second.value).toEqual(new Uint8Array([1])); + }); + + test("redis store uses prefixed string operations", async () => { + const values = new Map(); + const client = { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn(async (key: string, value: string) => { + values.set(key, value); + return "OK"; + }), + }; + const store = new DurableEvalRedisStore(client, { + keyPrefix: "evals:", + }); + + await store.write("run", new Uint8Array([0, 255, 1])); + + expect(client.set).toHaveBeenCalledWith("evals:run", "AP8B"); + expect(await store.read("run")).toEqual(new Uint8Array([0, 255, 1])); + expect(client.get).toHaveBeenCalledWith("evals:run"); + expect(await store.read("missing")).toBeUndefined(); + + await expect( + new DurableEvalRedisStore({ + get: async () => 42, + set: async () => "OK", + }).read("invalid"), + ).rejects.toThrow("expected GET to return a string"); + }); + + test("redis store uses node-redis atomic SET options", async () => { + const values = new Map(); + const client = { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn( + async ( + key: string, + value: string, + options?: { NX?: boolean; GET?: boolean }, + ) => { + expect(options).toEqual({ NX: true, GET: true }); + const existing = values.get(key) ?? null; + if (existing === null) values.set(key, value); + return existing; + }, + ), + sendCommand: vi.fn(), + }; + const store = new DurableEvalRedisStore(client); + + await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( + { value: new Uint8Array([1]), created: true }, + ); + await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( + { value: new Uint8Array([1]), created: false }, + ); + }); + + test("redis store uses ioredis atomic SET arguments", async () => { + const values = new Map(); + const client = { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn(async (key: string, value: string, ...options: string[]) => { + expect(options).toEqual(["NX", "GET"]); + const existing = values.get(key) ?? null; + if (existing === null) values.set(key, value); + return existing; + }), + defineCommand: vi.fn(), + }; + const store = new DurableEvalRedisStore(client); + + await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( + { value: new Uint8Array([1]), created: true }, + ); + await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( + { value: new Uint8Array([1]), created: false }, + ); + }); + + test("redis store uses Upstash atomic SET options", async () => { + const values = new Map(); + const client = { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn( + async ( + key: string, + value: string, + options?: { nx?: boolean; get?: boolean }, + ) => { + expect(options).toEqual({ nx: true, get: true }); + const existing = values.get(key) ?? null; + if (existing === null) values.set(key, value); + return existing; + }, + ), + createScript: vi.fn(), + }; + const store = new DurableEvalRedisStore(client); + + await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( + { value: new Uint8Array([1]), created: true }, + ); + await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( + { value: new Uint8Array([1]), created: false }, + ); + }); +}); + +describe("defineDurableEval", () => { + test("runs ordinary tasks and scorers", async () => { + const task = vi.fn((input: number) => input * 2); + const result = await defineDurableEval("local", { + store: new DurableEvalMemoryStore(), + data: [ + { id: "one", input: 1, expected: 2 }, + { id: "two", input: 2, expected: 4 }, + ], + task, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }).start({ noSendLogs: true }); + + expect(result).toMatchObject({ + status: "completed", + summary: { scores: { exact: { score: 1 } } }, + }); + expect(task).toHaveBeenCalledTimes(2); + }); + + test("generates a new run id for every start", async () => { + const durable = defineDurableEval("generated-runs", { + store: new DurableEvalMemoryStore(), + data: [{ input: 1 }], + task: (input) => input, + scores: [() => 1], + }); + + const first = await durable.start({ noSendLogs: true }); + const second = await durable.start({ noSendLogs: true }); + + expect(first.runId).not.toBe(second.runId); + }); + + test("stores run, case, and batch records separately", async () => { + const values = new Map(); + const store: DurableEvalStore = { + async read(key) { + return values.get(key); + }, + async write(key, value) { + values.set(key, value); + }, + async getOrSet(key, value) { + const existing = values.get(key); + if (existing) return { value: existing, created: false }; + values.set(key, value); + return { value, created: true }; + }, + }; + const task = BatchTask< + number, + number, + void, + void, + Record, + { id: string } + >({ + async submit() { + return { id: "provider-job" }; + }, + completion: { + mode: "poll", + async poll() { + return { status: "pending" }; + }, + }, + async collect() { + return []; + }, + }); + + await defineDurableEval("normalized-store", { + store, + data: [ + { id: "one", input: 1 }, + { id: "two", input: 2 }, + ], + task, + }).start({ noSendLogs: true }); + + const decoder = new TextDecoder(); + const records = [...values] + .filter(([key]) => !key.includes("/claims/")) + .map( + ([key, value]) => + [ + key, + JSON.parse(decoder.decode(value)) as Record, + ] as const, + ); + const run = records.find( + ([key]) => !key.includes("/cases/") && !key.includes("/batches/"), + )?.[1]; + expect(run).toMatchObject({ + schemaVersion: 1, + status: "running", + caseIds: ["one:trial:0", "two:trial:0"], + }); + expect(run).not.toHaveProperty("cases"); + expect(run).not.toHaveProperty("batches"); + expect(run).not.toHaveProperty("batchIds"); + expect(records.filter(([key]) => key.includes("/cases/"))).toHaveLength(2); + expect(records.filter(([key]) => key.includes("/batches/"))).toHaveLength( + 1, + ); + expect( + [...values.keys()].filter((key) => key.includes("/claims/")), + ).toHaveLength(1); + }); + + test("polls each existing task and scorer sub-batch once", async () => { + const taskJobs = new Map< + string, + DurableBatchTaskItem>[] + >(); + const scoreJobs = new Map< + string, + DurableBatchScorerItem[] + >(); + const taskPoll = vi.fn(async () => ({ status: "complete" as const })); + const scorePoll = vi.fn(async () => ({ status: "complete" as const })); + + const task = BatchTask< + number, + number, + number, + void, + Record, + { id: string } + >({ + batchSize: 2, + async submit(items) { + const id = `task-${taskJobs.size + 1}`; + taskJobs.set(id, items); + return { id }; + }, + completion: { + mode: "poll", + poll: taskPoll, + }, + async collect(handle) { + return (taskJobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: item.input * 2, + })); + }, + }); + const scorer = BatchScorer({ + name: "exact", + batchSize: 2, + async submit(items) { + const id = `score-${scoreJobs.size + 1}`; + scoreJobs.set(id, items); + return { id }; + }, + completion: { + mode: "poll", + poll: scorePoll, + }, + async collect(handle) { + return (scoreJobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + score: item.output === item.expected ? 1 : 0, + })); + }, + }); + + const store = new DurableEvalMemoryStore(); + const durable = defineDurableEval("polling-batches", { + store, + data: [1, 2, 3].map((input) => ({ + id: `case-${input}`, + input, + expected: input * 2, + })), + task, + scores: [scorer], + }); + const waiting = await durable.start({ noSendLogs: true }); + expect(waiting).toMatchObject({ + status: "waiting", + runId: expect.any(String), + pending: { poll: 2, webhook: 0 }, + }); + const options = { runId: waiting.runId }; + expect(taskJobs.size).toBe(2); + expect(scoreJobs.size).toBe(0); + await expect(durable.status(options)).resolves.toEqual({ + status: "waiting", + runId: waiting.runId, + pending: { poll: 2, webhook: 0 }, + }); + expect(taskJobs.size).toBe(2); + expect(taskPoll).not.toHaveBeenCalled(); + + await expect(durable.poll(options)).resolves.toEqual({ + status: "waiting", + runId: waiting.runId, + pending: { poll: 2, webhook: 0 }, + }); + expect(scoreJobs.size).toBe(2); + expect(taskPoll).toHaveBeenCalledTimes(2); + expect(scorePoll).not.toHaveBeenCalled(); + + const result = await durable.poll(options); + + expect(result).toMatchObject({ + status: "completed", + pending: { poll: 0, webhook: 0 }, + summary: { scores: { exact: { score: 1 } } }, + }); + await expect(durable.status(options)).resolves.toMatchObject({ + status: "completed", + pending: { poll: 0, webhook: 0 }, + summary: { scores: { exact: { score: 1 } } }, + }); + expect(scorePoll).toHaveBeenCalledTimes(2); + expect([...taskJobs.values()].map((items) => items.length)).toEqual([2, 1]); + expect([...scoreJobs.values()].map((items) => items.length)).toEqual([ + 2, 1, + ]); + }); + + test("processes task and scorer webhook batches through one method", async () => { + const store = new DurableEvalMemoryStore(); + const taskJobs = new Map< + string, + DurableBatchTaskItem>[] + >(); + const scoreJobs = new Map< + string, + DurableBatchScorerItem[] + >(); + let taskCollectCount = 0; + let releaseTaskCollect!: () => void; + const taskBatchesCollecting = new Promise((resolve) => { + releaseTaskCollect = resolve; + }); + const task = BatchTask< + number, + number, + number, + void, + Record, + { id: string } + >({ + batchSize: 2, + async submit(items) { + const id = `task-provider-${taskJobs.size + 1}`; + taskJobs.set(id, items); + return { id }; + }, + completion: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect(handle) { + taskCollectCount++; + if (taskCollectCount === 2) releaseTaskCollect(); + await taskBatchesCollecting; + return (taskJobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: item.input * 2, + })); + }, + }); + const scorer = BatchScorer({ + name: "exact", + batchSize: 2, + async submit(items) { + const id = `score-provider-${scoreJobs.size + 1}`; + scoreJobs.set(id, items); + return { id }; + }, + completion: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect(handle) { + return (scoreJobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + score: item.output === item.expected ? 1 : 0, + })); + }, + }); + const durable = defineDurableEval("webhook-batches", { + store, + data: [1, 2, 3].map((input) => ({ + id: `case-${input}`, + input, + expected: input * 2, + })), + task, + scores: [scorer], + }); + + const waiting = await durable.start({ noSendLogs: true }); + expect(waiting).toMatchObject({ + status: "waiting", + runId: expect.any(String), + pending: { poll: 0, webhook: 2 }, + }); + expect(taskJobs.size).toBe(2); + const runId = waiting.runId; + + const taskIds = [...taskJobs.keys()]; + await expect( + durable.processBatchResult({ + runId: "missing-run", + externalId: taskIds[0], + }), + ).rejects.toThrow("Durable eval run missing-run is missing"); + await Promise.all( + taskIds.map((externalId) => + durable.processBatchResult({ runId, externalId }), + ), + ); + expect(scoreJobs.size).toBe(2); + await expect(durable.status({ runId })).resolves.toMatchObject({ + status: "waiting", + pending: { poll: 0, webhook: 2 }, + }); + + const scoreIds = [...scoreJobs.keys()]; + let result; + for (const externalId of scoreIds) { + result = await durable.processBatchResult({ runId, externalId }); + } + expect(result).toMatchObject({ + status: "completed", + pending: { poll: 0, webhook: 0 }, + summary: { scores: { exact: { score: 1 } } }, + }); + }); + + test("claims downstream work once across concurrent webhook deliveries", async () => { + let taskItems: DurableBatchTaskItem< + number, + number, + void, + Record + >[] = []; + let collectCount = 0; + let releaseCollect!: () => void; + const bothCollecting = new Promise((resolve) => { + releaseCollect = resolve; + }); + const task = BatchTask< + number, + number, + number, + void, + Record, + { id: string } + >({ + async submit(items) { + taskItems = items; + return { id: "task-provider" }; + }, + completion: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect() { + collectCount++; + if (collectCount === 2) releaseCollect(); + await bothCollecting; + return taskItems.map((item) => ({ + id: item.id, + output: item.input * 2, + })); + }, + }); + const scoreSubmit = vi.fn(async () => ({ id: "score-provider" })); + const scorer = BatchScorer({ + name: "exact", + submit: scoreSubmit, + completion: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect() { + return []; + }, + }); + const durable = defineDurableEval("concurrent-webhooks", { + store: new DurableEvalMemoryStore(), + data: [{ id: "one", input: 2, expected: 4 }], + task, + scores: [scorer], + }); + const waiting = await durable.start({ noSendLogs: true }); + + await Promise.all([ + durable.processBatchResult({ + runId: waiting.runId, + externalId: "task-provider", + }), + durable.processBatchResult({ + runId: waiting.runId, + externalId: "task-provider", + }), + ]); + + expect(scoreSubmit).toHaveBeenCalledTimes(1); + await expect( + durable.status({ runId: waiting.runId }), + ).resolves.toMatchObject({ + status: "waiting", + pending: { poll: 0, webhook: 1 }, + }); + }); + + test("supports scorer names inherited from Object.prototype", async () => { + const jobs = new Map< + string, + Array<{ id: string; input: number; expected?: number; output?: number }> + >(); + const completion = { + mode: "poll" as const, + async poll() { + return { status: "complete" as const }; + }, + }; + const task = BatchTask>( + { + async submit(items) { + jobs.set("task", items); + return { id: "task" }; + }, + completion, + async collect() { + return (jobs.get("task") ?? []).map((item) => ({ + id: item.id, + output: item.input * 2, + })); + }, + }, + ); + const scorer = BatchScorer({ + name: "__proto__", + async submit(items) { + jobs.set("score", items); + return { id: "score" }; + }, + completion, + async collect() { + return (jobs.get("score") ?? []).map((item) => ({ + id: item.id, + score: item.output === item.expected ? 1 : 0, + })); + }, + }); + const durable = defineDurableEval("prototype-names", { + store: new DurableEvalMemoryStore(), + data: [{ id: "one", input: 2, expected: 4 }], + task, + scores: [scorer], + }); + + const waiting = await durable.start({ noSendLogs: true }); + await expect(durable.poll({ runId: waiting.runId })).resolves.toMatchObject( + { + status: "waiting", + }, + ); + const result = await durable.poll({ runId: waiting.runId }); + expect(result.status).toBe("completed"); + if (result.status !== "completed") throw new Error("Eval did not complete"); + expect(Object.hasOwn(result.summary.scores, "__proto__")).toBe(true); + expect(result.summary.scores.__proto__?.score).toBe(1); + }); + + test("requires stable case ids", async () => { + await expect( + defineDurableEval("missing-ids", { + store: new DurableEvalMemoryStore(), + data: [{ input: "hello" }], + task: BatchTask({ + async submit() { + return { id: "unused" }; + }, + completion: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect() { + return []; + }, + }), + scores: [], + }).start({ noSendLogs: true }), + ).rejects.toThrow("requires id, upsert_id, or caseId"); + }); +}); diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts new file mode 100644 index 000000000..380f385d6 --- /dev/null +++ b/js/src/durable-eval.ts @@ -0,0 +1,2187 @@ +import { + base64ToUint8Array, + makeScorerPropagatedEvent, + SpanTypeAttribute, + uint8ArrayToBase64, +} from "../util/index"; +import { + type EvalParameters, + type InferParameters, + validateParameters, +} from "./eval-parameters"; +import { + _internalInitEvaluatorExperiment, + _internalPrepareEvaluatorClassification, + _internalPrepareEvaluatorScore, + _internalResolveEvaluatorData, + _internalRunEvaluatorTask, + buildLocalSummary as buildEvaluatorLocalSummary, + callEvaluatorData, + classifierName, + type EvalClassifier, + type Evaluator, + type EvaluatorDef, + type EvalResult, + type EvalScorer, + type EvalScorerArgs, + type EvalTask, + type OneOrMoreScores, + runEvaluator, +} from "./framework"; +import iso from "./isomorph"; +import { + type BaseMetadata, + type DefaultMetadataType, + type EvalCase, + type Experiment, + type ExperimentSummary, + NOOP_SPAN, + type Span, + _internalResumeSpan, + _internalStartSpanWithInitialMerge, + logError as logSpanError, + newId, +} from "./logger"; + +const encoder = new TextEncoder(); +const decoder = new TextDecoder(); +const BATCH_TASK_KIND = "braintrust.durable.batch-task"; +const BATCH_SCORER_KIND = "braintrust.durable.batch-scorer"; +const CHECKPOINT_VERSION = 1; +const DEFAULT_BATCH_SIZE = 1_000; + +type JsonPrimitive = string | number | boolean | null; +type JsonValue = JsonPrimitive | JsonValue[] | { [key: string]: JsonValue }; + +/** + * Minimal persistence used to reconnect provider webhooks with submitted + * batches. Each run, case, and batch is stored under its own key. Durable + * evaluations do not require any Braintrust backend changes. + * + * @experimental - The API for this interface is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. + */ +export interface DurableEvalStore { + read(key: string): Promise; + write(key: string, value: Uint8Array): Promise; + /** Atomically stores `value` when `key` is absent and returns its stored value. */ + getOrSet( + key: string, + value: Uint8Array, + ): Promise<{ value: Uint8Array; created: boolean }>; +} + +/** + * Stores durable evaluation state in memory. State is lost when the current + * JavaScript process exits. + * + * @experimental - The API for this class is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. + */ +export class DurableEvalMemoryStore implements DurableEvalStore { + private readonly values = new Map(); + + async read(key: string): Promise { + return this.values.get(key)?.slice(); + } + + async write(key: string, value: Uint8Array): Promise { + this.values.set(key, value.slice()); + } + + async getOrSet(key: string, value: Uint8Array) { + const existing = this.values.get(key); + if (existing) return { value: existing.slice(), created: false }; + this.values.set(key, value.slice()); + return { value: value.slice(), created: true }; + } +} + +/** + * Stores durable evaluation state in Redis using an existing Redis client. + * Values are base64 encoded so only string `GET` and `SET` operations are + * required from the client. Clients from `redis` (node-redis), `ioredis`, and + * `@upstash/redis` can be passed directly. + * + * @experimental - The API for this class is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. + */ +export class DurableEvalRedisStore implements DurableEvalStore { + private readonly keyPrefix: string; + + constructor( + private readonly client: { + get(key: string): Promise; + set(key: string, value: string): Promise; + }, + options: { keyPrefix?: string } = {}, + ) { + this.keyPrefix = options.keyPrefix ?? "braintrust:"; + } + + async read(key: string): Promise { + const value = await this.client.get(`${this.keyPrefix}${key}`); + if (value == null) return undefined; + if (typeof value !== "string") { + throw new Error("DurableEvalRedisStore expected GET to return a string"); + } + return base64ToUint8Array(value); + } + + async write(key: string, value: Uint8Array): Promise { + await this.client.set(`${this.keyPrefix}${key}`, uint8ArrayToBase64(value)); + } + + async getOrSet(key: string, value: Uint8Array) { + const redisKey = `${this.keyPrefix}${key}`; + const encoded = uint8ArrayToBase64(value); + const client = this.client as typeof this.client & { + createScript?: unknown; + defineCommand?: unknown; + sendCommand?: unknown; + }; + const set = client.set as unknown as ( + ...args: unknown[] + ) => Promise; + let setOptions: unknown[]; + if (typeof client.defineCommand === "function") { + setOptions = ["NX", "GET"]; + } else if (typeof client.sendCommand === "function") { + setOptions = [{ NX: true, GET: true }]; + } else if (typeof client.createScript === "function") { + setOptions = [{ nx: true, get: true }]; + } else { + throw new Error( + "DurableEvalRedisStore getOrSet requires a node-redis, ioredis, or @upstash/redis client", + ); + } + const existing = await set.call(client, redisKey, encoded, ...setOptions); + if (existing === null) return { value: value.slice(), created: true }; + if (typeof existing !== "string") { + throw new Error( + "DurableEvalRedisStore expected atomic SET to return a string or null", + ); + } + return { value: base64ToUint8Array(existing), created: false }; + } +} + +interface DurableBatchContext { + runId: string; + batchId: string; +} + +type DurableBatchPoll = + | { status: "pending" } + | { status: "complete" } + | { status: "failed"; error: unknown }; + +type DurableBatchCompletion = + | { + mode: "poll"; + poll( + handle: Handle, + context: DurableBatchContext, + ): Promise; + } + | { + mode: "webhook"; + externalId(handle: Handle, context: DurableBatchContext): string; + }; + +export interface DurableBatchTaskItem< + Input, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +> { + id: string; + input: Input; + expected: Expected; + metadata: Metadata; + tags: string[] | undefined; + parameters: InferParameters; + trialIndex: number; +} + +type DurableBatchTaskResult = + | { + id: string; + output: Output; + metadata?: Metadata; + tags?: string[]; + } + | { id: string; error: unknown }; + +export type DurableBatchScorerItem< + Input, + Output, + Expected, + Metadata extends BaseMetadata, +> = EvalScorerArgs & { + id: string; + trialIndex: number; +}; + +type DurableBatchScorerResult = + | { id: string; score: OneOrMoreScores } + | { id: string; error: unknown }; + +interface DurableBatchProcessor { + batchSize?: number; + submit(items: Item[], context: DurableBatchContext): Promise; + completion: DurableBatchCompletion; + collect(handle: Handle, context: DurableBatchContext): Promise; +} + +interface DurableBatchTask< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, + Handle extends JsonValue, +> { + readonly kind: typeof BATCH_TASK_KIND; + readonly processor: DurableBatchProcessor< + DurableBatchTaskItem, + DurableBatchTaskResult, + Handle + >; +} + +interface DurableBatchScorer< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Handle extends JsonValue, +> { + readonly kind: typeof BATCH_SCORER_KIND; + name: string; + readonly processor: DurableBatchProcessor< + DurableBatchScorerItem, + DurableBatchScorerResult, + Handle + >; +} + +/** + * Defines a task that runs through asynchronous provider batch operations. + * + * @experimental - The API for this function is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. + */ +export function BatchTask< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, + Handle extends JsonValue = JsonValue, +>( + processor: DurableBatchProcessor< + DurableBatchTaskItem, + DurableBatchTaskResult, + Handle + >, +): DurableBatchTask { + return { kind: BATCH_TASK_KIND, processor }; +} + +/** + * Defines a scorer that runs through asynchronous provider batch operations. + * + * @experimental - The API for this function is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. + */ +export function BatchScorer< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Handle extends JsonValue = JsonValue, +>( + processor: DurableBatchProcessor< + DurableBatchScorerItem, + DurableBatchScorerResult, + Handle + > & { name: string }, +): DurableBatchScorer { + return { + kind: BATCH_SCORER_KIND, + name: processor.name, + processor, + }; +} + +type DurableEvaluator< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +> = Omit< + Evaluator, + "task" | "scores" | "timeout" | "signal" | "maxConcurrency" | "update" +> & { + store: DurableEvalStore; + caseId?: ( + datum: EvalCase, + ) => string | Promise; + task: + | EvalTask + | DurableBatchTask< + Input, + Output, + Expected, + Metadata, + Parameters, + JsonValue + >; + scores?: Array< + | EvalScorer + | DurableBatchScorer + >; +}; + +interface DurableEvalStartOptions< + Parameters extends EvalParameters = EvalParameters, +> { + parameters?: InferParameters; + noSendLogs?: boolean; +} + +type DurableBatchResult = { + runId: string; + batchId?: string; + externalId?: string; +}; + +type DurableEvalResult = + | { + status: "waiting"; + runId: string; + pending: { + poll: number; + webhook: number; + }; + } + | { + status: "completed"; + runId: string; + pending: { + poll: 0; + webhook: 0; + }; + summary: ExperimentSummary; + }; + +interface DurableEvalDefinition< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +> { + readonly projectName: string; + readonly evalName: string; + readonly evaluator: DurableEvaluator< + Input, + Output, + Expected, + Metadata, + Parameters + >; + start( + options?: DurableEvalStartOptions, + ): Promise; + status(options: { runId: string }): Promise; + poll(options: { runId: string }): Promise; + processBatchResult(result: DurableBatchResult): Promise; +} + +type DurableCaseRecord = { + id: string; + caseId: string; + trialIndex: number; + datum: JsonValue; + metadata: JsonValue; + tags?: string[]; + taskComplete: boolean; + taskLogged: boolean; + output?: JsonValue; + rootSpan?: string; + scores: Record; + loggedScores: Record; + classifications: Record; + loggedClassifications: Record; +}; + +type DurableCaseBaseRecord = Pick< + DurableCaseRecord, + "id" | "caseId" | "trialIndex" | "datum" | "metadata" | "tags" +>; + +type DurableTaskResultRecord = Pick< + DurableCaseRecord, + "output" | "metadata" | "tags" +> & { taskComplete: true }; + +type DurableTaskLogRecord = Pick & { + taskLogged: true; +}; + +type DurableBatchRecord = { + id: string; + kind: "task" | "score"; + scorerName?: string; + itemIds: string[]; + handle: JsonValue; + externalId?: string; + status: "submitted" | "complete"; +}; + +type DurableRunState = { + schemaVersion: number; + runId: string; + projectName: string; + evalName: string; + experimentName: string; + noSendLogs: boolean; + parameters: JsonValue; + status: "running" | "completed"; + summary?: ExperimentSummary; + cases: DurableCaseRecord[]; + batches: DurableBatchRecord[]; +}; + +type DurableRunRecord = Omit & { + caseIds: string[]; +}; + +class DurableEvalDefinitionImpl< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +> implements DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters +> { + readonly evalName: string; + + constructor( + readonly projectName: string, + readonly evaluator: DurableEvaluator< + Input, + Output, + Expected, + Metadata, + Parameters + >, + ) { + this.evalName = evaluator.experimentName ?? projectName; + } + + start( + options: DurableEvalStartOptions = {}, + ): Promise { + return startDurableEval(this, options); + } + + status(options: { runId: string }): Promise { + return getDurableEvalStatus(this, options); + } + + poll(options: { runId: string }): Promise { + return pollDurableEval(this, options); + } + + processBatchResult(result: DurableBatchResult): Promise { + return processDurableBatchResult(this, result); + } +} + +/* + * Internal usage notes. Keep these out of the public README while + * defineDurableEval() is experimental. + * + * ## Durable evaluations + * + * `defineDurableEval()` runs tasks and scorers through asynchronous provider + * batch APIs. + * `batchSize` splits a dataset into provider-sized sub-batches. A small external + * store connects submitted jobs with later webhook callbacks; it is required on + * the eval definition so every invocation uses the same persistence authority. + * No Braintrust backend changes are required. + * + * Every case needs a stable `id` (or a `caseId` function). + * + * For local or single-process runs, use the built-in memory store. For durable + * deployments, the Redis adapter accepts any existing client with asynchronous + * `get(key)` and `set(key, value)` methods. The adapter itself adds no Redis + * dependency, so install and configure whichever client your application already + * uses. + * + * The following popular clients can be passed directly to + * `DurableEvalRedisStore`: + * + * - [`redis`](https://github.com/redis/node-redis) (node-redis), including the + * lower-level `@redis/client` package + * - [`ioredis`](https://github.com/redis/ioredis) + * - [`@upstash/redis`](https://upstash.com/docs/redis/sdks/ts/deployment), + * including its platform-specific entrypoints + * + * Choose the example for your client: + * + * node-redis (`redis` or `@redis/client`): + * + * ```typescript + * import { createClient } from "redis"; + * import { DurableEvalRedisStore } from "braintrust"; + * + * const nodeRedis = await createClient({ url: process.env.REDIS_URL! }).connect(); + * const redisStore = new DurableEvalRedisStore(nodeRedis); + * ``` + * + * ioredis: + * + * ```typescript + * import Redis from "ioredis"; + * import { DurableEvalRedisStore } from "braintrust"; + * + * const ioRedis = new Redis(process.env.REDIS_URL!); + * const redisStore = new DurableEvalRedisStore(ioRedis); + * ``` + * + * Upstash: + * + * ```typescript + * import { Redis } from "@upstash/redis"; + * import { DurableEvalRedisStore } from "braintrust"; + * + * const upstashRedis = Redis.fromEnv(); + * const redisStore = new DurableEvalRedisStore(upstashRedis); + * ``` + * + * Other clients are compatible when `get(key)` resolves to a string, `null`, or + * `undefined`, and `set(key, value)` accepts a string value. For local testing, + * `new DurableEvalMemoryStore()` requires no external client, but is process-local + * and loses its state when the process exits, so it should not be used to + * reconnect webhooks across serverless invocations. + * + * ```typescript + * import { BatchTask, defineDurableEval } from "braintrust"; + * + * const supportEval = defineDurableEval("Support bot", { + * store: redisStore, + * data: [ + * { + * id: "password-reset", + * input: "How do I reset my password?", + * expected: "Open account settings...", + * }, + * ], + * task: BatchTask({ + * // Each provider job contains at most 500 eval cases. + * batchSize: 500, + * + * // Submit one sub-batch and return a JSON-serializable provider handle. + * async submit(items, context) { + * const batch = await provider.submit({ + * idempotencyKey: context.batchId, + * metadata: { + * durableRunId: context.runId, + * durableBatchId: context.batchId, + * }, + * items, + * }); + * return { id: batch.id }; + * }, + * + * completion: { + * // "webhook" waits for processBatchResult(). Use "poll" with a poll() + * // callback when the provider does not send completion events. + * mode: "webhook", + * externalId: (handle) => handle.id, + * }, + * + * async collect(handle) { + * return (await provider.results(handle.id)).map((item) => ({ + * id: item.id, + * output: item.output, + * })); + * }, + * }), + * scores: [ + * function exact({ output, expected }) { + * return output === expected ? 1 : 0; + * }, + * ], + * }); + * + * const result = await supportEval.start(); + * const { runId } = result; + * ``` + * + * `start()` initializes the run, submits every ready task sub-batch, and returns. + * It never waits in a polling loop. When all task results are available, scoring + * begins. `BatchScorer` uses the same `batchSize`, `submit`, `completion`, and + * array-returning `collect` contract. + * + * ### Polling + * + * Polling adapters report the provider's current status through `completion`: + * + * ```typescript + * completion: { + * mode: "poll", + * async poll(handle) { + * const batch = await provider.getBatch(handle.id); + * if (batch.status === "completed") return { status: "complete" }; + * if (batch.status === "failed") { + * return { status: "failed", error: batch.error }; + * } + * return { status: "pending" }; + * }, + * }, + * ``` + * + * Call `poll()` from a cron, queue worker, or another short-lived invocation. It + * checks every previously submitted polling batch once, collects completed + * results, submits newly ready work, and returns without sleeping: + * + * ```typescript + * const result = await supportEval.poll({ + * runId, + * }); + * + * if (result.status === "waiting" && result.pending.poll > 0) { + * scheduleAnotherPoll(); + * } + * ``` + * + * `start()`, `poll()`, and `processBatchResult()` return the current eval status. + * A waiting result includes the number of submitted batches using each completion + * mode: + * + * ```typescript + * { + * status: "waiting", + * runId, + * pending: { poll: 2, webhook: 1 }, + * } + * ``` + * + * Use `status()` to read the same information without polling providers, + * collecting results, or advancing the evaluation: + * + * ```typescript + * const status = await supportEval.status({ + * runId, + * }); + * ``` + * + * Completed statuses have zero pending batches and include the saved experiment + * summary. They can be read repeatedly without logging the eval again. + * + * ### Webhook processing + * + * When the provider reports that any task or scorer batch completed, fetch and + * store its results through `processBatchResult()`: + * + * ```typescript + * app.post("/webhooks/provider", async (request, response) => { + * const event = request.body; + * const batch = await provider.getBatch(event.batchId); + * const runId = batch.metadata.durableRunId; + * + * const result = await supportEval.processBatchResult({ + * // Returned by start() and saved alongside the provider job. + * runId, + * // The provider's batch ID. The durable eval saved it from submit()'s handle. + * externalId: batch.id, + * // The SDK-generated ID passed to submit(); include it in provider metadata + * // when the webhook cannot provide the external ID used by the handle. + * batchId: batch.metadata?.durableBatchId, + * }); + * + * response.status(result.status === "waiting" ? 202 : 200).end(); + * }); + * ``` + * + * The method accepts either `externalId` or `batchId`. The stored batch locator + * identifies the task or scorer batch, whose `collect()` results are stored + * before the eval advances. Provider failure handling remains the application's + * responsibility for now. + */ + +/** + * Defines a durable evaluation backed by a user-provided store. + * + * @experimental - The API for this function is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. + */ +export function defineDurableEval< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +>( + projectName: string, + evaluator: DurableEvaluator, +): DurableEvalDefinition { + return new DurableEvalDefinitionImpl(projectName, evaluator); +} + +async function startDurableEval< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + options: DurableEvalStartOptions, +): Promise { + const store = definition.evaluator.store; + const runId = newId(); + const key = runKey(definition.projectName, definition.evalName, runId); + const { data } = callEvaluatorData(definition.evaluator.data); + const parameters = await validateParameters( + options.parameters ?? {}, + definition.evaluator.parameters, + ); + const experimentName = + definition.evaluator.experimentName ?? `${definition.evalName}-${runId}`; + const experiment = await _internalInitEvaluatorExperiment( + definition.projectName, + { ...definition.evaluator, data } as unknown as Evaluator< + Input, + Output, + Expected, + Metadata, + Parameters + >, + data, + { + disabled: options.noSendLogs ?? false, + experimentName, + update: true, + }, + ); + if ( + !isBatchTask(definition.evaluator.task) && + !(definition.evaluator.scores ?? []).some(isBatchScorer) + ) { + const result = await runEvaluator( + experiment, + { + ...definition.evaluator, + projectName: definition.projectName, + evalName: definition.evalName, + data, + } as unknown as EvaluatorDef< + Input, + Output, + Expected, + Metadata, + Parameters + >, + { + start: () => undefined, + stop: () => undefined, + increment: () => undefined, + }, + [], + undefined, + parameters, + true, + true, + ); + const state: DurableRunState = { + schemaVersion: CHECKPOINT_VERSION, + runId, + projectName: definition.projectName, + evalName: definition.evalName, + experimentName, + noSendLogs: options.noSendLogs ?? false, + parameters: assertJsonValue(parameters, "eval parameters"), + status: "completed", + summary: result.summary, + cases: [], + batches: [], + }; + await experiment?.flush(); + await writeRunRecord(store, key, state); + return currentStatus(definition, state); + } + const state: DurableRunState = { + schemaVersion: CHECKPOINT_VERSION, + runId, + projectName: definition.projectName, + evalName: definition.evalName, + experimentName, + noSendLogs: options.noSendLogs ?? false, + parameters: assertJsonValue(parameters, "eval parameters"), + status: "running", + cases: await materializeCases(definition, data, experiment), + batches: [], + }; + await writeCaseBaseRecords(store, key, state.cases); + await writeRunRecord(store, key, state); + return advanceDurableEval(definition, state, store, key, experiment); +} + +async function getDurableEvalStatus< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + options: { runId: string }, +): Promise { + const store = definition.evaluator.store; + const state = await readRunState( + definition, + store, + runKey(definition.projectName, definition.evalName, options.runId), + ); + if (!state) throw new Error(`Durable eval run ${options.runId} is missing`); + return currentStatus(definition, state); +} + +async function processDurableBatchResult< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + result: DurableBatchResult, +): Promise { + if (!result.batchId && !result.externalId) { + throw new Error("Batch results require batchId or externalId"); + } + const store = definition.evaluator.store; + const key = runKey(definition.projectName, definition.evalName, result.runId); + const state = await readRunState(definition, store, key); + if (!state) throw new Error(`Durable eval run ${result.runId} is missing`); + const byBatch = result.batchId + ? state.batches.find((candidate) => candidate.id === result.batchId) + : undefined; + const byExternal = result.externalId + ? state.batches.find( + (candidate) => candidate.externalId === result.externalId, + ) + : undefined; + if (byBatch && byExternal && byBatch.id !== byExternal.id) { + throw new Error("batchId and externalId identify different batches"); + } + const batch = byBatch ?? byExternal; + if (!batch) throw new Error("No submitted batch matches this result"); + if (batch.status !== "complete") { + const records = await collectBatch(definition, state, batch); + batch.status = "complete"; + await writeCaseRecords(store, key, records); + await writeBatchRecords(store, key, [batch]); + } + return advanceDurableEval( + definition, + (await readRunState(definition, store, key))!, + store, + key, + ); +} + +async function pollDurableEval< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + options: { runId: string }, +): Promise { + const store = definition.evaluator.store; + const key = runKey( + definition.projectName, + definition.evalName, + options.runId, + ); + const state = await readRunState(definition, store, key); + if (!state) throw new Error(`Durable eval run ${options.runId} is missing`); + + const batches = state.batches.filter((batch) => { + if (batch.status === "complete") return false; + return ( + processorForStage(definition, batch.kind, batch.scorerName).completion + .mode === "poll" + ); + }); + const results = await Promise.all( + batches.map(async (batch) => ({ + batch, + result: await ( + processorForStage(definition, batch.kind, batch.scorerName) + .completion as Extract< + DurableBatchCompletion, + { mode: "poll" } + > + ).poll(batch.handle, { + runId: state.runId, + batchId: batch.id, + }), + })), + ); + const changedCases = new Map(); + const changedBatches: DurableBatchRecord[] = []; + for (const { batch, result } of results) { + if (result.status === "failed") throw asError(result.error); + if (result.status !== "complete") continue; + for (const record of await collectBatch(definition, state, batch)) { + changedCases.set(record.id, record); + } + batch.status = "complete"; + changedBatches.push(batch); + } + if (changedBatches.length > 0) { + await writeCaseRecords(store, key, [...changedCases.values()]); + await writeBatchRecords(store, key, changedBatches); + } + const currentState = + changedBatches.length > 0 + ? (await readRunState(definition, store, key))! + : state; + return advanceDurableEval(definition, currentState, store, key); +} + +async function openDurableExperiment( + definition: DurableEvalDefinition, + state: DurableRunState, +) { + const data: EvalCase[] = []; + return await _internalInitEvaluatorExperiment( + definition.projectName, + { ...definition.evaluator, data } as unknown as Evaluator< + unknown, + unknown, + unknown, + BaseMetadata, + EvalParameters + >, + data, + { + disabled: state.noSendLogs, + experimentName: state.experimentName, + update: true, + }, + ); +} + +async function advanceDurableEval< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + state: DurableRunState, + store: DurableEvalStore, + key: string, + existingExperiment?: Experiment | null, +): Promise { + if (state.status === "completed") return currentStatus(definition, state); + const experiment = + existingExperiment === undefined + ? await openDurableExperiment(definition, state) + : existingExperiment; + await runTaskStage(definition, state, store, key, experiment); + await logCompletedTasks(definition, state, store, key, experiment); + state = (await readRunState(definition, store, key)) ?? state; + if ( + state.cases.some((record) => !record.taskComplete || !record.taskLogged) + ) { + return currentStatus(definition, state); + } + + await runScoreStages(definition, state, store, key, experiment); + state = (await readRunState(definition, store, key)) ?? state; + const scorerNames = resolveScorers(definition.evaluator.scores ?? []).map( + ({ name }) => name, + ); + const classifierNames = (definition.evaluator.classifiers ?? []).map( + classifierName, + ); + if ( + state.cases.some( + (record) => + scorerNames.some( + (name) => + !Object.hasOwn(record.scores, name) || + !Object.hasOwn(record.loggedScores, name), + ) || + classifierNames.some( + (name) => !Object.hasOwn(record.loggedClassifications, name), + ), + ) + ) { + return currentStatus(definition, state); + } + + if (!(await claimAction(store, key, "finish"))) { + const latest = await readRunState(definition, store, key); + return currentStatus(definition, latest ?? state); + } + state.summary = await finishExperiment(definition, state, experiment); + state.status = "completed"; + await writeRunRecord(store, key, state); + return currentStatus(definition, state); +} + +function currentStatus( + definition: DurableEvalDefinition, + state: DurableRunState, +): DurableEvalResult { + if (state.status === "completed") { + if (!state.summary) { + throw new Error(`Durable eval run ${state.runId} has no saved summary`); + } + return { + status: "completed", + runId: state.runId, + pending: { poll: 0, webhook: 0 }, + summary: state.summary, + }; + } + const pending = { poll: 0, webhook: 0 }; + for (const batch of state.batches) { + if (batch.status === "complete") continue; + pending[ + processorForStage(definition, batch.kind, batch.scorerName).completion + .mode + ]++; + } + return { status: "waiting", runId: state.runId, pending }; +} + +async function startCaseRoot( + definition: DurableEvalDefinition, + state: DurableRunState, + record: DurableCaseRecord, + experiment: Experiment | null, +): Promise { + if (!experiment) return NOOP_SPAN; + const datum = record.datum as EvalCase; + return _internalStartSpanWithInitialMerge({ + ...(definition.evaluator.state + ? { state: definition.evaluator.state } + : {}), + parent: await experiment.export(), + name: "eval", + spanId: deterministicId(`${state.runId}:${record.id}:span`), + spanAttributes: { type: SpanTypeAttribute.EVAL }, + event: { + id: deterministicId(`${state.runId}:${record.id}:row`), + input: datum.input, + expected: "expected" in datum ? datum.expected : undefined, + tags: datum.tags, + origin: datum.origin, + }, + }); +} + +async function logTaskResult( + definition: DurableEvalDefinition, + state: DurableRunState, + record: DurableCaseRecord, + experiment: Experiment | null, + task?: EvalTask, +) { + const datum = record.datum as EvalCase; + const root = await startCaseRoot(definition, state, record, experiment); + try { + if (task) { + const result = await root.traced( + (span) => + _internalRunEvaluatorTask( + task, + datum, + record.trialIndex, + state.parameters as Record, + span, + ), + { + name: "task", + spanId: deterministicId(`${state.runId}:${record.id}:task`), + spanAttributes: { type: SpanTypeAttribute.TASK }, + event: { input: datum.input }, + }, + ); + record.output = assertJsonValue( + result.output, + `task output for ${record.caseId}`, + ); + record.metadata = assertJsonValue(result.metadata, "task metadata"); + record.tags = result.tags; + record.taskComplete = true; + } else { + await root.traced((span) => span.log({ output: record.output }), { + name: "task", + spanId: deterministicId(`${state.runId}:${record.id}:task`), + spanAttributes: { type: SpanTypeAttribute.TASK }, + event: { input: datum.input }, + }); + } + root.log({ + output: record.output, + expected: "expected" in datum ? datum.expected : undefined, + metadata: { + ...(record.metadata as Record), + durable_eval: { + run_id: state.runId, + case_id: record.caseId, + trial_index: record.trialIndex, + }, + }, + tags: record.tags, + }); + record.rootSpan = await root.export(); + record.taskLogged = true; + } catch (error) { + logSpanError(root, error); + throw error; + } finally { + root.end(); + } +} + +async function logCompletedTasks( + definition: DurableEvalDefinition, + state: DurableRunState, + store: DurableEvalStore, + key: string, + experiment: Experiment | null, +) { + const changed: DurableCaseRecord[] = []; + for (const record of state.cases) { + if (!record.taskComplete || record.taskLogged) continue; + if (!(await claimAction(store, key, "task-log", record.id))) continue; + await logTaskResult(definition, state, record, experiment); + changed.push(record); + } + if (changed.length > 0) { + await experiment?.flush(); + await writeCaseRecords(store, key, changed); + } +} + +async function runTaskStage< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + state: DurableRunState, + store: DurableEvalStore, + key: string, + experiment: Experiment | null, +) { + if (isBatchTask(definition.evaluator.task)) { + await ensureBatches(definition, state, store, key, "task"); + return; + } + + const task = definition.evaluator.task as EvalTask< + Input, + Output, + Expected, + Metadata, + Parameters + >; + const changed: DurableCaseRecord[] = []; + for (const record of state.cases) { + if (record.taskComplete) continue; + if (!(await claimAction(store, key, "task", record.id))) continue; + await logTaskResult(definition, state, record, experiment, task); + changed.push(record); + } + if (changed.length > 0) { + await experiment?.flush(); + await writeCaseRecords(store, key, changed); + } +} + +async function runScoreStages< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + state: DurableRunState, + store: DurableEvalStore, + key: string, + experiment: Experiment | null, +) { + const scorers = resolveScorers(definition.evaluator.scores ?? []); + const changed = new Map(); + const persistChangedCases = async () => { + if (changed.size === 0) return; + await experiment?.flush(); + await writeCaseRecords(store, key, [...changed.values()]); + changed.clear(); + }; + for (const { name, scorer } of scorers) { + if (isBatchScorer(scorer)) { + for (const record of state.cases) { + if ( + Object.hasOwn(record.scores, name) && + !Object.hasOwn(record.loggedScores, name) + ) { + if (!(await claimAction(store, key, "score-log", record.id, name))) { + continue; + } + await evaluateAndLogScore( + definition, + state, + record, + name, + experiment, + ); + changed.set(record.id, record); + } + } + await persistChangedCases(); + await ensureBatches(definition, state, store, key, "score", name); + continue; + } + for (const record of state.cases) { + if (Object.hasOwn(record.loggedScores, name)) continue; + if (!(await claimAction(store, key, "score", record.id, name))) continue; + await evaluateAndLogScore( + definition, + state, + record, + name, + experiment, + scorer, + ); + changed.set(record.id, record); + } + } + + for (const [index, classifier] of ( + definition.evaluator.classifiers ?? [] + ).entries()) { + const name = classifierName(classifier, index); + for (const record of state.cases) { + if (Object.hasOwn(record.loggedClassifications, name)) continue; + if (!(await claimAction(store, key, "classification", record.id, name))) { + continue; + } + await evaluateAndLogClassification( + definition, + state, + record, + name, + classifier, + experiment, + ); + changed.set(record.id, record); + } + } + await persistChangedCases(); +} + +function scorerArgs(record: DurableCaseRecord) { + const datum = record.datum as EvalCase; + return { + ...datum, + metadata: record.metadata, + output: record.output, + } as EvalScorerArgs; +} + +function resumeCaseRoot( + definition: DurableEvalDefinition, + record: DurableCaseRecord, + experiment: Experiment | null, +) { + if (!experiment) return NOOP_SPAN; + if (!record.rootSpan) { + throw new Error(`Durable eval case ${record.caseId} has no root span`); + } + return _internalResumeSpan({ + exported: record.rootSpan, + state: definition.evaluator.state, + }); +} + +async function evaluateAndLogScore( + definition: DurableEvalDefinition, + state: DurableRunState, + record: DurableCaseRecord, + name: string, + experiment: Experiment | null, + scorer?: EvalScorer, +) { + const root = resumeCaseRoot(definition, record, experiment); + try { + const rootExport = await root.export(); + const prepared = await root.traced( + async (span) => { + const value = scorer + ? await scorer(scorerArgs(record)) + : (record.scores[name] as OneOrMoreScores); + if (scorer) { + record.scores[name] = assertJsonValue(value, `scorer ${name} output`); + } + const result = _internalPrepareEvaluatorScore(value, name); + if (result.results !== null) { + span.log({ + output: result.output, + metadata: result.metadata, + scores: result.scores, + }); + } + return result; + }, + { + name, + spanId: deterministicId(`${state.runId}:${record.id}:score:${name}`), + spanAttributes: { + type: SpanTypeAttribute.SCORE, + purpose: "scorer", + }, + propagatedEvent: makeScorerPropagatedEvent(rootExport || undefined), + event: { input: scorerArgs(record) }, + }, + ); + if (prepared.scores) root.log({ scores: prepared.scores }); + record.loggedScores[name] = true; + } catch (error) { + logSpanError(root, error); + throw error; + } finally { + root.end(); + } +} + +async function evaluateAndLogClassification( + definition: DurableEvalDefinition, + state: DurableRunState, + record: DurableCaseRecord, + name: string, + classifier: EvalClassifier, + experiment: Experiment | null, +) { + const root = resumeCaseRoot(definition, record, experiment); + try { + const rootExport = await root.export(); + const prepared = await root.traced( + async (span) => { + const value = await classifier(scorerArgs(record)); + record.classifications[name] = assertJsonValue( + value, + `classifier ${name} output`, + ); + const result = _internalPrepareEvaluatorClassification(value, name); + if (result.results !== null) { + span.log({ output: result.output, metadata: result.metadata }); + } + return result; + }, + { + name, + spanId: deterministicId( + `${state.runId}:${record.id}:classification:${name}`, + ), + spanAttributes: { + type: SpanTypeAttribute.CLASSIFIER, + purpose: "scorer", + }, + propagatedEvent: makeScorerPropagatedEvent(rootExport || undefined), + event: { input: scorerArgs(record) }, + }, + ); + if (prepared.classifications) { + root.log({ classifications: prepared.classifications }); + } + record.loggedClassifications[name] = true; + } catch (error) { + logSpanError(root, error); + throw error; + } finally { + root.end(); + } +} + +async function ensureBatches( + definition: DurableEvalDefinition, + state: DurableRunState, + store: DurableEvalStore, + key: string, + kind: "task" | "score", + scorerName?: string, +) { + const processor = processorForStage(definition, kind, scorerName); + const plans = plannedBatches( + definition, + state.runId, + state.cases.map(({ id }) => id), + ); + const casesById = new Map(state.cases.map((record) => [record.id, record])); + for (const plan of plans) { + if (plan.kind !== kind || plan.scorerName !== scorerName) continue; + if (state.batches.some(({ id }) => id === plan.id)) continue; + const records = plan.itemIds.map((id) => casesById.get(id)!); + const ready = records.every((record) => + kind === "task" + ? !record.taskComplete + : !Object.hasOwn(record.scores, scorerName!), + ); + if (!ready) continue; + const batchId = plan.id; + const claim = await store.getOrSet( + claimRecordKey(key, "batch", batchId), + encoder.encode(batchId), + ); + if (!claim.created) continue; + const context = { runId: state.runId, batchId }; + const items = records.map((record) => + kind === "task" + ? taskBatchItem(record, state.parameters) + : scorerBatchItem(record), + ); + const handle = assertJsonValue( + await processor.submit(items, context), + `handle for batch ${batchId}`, + ); + const externalId = + processor.completion.mode === "webhook" + ? processor.completion.externalId(handle, context) + : undefined; + if (externalId !== undefined && !externalId.trim()) { + throw new Error(`Batch ${batchId} produced an empty externalId`); + } + const batch: DurableBatchRecord = { + id: batchId, + kind, + scorerName, + itemIds: records.map((record) => record.id), + handle, + externalId, + status: "submitted", + }; + state.batches.push(batch); + await writeBatchRecords(store, key, [batch]); + } +} + +async function collectBatch( + definition: DurableEvalDefinition, + state: DurableRunState, + batch: DurableBatchRecord, +) { + const processor = processorForStage(definition, batch.kind, batch.scorerName); + const context = { runId: state.runId, batchId: batch.id }; + const results = await processor.collect(batch.handle, context); + if (!Array.isArray(results)) { + throw new Error(`collect for batch ${batch.id} must return an array`); + } + const expectedIds = new Set(batch.itemIds); + const seen = new Set(); + const records: DurableCaseRecord[] = []; + for (const result of results) { + const id = resultItemId(result); + if (!expectedIds.has(id)) { + throw new Error(`Batch ${batch.id} returned unknown item ${id}`); + } + if (seen.has(id)) { + throw new Error(`Batch ${batch.id} returned item ${id} more than once`); + } + seen.add(id); + if ("error" in result) throw asError(result.error); + const record = state.cases.find((candidate) => candidate.id === id)!; + records.push(record); + if (batch.kind === "task") { + record.output = assertJsonValue( + result.output, + `task output for item ${id}`, + ); + if ("metadata" in result && result.metadata !== undefined) { + record.metadata = assertJsonValue( + result.metadata, + `metadata for ${id}`, + ); + } + if ("tags" in result && result.tags !== undefined) + record.tags = result.tags; + record.taskComplete = true; + } else { + record.scores[batch.scorerName!] = assertJsonValue( + result.score, + `score output for item ${id}`, + ); + } + } + const missing = batch.itemIds.filter((id) => !seen.has(id)); + if (missing.length > 0) { + throw new Error( + `Batch ${batch.id} did not return results for: ${missing.join(", ")}`, + ); + } + return records; +} + +function processorForStage( + definition: DurableEvalDefinition, + kind: "task" | "score", + scorerName?: string, +): DurableBatchProcessor { + if (kind === "task") { + if (!isBatchTask(definition.evaluator.task)) { + throw new Error("Definition no longer contains the batch task"); + } + return definition.evaluator.task.processor as DurableBatchProcessor< + any, + any, + JsonValue + >; + } + const scorer = resolveScorers(definition.evaluator.scores ?? []).find( + ({ name }) => name === scorerName, + )?.scorer; + if (!isBatchScorer(scorer)) { + throw new Error(`Definition no longer contains scorer ${scorerName}`); + } + return scorer.processor as DurableBatchProcessor; +} + +async function materializeCases< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + data: Evaluator["data"], + experiment: Experiment | null, +): Promise { + const evaluator = definition.evaluator; + const iterable = await _internalResolveEvaluatorData( + { + data, + projectName: definition.projectName, + projectId: evaluator.projectId, + state: evaluator.state, + }, + experiment, + ); + const records: DurableCaseRecord[] = []; + const seen = new Set(); + for await (const datum of iterable) { + const caseId = + datum.id ?? + datum.upsert_id ?? + (evaluator.caseId + ? await evaluator.caseId(datum as EvalCase) + : undefined); + if (!caseId) { + throw new Error( + "Every durable eval case requires id, upsert_id, or caseId", + ); + } + if (seen.has(caseId)) + throw new Error(`Duplicate durable eval case id: ${caseId}`); + seen.add(caseId); + const trialCount = datum.trialCount ?? evaluator.trialCount ?? 1; + if (!Number.isInteger(trialCount) || trialCount < 1) { + throw new Error(`Invalid trialCount for durable eval case ${caseId}`); + } + for (let trialIndex = 0; trialIndex < trialCount; trialIndex++) { + records.push({ + id: `${caseId}:trial:${trialIndex}`, + caseId, + trialIndex, + datum: assertJsonValue(datum, `case ${caseId}`), + metadata: assertJsonValue( + "metadata" in datum ? datum.metadata : {}, + `metadata for ${caseId}`, + ), + tags: datum.tags, + taskComplete: false, + taskLogged: false, + scores: Object.create(null), + loggedScores: Object.create(null), + classifications: Object.create(null), + loggedClassifications: Object.create(null), + }); + } + } + return records; +} + +function taskBatchItem(record: DurableCaseRecord, parameters: JsonValue) { + const datum = record.datum as EvalCase; + return { + id: record.id, + input: datum.input, + expected: "expected" in datum ? datum.expected : undefined, + metadata: record.metadata, + tags: record.tags, + parameters, + trialIndex: record.trialIndex, + }; +} + +function scorerBatchItem(record: DurableCaseRecord) { + const datum = record.datum as EvalCase; + return { + id: record.id, + input: datum.input, + output: record.output, + expected: "expected" in datum ? datum.expected : undefined, + metadata: record.metadata, + tags: record.tags, + trialIndex: record.trialIndex, + }; +} + +async function finishExperiment( + definition: DurableEvalDefinition, + state: DurableRunState, + experiment: Experiment | null, +) { + const scorerNames = resolveScorers(definition.evaluator.scores ?? []).map( + ({ name }) => name, + ); + const results = state.cases.map((record) => { + const datum = record.datum as EvalCase; + const scores = Object.fromEntries( + scorerNames.flatMap((name) => + Object.entries( + _internalPrepareEvaluatorScore( + record.scores[name] as OneOrMoreScores, + name, + ).scores ?? {}, + ), + ), + ); + const classifications = Object.fromEntries( + Object.entries(record.classifications).flatMap(([name, value]) => + Object.entries( + _internalPrepareEvaluatorClassification(value as never, name) + .classifications ?? {}, + ), + ), + ); + return { + ...datum, + output: record.output, + metadata: record.metadata, + tags: record.tags, + scores, + error: undefined, + ...(Object.keys(classifications).length > 0 ? { classifications } : {}), + } as EvalResult; + }); + if (!experiment) { + return buildEvaluatorLocalSummary( + { + ...definition.evaluator, + projectName: definition.projectName, + evalName: state.experimentName, + } as unknown as EvaluatorDef, + results, + ); + } + await experiment.flush(); + let comparisonExperimentId = definition.evaluator.baseExperimentId; + if (!comparisonExperimentId) { + try { + comparisonExperimentId = await experiment._getBaseExperimentId(); + } catch { + comparisonExperimentId = undefined; + } + } + return await experiment.summarize({ + summarizeScores: definition.evaluator.summarizeScores, + ...(comparisonExperimentId ? { comparisonExperimentId } : {}), + }); +} + +function resolveScorers( + scorers: Array< + | EvalScorer + | DurableBatchScorer + >, +) { + return scorers.map((scorer, index) => ({ + name: isBatchScorer(scorer) + ? scorer.name + : scorer.name || `scorer_${index}`, + scorer, + })); +} + +function runKey(projectName: string, evalName: string, runId: string) { + return `durable-eval/v1/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; +} + +function encodedKeyPart(value: string) { + return uint8ArrayToBase64(encoder.encode(value)) + .replaceAll("+", "-") + .replaceAll("/", "_") + .replace(/=+$/, ""); +} + +function caseRecordKey( + key: string, + caseId: string, + kind: + | "base" + | "task" + | "task-log" + | "score" + | "score-log" + | "classification" + | "classification-log", + ...names: string[] +) { + const suffix = names.map(encodedKeyPart).join("/"); + return `${key}/cases/${encodedKeyPart(caseId)}/${kind}${suffix ? `/${suffix}` : ""}`; +} + +function batchRecordKey(key: string, batchId: string) { + return `${key}/batches/${encodedKeyPart(batchId)}`; +} + +function claimRecordKey(key: string, kind: string, ...parts: string[]) { + const identity = stableStringify([kind, parts]); + return `${key}/claims/${contentVersion(encoder.encode(identity))}`; +} + +async function claimAction( + store: DurableEvalStore, + key: string, + kind: string, + ...parts: string[] +) { + return ( + await store.getOrSet( + claimRecordKey(key, kind, ...parts), + encoder.encode("claimed"), + ) + ).created; +} + +type DurableBatchPlan = Omit< + DurableBatchRecord, + "handle" | "externalId" | "status" +>; + +function plannedBatches( + definition: DurableEvalDefinition, + runId: string, + caseIds: string[], +) { + const stages: Array<{ kind: "task" | "score"; scorerName?: string }> = []; + if (isBatchTask(definition.evaluator.task)) stages.push({ kind: "task" }); + for (const { name, scorer } of resolveScorers( + definition.evaluator.scores ?? [], + )) { + if (isBatchScorer(scorer)) { + stages.push({ kind: "score", scorerName: name }); + } + } + const plans: DurableBatchPlan[] = []; + for (const { kind, scorerName } of stages) { + const batchSize = + processorForStage(definition, kind, scorerName).batchSize ?? + DEFAULT_BATCH_SIZE; + if (!Number.isInteger(batchSize) || batchSize < 1) { + throw new Error( + `Invalid batchSize for ${scorerName ?? "task"}: ${batchSize}`, + ); + } + for (let offset = 0; offset < caseIds.length; offset += batchSize) { + const itemIds = caseIds.slice(offset, offset + batchSize); + plans.push({ + id: deterministicId( + stableStringify([runId, kind, scorerName, itemIds]), + ), + kind, + scorerName, + itemIds, + }); + } + } + return plans; +} + +async function readCaseRecord( + definition: DurableEvalDefinition, + store: DurableEvalStore, + key: string, + id: string, +) { + const scorers = resolveScorers(definition.evaluator.scores ?? []); + const classifiers = (definition.evaluator.classifiers ?? []).map( + classifierName, + ); + const [ + base, + task, + taskLog, + scoreValues, + scoreLogValues, + classificationValues, + classificationLogValues, + ] = await Promise.all([ + readJson(store, caseRecordKey(key, id, "base")), + readJson(store, caseRecordKey(key, id, "task")), + readJson(store, caseRecordKey(key, id, "task-log")), + Promise.all( + scorers.map(async ({ name }) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "score", name), + ), + })), + ), + Promise.all( + scorers.map(async ({ name }) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "score-log", name), + ), + })), + ), + Promise.all( + classifiers.map(async (name) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "classification", name), + ), + })), + ), + Promise.all( + classifiers.map(async (name) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "classification-log", name), + ), + })), + ), + ]); + if (!base) throw new Error(`Durable eval case ${id} is missing`); + const scores: Record = Object.create(null); + for (const { name, value } of scoreValues) { + if (value !== undefined) scores[name] = value; + } + const loggedScores: Record = Object.create(null); + for (const { name, value } of scoreLogValues) { + if (value) loggedScores[name] = true; + } + const classifications: Record = Object.create(null); + for (const { name, value } of classificationValues) { + if (value !== undefined) classifications[name] = value; + } + const loggedClassifications: Record = Object.create(null); + for (const { name, value } of classificationLogValues) { + if (value) loggedClassifications[name] = true; + } + return { + ...base, + metadata: task?.metadata ?? base.metadata, + tags: task ? task.tags : base.tags, + taskComplete: task !== undefined, + taskLogged: taskLog !== undefined, + output: task?.output, + rootSpan: taskLog?.rootSpan, + scores, + loggedScores, + classifications, + loggedClassifications, + } satisfies DurableCaseRecord; +} + +async function readRunState( + definition: DurableEvalDefinition, + store: DurableEvalStore, + key: string, +) { + const record = await readJson(store, key); + if (!record) return undefined; + const plans = plannedBatches(definition, record.runId, record.caseIds); + const [cases, batchRecords] = await Promise.all([ + Promise.all( + record.caseIds.map((id) => readCaseRecord(definition, store, key, id)), + ), + Promise.all( + plans.map(async ({ id }) => { + return readJson(store, batchRecordKey(key, id)); + }), + ), + ]); + const batches = batchRecords.filter( + (value): value is DurableBatchRecord => value !== undefined, + ); + const { caseIds: _caseIds, ...state } = record; + return { ...state, cases, batches }; +} + +async function writeRunRecord( + store: DurableEvalStore, + key: string, + state: DurableRunState, +) { + const { cases, batches: _batches, ...record } = state; + await writeJson(store, key, { + ...record, + caseIds: cases.map(({ id }) => id), + } satisfies DurableRunRecord); +} + +async function writeCaseBaseRecords( + store: DurableEvalStore, + key: string, + records: DurableCaseRecord[], +) { + await Promise.all( + records.map(({ id, caseId, trialIndex, datum, metadata, tags }) => + writeJson(store, caseRecordKey(key, id, "base"), { + id, + caseId, + trialIndex, + datum, + metadata, + tags, + } satisfies DurableCaseBaseRecord), + ), + ); +} + +async function writeCaseRecords( + store: DurableEvalStore, + key: string, + records: DurableCaseRecord[], +) { + const writes: Promise[] = []; + for (const record of records) { + if (record.taskComplete) { + writes.push( + writeJson(store, caseRecordKey(key, record.id, "task"), { + output: record.output, + metadata: record.metadata, + tags: record.tags, + taskComplete: true, + } satisfies DurableTaskResultRecord), + ); + } + if (record.taskLogged) { + writes.push( + writeJson(store, caseRecordKey(key, record.id, "task-log"), { + rootSpan: record.rootSpan, + taskLogged: true, + } satisfies DurableTaskLogRecord), + ); + } + for (const [name, value] of Object.entries(record.scores)) { + writes.push( + writeJson(store, caseRecordKey(key, record.id, "score", name), value), + ); + } + for (const name of Object.keys(record.loggedScores)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "score-log", name), + true, + ), + ); + } + for (const [name, value] of Object.entries(record.classifications)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "classification", name), + value, + ), + ); + } + for (const name of Object.keys(record.loggedClassifications)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "classification-log", name), + true, + ), + ); + } + } + await Promise.all(writes); +} + +async function writeBatchRecords( + store: DurableEvalStore, + key: string, + records: DurableBatchRecord[], +) { + await Promise.all( + records.map((record) => + writeJson(store, batchRecordKey(key, record.id), record), + ), + ); +} + +function deterministicId(value: string) { + const hex = contentVersion(encoder.encode(value)) + .padEnd(32, "0") + .slice(0, 32); + return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-5${hex.slice(13, 16)}-a${hex.slice(17, 20)}-${hex.slice(20, 32)}`; +} + +function contentVersion(value: Uint8Array) { + if (iso.hash) return iso.hash(decoder.decode(value)); + let hash = 2166136261; + for (const byte of value) { + hash ^= byte; + hash = Math.imul(hash, 16777619); + } + return (hash >>> 0).toString(16).padStart(8, "0"); +} + +async function readJson(store: DurableEvalStore, key: string) { + const value = await store.read(key); + return value ? (JSON.parse(decoder.decode(value)) as T) : undefined; +} + +async function writeJson(store: DurableEvalStore, key: string, value: unknown) { + await store.write(key, encoder.encode(stableStringify(value))); +} + +function stableStringify(value: unknown) { + return JSON.stringify(value, (_key, nested) => { + if (nested && typeof nested === "object" && !Array.isArray(nested)) { + return Object.fromEntries( + Object.entries(nested).sort(([left], [right]) => + left.localeCompare(right), + ), + ); + } + return nested; + }); +} + +function assertJsonValue(value: unknown, label: string): JsonValue { + try { + const serialized = JSON.stringify(value); + if (serialized === undefined) + throw new Error("value serializes to undefined"); + return JSON.parse(serialized) as JsonValue; + } catch (error) { + throw new Error(`${label} must be JSON serializable`, { cause: error }); + } +} + +function resultItemId(value: unknown) { + if ( + typeof value !== "object" || + value === null || + !("id" in value) || + typeof value.id !== "string" + ) { + throw new Error("Batch results must contain a string id"); + } + return value.id; +} + +function isBatchTask( + value: unknown, +): value is DurableBatchTask { + return ( + typeof value === "object" && + value !== null && + "kind" in value && + value.kind === BATCH_TASK_KIND + ); +} + +function isBatchScorer( + value: unknown, +): value is DurableBatchScorer { + return ( + typeof value === "object" && + value !== null && + "kind" in value && + value.kind === BATCH_SCORER_KIND + ); +} + +function asError(error: unknown) { + return error instanceof Error ? error : new Error(String(error)); +} diff --git a/js/src/exports.ts b/js/src/exports.ts index abd2bf611..bccb0bd37 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -266,6 +266,16 @@ export { defaultErrorScoreHandler, } from "./framework"; +export type { DurableEvalStore } from "./durable-eval"; + +export { + BatchScorer, + BatchTask, + defineDurableEval, + DurableEvalMemoryStore, + DurableEvalRedisStore, +} from "./durable-eval"; + export { agentAssertionScorer } from "./agent-assertions"; export { DatasetPipeline } from "./dataset-pipeline"; diff --git a/js/src/framework.ts b/js/src/framework.ts index 8628eb0e9..6371afa9c 100644 --- a/js/src/framework.ts +++ b/js/src/framework.ts @@ -1,6 +1,5 @@ import { makeScorerPropagatedEvent, - mergeDicts, Classification, ClassificationItem, Score, @@ -498,6 +497,38 @@ async function getExperimentParametersRef( }; } +export async function _internalInitEvaluatorExperiment( + projectName: string, + evaluator: Evaluator, + data: EvalData, + options: { + disabled?: boolean; + experimentName?: string; + update?: boolean; + } = {}, +): Promise { + if (options.disabled) return null; + const { baseExperiment } = callEvaluatorData(data); + const parameters = await getExperimentParametersRef(evaluator.parameters); + return initExperiment(evaluator.state, { + ...(evaluator.projectId + ? { projectId: evaluator.projectId } + : { project: projectName }), + experiment: options.experimentName ?? evaluator.experimentName, + description: evaluator.description, + metadata: evaluator.metadata, + tags: evaluator.tags, + isPublic: evaluator.isPublic, + update: options.update ?? evaluator.update, + baseExperiment: evaluator.baseExperimentName ?? baseExperiment, + baseExperimentId: evaluator.baseExperimentId, + gitMetadataSettings: evaluator.gitMetadataSettings, + repoInfo: evaluator.repoInfo, + dataset: Dataset.isDataset(data) ? data : undefined, + parameters, + }); +} + export function callEvaluatorData< Input, Expected, @@ -546,6 +577,65 @@ function isIterable(value: unknown): value is Iterable { ); } +export async function _internalResolveEvaluatorData( + evaluator: Pick< + EvaluatorDef, + "data" | "projectName" | "projectId" | "state" + >, + experiment: Experiment | null, +): Promise>> { + if (typeof evaluator.data === "string") { + throw new Error("Unimplemented: string data paths"); + } + let dataResult = + typeof evaluator.data === "function" ? evaluator.data() : evaluator.data; + + if ("_type" in dataResult) { + if (dataResult._type !== "BaseExperiment") { + throw new Error("Invalid _type"); + } + if (!experiment) { + throw new Error( + "Cannot use BaseExperiment() without connecting to Braintrust (you most likely set --no-send-logs)", + ); + } + let name = dataResult.name; + if (isEmpty(name)) { + const baseExperiment = await experiment.fetchBaseExperiment(); + if (!baseExperiment) { + throw new Error("BaseExperiment() failed to fetch base experiment"); + } + name = baseExperiment.name; + } + + dataResult = initExperiment(evaluator.state, { + ...(evaluator.projectId + ? { projectId: evaluator.projectId } + : { project: evaluator.projectName }), + experiment: name, + open: true, + }).asDataset(); + } + + const resolvedDataResult = + dataResult instanceof Promise ? await dataResult : dataResult; + if (isAsyncIterable>(resolvedDataResult)) { + return resolvedDataResult; + } + if ( + Array.isArray(resolvedDataResult) || + isIterable>(resolvedDataResult) + ) { + const iterable = resolvedDataResult as Iterable>; + return (async function* () { + for (const datum of iterable) yield datum; + })(); + } + throw new Error( + "Evaluator data must be an array, iterable, or async iterable", + ); +} + declare global { var _evals: EvaluatorFile; @@ -716,33 +806,13 @@ export async function Eval< const resolvedReporter = options.reporter || defaultReporter; try { - const { data, baseExperiment: defaultBaseExperiment } = callEvaluatorData( - evaluator.data, + const { data } = callEvaluatorData(evaluator.data); + const experiment = await _internalInitEvaluatorExperiment( + name, + evaluator, + data, + { disabled: Boolean(options.parent || options.noSendLogs) }, ); - const parameters = await getExperimentParametersRef(evaluator.parameters); - // NOTE: This code is duplicated with initExperiment in js/src/cli.ts. Make sure - // to update that if you change this. - const experiment = - options.parent || options.noSendLogs - ? null - : initExperiment(evaluator.state, { - ...(evaluator.projectId - ? { projectId: evaluator.projectId } - : { project: name }), - experiment: evaluator.experimentName, - description: evaluator.description, - metadata: evaluator.metadata, - tags: evaluator.tags, - isPublic: evaluator.isPublic, - update: evaluator.update, - baseExperiment: - evaluator.baseExperimentName ?? defaultBaseExperiment, - baseExperimentId: evaluator.baseExperimentId, - gitMetadataSettings: evaluator.gitMetadataSettings, - repoInfo: evaluator.repoInfo, - dataset: Dataset.isDataset(data) ? data : undefined, - parameters, - }); // Ensure experiment ID is resolved before tasks start for OTEL parent attribute support // The Experiment constructor starts resolution (fire-and-forget), but we await here to ensure completion @@ -905,14 +975,49 @@ export function classifierName( return classifier.name || `classifier_${classifier_idx}`; } +export async function _internalRunEvaluatorTask( + task: EvalTask, + datum: EvalCase, + trialIndex: number, + parameters: Record, + span: Span, + reportProgress: (event: TaskProgressEvent) => void = () => undefined, +): Promise<{ + output: unknown; + metadata: Record; + tags: string[]; +}> { + const metadata: Record = { + ...("metadata" in datum ? datum.metadata : {}), + }; + const hooks: EvalHooks, EvalParameters> = { + meta(value) { + Object.assign(metadata, value); + }, + metadata, + expected: "expected" in datum ? datum.expected : undefined, + span, + parameters, + reportProgress, + trialIndex, + tags: [...(datum.tags ?? [])], + }; + const output = await task(datum.input, hooks); + span.log({ output }); + return { + output, + metadata: hooks.metadata, + tags: hooks.tags ?? [], + }; +} + function buildSpanMetadata( results: Array<{ name: string; metadata?: Record }>, ) { return results.length === 1 ? results[0].metadata - : results.reduce( - (prev, s) => mergeDicts(prev, { [s.name]: s.metadata }), - {}, + : Object.fromEntries( + results.map((result) => [result.name, result.metadata]), ); } @@ -923,13 +1028,57 @@ function buildSpanScores( metadata?: Record; }>, ) { - const scoresRecord = results.reduce( - (prev, s) => mergeDicts(prev, { [s.name]: s.score }), - {}, + const scoresRecord = Object.fromEntries( + results.map((result) => [result.name, result.score]), ); return { resultMetadata: buildSpanMetadata(results), scoresRecord }; } +export function _internalPrepareEvaluatorScore( + scoreValue: OneOrMoreScores, + name: string, +): { + results: Score[] | null; + output?: unknown; + metadata?: Record; + scores?: Record; +} { + if (scoreValue === null) return { results: null }; + if (Array.isArray(scoreValue)) { + for (const score of scoreValue) { + if (!(typeof score === "object" && !isEmpty(score))) { + throw new Error( + `When returning an array of scores, each score must be a non-empty object. Got: ${JSON.stringify(score)}`, + ); + } + } + } + let results: Score[]; + if (Array.isArray(scoreValue)) { + results = scoreValue; + } else if (typeof scoreValue === "object" && !isEmpty(scoreValue)) { + results = [scoreValue]; + } else { + results = [{ name, score: scoreValue }]; + } + const { resultMetadata, scoresRecord } = buildSpanScores(results); + const fields = (score: Score) => { + const { metadata: _metadata, name: _name, ...rest } = score; + return rest; + }; + return { + results, + output: + results.length === 1 + ? fields(results[0]) + : Object.fromEntries( + results.map((score) => [score.name ?? name, fields(score)]), + ), + metadata: resultMetadata, + scores: scoresRecord, + }; +} + async function runInScorerSpan( rootSpan: Span, spanName: string, @@ -996,6 +1145,40 @@ function toClassificationItem(c: Classification): ClassificationItem { }; } +export function _internalPrepareEvaluatorClassification( + value: OneOrMoreClassifications, + name: string, +): { + results: Classification[] | null; + output?: unknown; + metadata?: Record; + classifications?: Record; +} { + if (value === null) return { results: null }; + const results = (Array.isArray(value) ? value : [value]).map((result) => + validateClassificationResult(result, name), + ); + const classifications: Record = + Object.create(null); + for (const result of results) { + (classifications[result.name] ??= []).push(toClassificationItem(result)); + } + return { + results, + output: + results.length === 1 + ? toClassificationItem(results[0]) + : Object.fromEntries( + results.map((result) => [ + result.name, + toClassificationItem(result), + ]), + ), + metadata: buildSpanMetadata(results), + classifications, + }; +} + function logScoringFailures( kind: string, failures: { name: string; error: unknown }[], @@ -1076,69 +1259,14 @@ async function runEvaluatorInternal( (evaluator.state ?? _internalGetGlobalState())?.spanCache?.start(); } try { - if (typeof evaluator.data === "string") { - throw new Error("Unimplemented: string data paths"); - } - let dataResult = - typeof evaluator.data === "function" ? evaluator.data() : evaluator.data; - parameters = await validateParameters( parameters ?? {}, evaluator.parameters, ); - - if ("_type" in dataResult) { - if (dataResult._type !== "BaseExperiment") { - // For some reason, the typesystem won't let me check if dataResult._type === "BaseExperiment" - throw new Error("Invalid _type"); - } - if (!experiment) { - throw new Error( - "Cannot use BaseExperiment() without connecting to Braintrust (you most likely set --no-send-logs)", - ); - } - let name = dataResult.name; - if (isEmpty(name)) { - const baseExperiment = await experiment.fetchBaseExperiment(); - if (!baseExperiment) { - throw new Error("BaseExperiment() failed to fetch base experiment"); - } - name = baseExperiment.name; - } - - dataResult = initExperiment(evaluator.state, { - ...(evaluator.projectId - ? { projectId: evaluator.projectId } - : { project: evaluator.projectName }), - experiment: name, - open: true, - }).asDataset(); - } - - const resolvedDataResult = - dataResult instanceof Promise ? await dataResult : dataResult; - - const dataIterable: AsyncIterable> = (() => { - if (isAsyncIterable>(resolvedDataResult)) { - return resolvedDataResult; - } - if ( - Array.isArray(resolvedDataResult) || - isIterable>(resolvedDataResult) - ) { - const iterable = resolvedDataResult as Iterable< - EvalCase - >; - return (async function* () { - for (const datum of iterable) { - yield datum; - } - })(); - } - throw new Error( - "Evaluator data must be an array, iterable, or async iterable", - ); - })(); + const dataIterable = await _internalResolveEvaluatorData( + evaluator, + experiment, + ); progressReporter.start(evaluator.evalName, 0); @@ -1252,37 +1380,29 @@ async function runEvaluatorInternal( }) : undefined; - let metadata: Record = { - ...("metadata" in datum ? datum.metadata : {}), - }; + let metadata: Record = {}; const expected = "expected" in datum ? datum.expected : undefined; let output: unknown = undefined; let error: unknown | undefined = undefined; - let tags: string[] = [...(datum.tags ?? [])]; - const scores: Record = {}; - const classifications: Record = {}; + let tags: string[] = []; + const scores: Record = Object.create(null); + const classifications: Record = + Object.create(null); const scorerNames = (evaluator.scores ?? []).map(scorerName); const classifierNames = (evaluator.classifiers ?? []).map( classifierName, ); let unhandledScores: string[] | null = scorerNames; try { - const meta = (o: Record) => - (metadata = { ...metadata, ...o }); - - await rootSpan.traced( - async (span: Span) => { - const hooksForTask: EvalHooks< - unknown, - Record, - EvalParameters - > = { - meta, - metadata, - expected, + const taskResult = await rootSpan.traced( + (span: Span) => + _internalRunEvaluatorTask( + evaluator.task, + datum, + trialIndex, + parameters ?? {}, span, - parameters: parameters ?? {}, - reportProgress: (event: TaskProgressEvent) => { + (event) => { stream?.({ ...event, id: rootSpan.id, @@ -1291,27 +1411,16 @@ async function runEvaluatorInternal( object_type: "task", }); }, - trialIndex, - tags, - }; - - const outputResult = evaluator.task(datum.input, hooksForTask); - if (outputResult instanceof Promise) { - output = await outputResult; - } else { - output = outputResult; - } - - tags = hooksForTask.tags ?? []; - - span.log({ output }); - }, + ), { name: "task", spanAttributes: { type: SpanTypeAttribute.TASK }, event: { input: datum.input }, }, ); + output = taskResult.output; + metadata = taskResult.metadata; + tags = taskResult.tags; if (tags.length) { rootSpan.log({ output, metadata, expected, tags }); } else { @@ -1334,11 +1443,6 @@ async function runEvaluatorInternal( await rootSpan.export(), ); - const getOtherFields = (s: Score) => { - const { metadata: _metadata, name: _name, ...rest } = s; - return rest; - }; - const [scoreResults, classificationResults] = await Promise.all([ Promise.all( (evaluator.scores ?? []).map((score, score_idx) => @@ -1352,44 +1456,17 @@ async function runEvaluatorInternal( const scoreValue = await Promise.resolve( score(scoringArgs), ); - if (scoreValue === null) return null; - if (Array.isArray(scoreValue)) { - for (const s of scoreValue) { - if (!(typeof s === "object" && !isEmpty(s))) { - throw new Error( - `When returning an array of scores, each score must be a non-empty object. Got: ${JSON.stringify(s)}`, - ); - } - } - } - const results: Score[] = Array.isArray(scoreValue) - ? scoreValue - : typeof scoreValue === "object" && !isEmpty(scoreValue) - ? [scoreValue] - : [ - { - name: scorerNames[score_idx], - score: scoreValue, - }, - ]; - const { resultMetadata, scoresRecord } = - buildSpanScores(results); - const resultOutput = - results.length === 1 - ? getOtherFields(results[0]) - : results.reduce( - (prev, s) => - mergeDicts(prev, { - [s.name]: getOtherFields(s), - }), - {}, - ); + const prepared = _internalPrepareEvaluatorScore( + scoreValue, + scorerNames[score_idx], + ); + if (prepared.results === null) return null; span.log({ - output: resultOutput, - metadata: resultMetadata, - scores: scoresRecord, + output: prepared.output, + metadata: prepared.metadata, + scores: prepared.scores, }); - return results; + return prepared.results; }, ), ), @@ -1406,32 +1483,16 @@ async function runEvaluatorInternal( const classifierValue = await Promise.resolve( classifier(scoringArgs), ); - if (classifierValue === null) return null; - const rawResults = ( - Array.isArray(classifierValue) - ? classifierValue - : [classifierValue] - ).map((result) => - validateClassificationResult( - result, - classifierNames[idx], - ), + const prepared = _internalPrepareEvaluatorClassification( + classifierValue, + classifierNames[idx], ); - const resultOutput = - rawResults.length === 1 - ? toClassificationItem(rawResults[0]) - : rawResults.reduce( - (prev, r) => - mergeDicts(prev, { - [r.name]: toClassificationItem(r), - }), - {}, - ); + if (prepared.results === null) return null; span.log({ - output: resultOutput, - metadata: buildSpanMetadata(rawResults), + output: prepared.output, + metadata: prepared.metadata, }); - return rawResults; + return prepared.results; }, ), ), @@ -1733,7 +1794,7 @@ function ensureScoreAccumulator( // eslint-disable-next-line @typescript-eslint/no-explicit-any results: EvalResult[], ) { - const accumulator: ScoreAccumulator = {}; + const accumulator: ScoreAccumulator = Object.create(null); for (const result of results) { accumulateScores(accumulator, result.scores); } diff --git a/js/src/isomorph.ts b/js/src/isomorph.ts index 231d7e6f2..d23ddcc7e 100644 --- a/js/src/isomorph.ts +++ b/js/src/isomorph.ts @@ -76,7 +76,7 @@ interface Common { path: string, opts?: { recursive?: boolean }, ) => Promise; - writeFile?: (filename: string, data: string) => Promise; + writeFile?: (filename: string, data: string | Uint8Array) => Promise; readFile?: (filename: string) => Promise; readdir?: (path: string) => Promise; utimes?: (path: string, atime: Date, mtime: Date) => Promise; diff --git a/js/src/logger.ts b/js/src/logger.ts index 8bdb8a67b..10662e68d 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -272,10 +272,14 @@ type StartSpanEventArgs = ExperimentLogPartialArgs & Partial; const INITIAL_SPAN_WRITE_AS_MERGE = Symbol( "braintrust.initial-span-write-as-merge", ); +const RESUME_SPAN_WITHOUT_INITIAL_WRITE = Symbol( + "braintrust.resume-span-without-initial-write", +); const INTERNAL_SPAN_CONTEXT = Symbol("braintrust.internal-span-context"); type InitialSpanWriteAsMergeArg = { readonly [INITIAL_SPAN_WRITE_AS_MERGE]?: true; + readonly [RESUME_SPAN_WITHOUT_INITIAL_WRITE]?: true; }; type InternalSpanContextArg = { @@ -2209,6 +2213,37 @@ export function updateSpan({ }); } +/** @internal Rehydrate an exported root span so work can continue in another process. */ +export function _internalResumeSpan({ + exported, + state, +}: { + exported: string; + state?: BraintrustState; +}): Span { + const resolvedState = state ?? _globalState; + const components = SpanComponentsV4.fromStr(exported); + const { row_id, root_span_id, span_id } = components.data; + if (!row_id || !root_span_id || !span_id) { + throw new Error("Only exported root spans can be resumed"); + } + return new SpanImpl({ + state: resolvedState, + parentObjectType: components.data.object_type, + parentObjectId: new LazyValue( + spanComponentsToObjectIdLambda(resolvedState, components), + ), + parentComputeObjectMetadataArgs: undefined, + parentSpanIds: { parentSpanIds: [], rootSpanId: root_span_id }, + spanId: span_id, + event: { id: row_id }, + propagatedEvent: (components.data.propagated_event ?? undefined) as + | StartSpanEventArgs + | undefined, + [RESUME_SPAN_WITHOUT_INITIAL_WRITE]: true, + }); +} + /** * An opaque W3C trace-context, as returned by * {@link extractTraceContextFromHeaders}. @@ -7716,7 +7751,9 @@ export class SpanImpl implements Span { // Deterministic spans can be initialized concurrently by separate // workflow executions, so their first write must not replace later merges. this.isMerge = args[INITIAL_SPAN_WRITE_AS_MERGE] === true; - this.logInternal({ event, internalData }); + if (!args[RESUME_SPAN_WITHOUT_INITIAL_WRITE]) { + this.logInternal({ event, internalData }); + } this.isMerge = true; }