Workflow evaluations

import asyncio
import json
from typing import TypedDict

from openai import AsyncOpenAI
from openai.types.responses import Response

from braintrust.workflow_eval import (
    WorkflowEvalMemoryStore,
    WorkflowEvalStore,
    WorkflowSubmissionCompletion,
    WorkflowSubmissionCompletionPoll,
    WorkflowSubmissionContext,
    WorkflowSubmissionPoll,
    WorkflowTask,
    WorkflowTaskItem,
    WorkflowTaskResult,
    define_workflow_eval,
)

client = AsyncOpenAI()

# Data your submit callback returns and the SDK saves between calls.
# The SDK passes it back to the polling and collection callbacks.
# Here it holds the OpenAI batch ID. You choose the fields.
class Submission(TypedDict):
    batch_id: str

# Sends one prompt to OpenAI as a single-request batch.
async def submit_request(
    prompt: str, context: WorkflowSubmissionContext
) -> Submission:
    request = {
        # Ties the batch request to this SDK submission.
        "custom_id": context.submission_id,
        "method": "POST",
        "url": "/v1/responses",
        "body": {"model": "gpt-5-mini", "input": prompt},
    }
    # The Batch API reads its requests from an uploaded JSONL file.
    batch_file = await client.files.create(
        file=("requests.jsonl", (json.dumps(request) + "\n").encode()),
        purpose="batch",
    )
    batch = await client.batches.create(
        input_file_id=batch_file.id,
        endpoint="/v1/responses",
        completion_window="24h",
        # Lets a webhook handler find the run.
        metadata={"braintrust_run_id": context.run_id},
    )
    # The SDK saves this dictionary and passes it to
    # poll_submission() and collect_task().
    return {"batch_id": batch.id}

# The task's submit callback. The SDK calls it once for each case in data,
# or once per trial if you set trial_count. item holds the case's input and
# expected output. context holds the SDK's run ID and submission ID.
async def submit_task(
    item: WorkflowTaskItem[str, str], context: WorkflowSubmissionContext
) -> Submission:
    prompt = f"{item.input} Reply with only the city name."
    return await submit_request(prompt, context)

# The polling callback. Called during each poll() until the batch
# reports complete.
async def poll_submission(
    submission: Submission, _context: WorkflowSubmissionContext
) -> WorkflowSubmissionPoll:
    batch = await client.batches.retrieve(submission["batch_id"])
    if batch.status == "completed":
        return WorkflowSubmissionPoll("complete")
    if batch.status in {"failed", "expired", "cancelled"}:
        message = f"Batch {batch.id}: {batch.status}, {batch.errors}"
        return WorkflowSubmissionPoll("failed", error=RuntimeError(message))
    return WorkflowSubmissionPoll("pending")

# Reads the model's answer from the batch output file. The file has one
# JSON line because each batch holds one request.
async def collect_text(submission: Submission) -> str:
    batch = await client.batches.retrieve(submission["batch_id"])
    if batch.error_file_id:
        errors = await client.files.content(batch.error_file_id)
        raise RuntimeError(errors.text)
    assert batch.output_file_id is not None
    content = await client.files.content(batch.output_file_id)
    result = json.loads(content.text)
    response = Response.model_validate(result["response"]["body"])
    return response.output_text

# The task's collect callback. Called after poll_submission() reports
# the batch complete.
async def collect_task(
    submission: Submission, _context: WorkflowSubmissionContext
) -> WorkflowTaskResult[str]:
    return WorkflowTaskResult(output=await collect_text(submission))

# A plain scorer function, like the ones you pass to Eval(). The SDK runs
# it on each case's collected output.
def exact_match(input: str, output: str, expected: str) -> float:
    return float(output.strip().casefold() == expected.casefold())

# The store and completion method are parameters so the Redis and
# webhook examples later on this page can reuse this function.
def make_evaluation(
    store: WorkflowEvalStore,
    completion: WorkflowSubmissionCompletion[Submission],
):
    return define_workflow_eval(
        "capital-cities-eval",
        store=store,
        # Each case needs a stable, unique id so the run can resume.
        data=[
            {
                "id": "france",
                "input": "Capital of France?",
                "expected": "Paris",
            },
            {
                "id": "japan",
                "input": "Capital of Japan?",
                "expected": "Tokyo",
            },
        ],
        # Bundles the task's submit and collect callbacks with its
        # completion method.
        task=WorkflowTask(
            submit=submit_task,
            completion=completion,
            collect=collect_task,
        ),
        # Plain scorer functions and WorkflowScorer instances can go here.
        scores=[exact_match],
    )

async def main():
    async with client:
        # The memory store works in this example because the whole run
        # happens in one process.
        evaluation = make_evaluation(
            WorkflowEvalMemoryStore(),
            WorkflowSubmissionCompletionPoll(poll=poll_submission),
        )
        # Submits one batch per case and returns right away with status
        # "waiting".
        result = await evaluation.start()
        print(result.run_id, result.status)
        # In an application, schedule poll() instead of sleeping in one process.
        while result.status == "waiting":
            await asyncio.sleep(60)
            result = await evaluation.poll(result.run_id)
            # pending counts submitted requests that haven't been collected.
            print(result.status, result.pending)
        print(result.summary)

if __name__ == "__main__":
    asyncio.run(main())