feat: Add batch evals api

[Luca Forstner (lforst)]

Batch eval

# openai_batch_eval.py
import json
from dataclasses import dataclass
from typing import Any

from braintrust import (
    BatchCompletionPoll,
    BatchCompletionWebhook,
    BatchContext,
    BatchScorer,
    BatchScorerItem,
    BatchScorerResult,
    BatchTask,
    BatchTaskItem,
    BatchTaskResult,
    DurableEvalStore,
    define_workflow_eval,
)
from openai import AsyncOpenAI


openai = AsyncOpenAI()


@dataclass(frozen=True)
class Submission:
    batch_id: str


async def submit_openai_batch(
    requests: list[dict[str, Any]],
    context: BatchContext,
) -> Submission:
    jsonl = "\n".join(
        json.dumps(
            {
                "custom_id": request["id"],
                "method": "POST",
                "url": "/v1/chat/completions",
                "body": request["body"],
            }
        )
        for request in requests
    )

    file = await openai.files.create(
        file=(
            f"{context.batch_id}.jsonl",
            jsonl.encode("utf-8"),
            "application/jsonl",
        ),
        purpose="batch",
    )

    batch = await openai.batches.create(
        input_file_id=file.id,
        endpoint="/v1/chat/completions",
        completion_window="24h",
        metadata={
            "braintrust_run_id": context.run_id,
            "braintrust_batch_id": context.batch_id,
        },
    )

    return Submission(batch_id=batch.id)


async def collect_openai_batch(submission: Submission) -> list[dict[str, Any]]:
    batch = await openai.batches.retrieve(submission.batch_id)

    if not batch.output_file_id:
        raise RuntimeError(f"OpenAI batch {batch.id} has no output file")

    response = await openai.files.content(batch.output_file_id)
    return [json.loads(line) for line in response.text.splitlines() if line]


def row_to_output(row: dict[str, Any]) -> str:
    return row["response"]["body"]["choices"][0]["message"]["content"]


def make_openai_batch_eval(
    store: DurableEvalStore,
    completion: BatchCompletionPoll[Submission] | BatchCompletionWebhook[Submission],
):
    async def submit_task(
        items: list[BatchTaskItem[str, str]],
        context: BatchContext,
    ) -> Submission:
        return await submit_openai_batch(
            [
                {
                    "id": item.id,
                    "body": {
                        "model": "gpt-4o-mini",
                        "messages": [{"role": "user", "content": item.input}],
                    },
                }
                for item in items
            ],
            context,
        )

    async def collect_task(
        submission: Submission,
        _context: BatchContext,
    ) -> list[BatchTaskResult[str]]:
        return [
            BatchTaskResult(id=row["custom_id"], output=row_to_output(row))
            for row in await collect_openai_batch(submission)
        ]

    async def submit_score(
        items: list[BatchScorerItem[str, str, str]],
        context: BatchContext,
    ) -> Submission:
        return await submit_openai_batch(
            [
                {
                    "id": item.id,
                    "body": {
                        "model": "gpt-4o-mini",
                        "messages": [
                            {
                                "role": "system",
                                "content": (
                                    "Return exactly 1 if the answer is correct, "
                                    "otherwise return exactly 0."
                                ),
                            },
                            {
                                "role": "user",
                                "content": json.dumps(
                                    {
                                        "answer": item.output,
                                        "expected": item.expected,
                                    }
                                ),
                            },
                        ],
                    },
                }
                for item in items
            ],
            context,
        )

    async def collect_score(
        submission: Submission,
        _context: BatchContext,
    ) -> list[BatchScorerResult]:
        return [
            BatchScorerResult(
                id=row["custom_id"],
                score=float(row_to_output(row)),
            )
            for row in await collect_openai_batch(submission)
        ]

    return define_workflow_eval(
        "OpenAI batch demo",
        store=store,
        data=[
            {
                "id": "france",
                "input": "What is the capital of France?",
                "expected": "Paris",
            },
            {
                "id": "japan",
                "input": "What is the capital of Japan?",
                "expected": "Tokyo",
            },
        ],
        task=BatchTask(
            batch_size=100,
            submit=submit_task,
            completion=completion,
            collect=collect_task,
        ),
        scores=[
            BatchScorer(
                name="openai_judge",
                batch_size=100,
                submit=submit_score,
                completion=completion,
                collect=collect_score,
            )
        ],
    )

Polling example

# openai_poll.py
import asyncio
import os

from braintrust import (
    BatchCompletionPoll,
    BatchContext,
    BatchPollResult,
    DurableEvalRedisStore,
)
from openai import AsyncOpenAI
from redis.asyncio import Redis

from openai_batch_eval import Submission, make_openai_batch_eval


openai = AsyncOpenAI()


async def poll_batch(
    submission: Submission,
    _context: BatchContext,
) -> BatchPollResult:
    batch = await openai.batches.retrieve(submission.batch_id)

    if batch.status == "completed":
        return BatchPollResult(status="complete")

    if batch.status in {"failed", "expired", "cancelled"}:
        return BatchPollResult(
            status="failed",
            error=RuntimeError(f"OpenAI batch {batch.id} {batch.status}"),
        )

    return BatchPollResult(status="pending")


async def main() -> None:
    redis = Redis.from_url(os.environ["REDIS_URL"])
    try:
        durable_eval = make_openai_batch_eval(
            DurableEvalRedisStore(redis),
            BatchCompletionPoll(poll=poll_batch),
        )

        result = await durable_eval.start()
        run_id = result.run_id

        while result.status != "completed":
            result = await durable_eval.poll(run_id)
            await asyncio.sleep(10)
    finally:
        await redis.aclose()


asyncio.run(main())

Webhook example

# openai_webhook_start.py
import asyncio
import os

from braintrust import BatchCompletionWebhook, DurableEvalRedisStore
from redis.asyncio import Redis

from openai_batch_eval import make_openai_batch_eval


async def main() -> None:
    redis = Redis.from_url(os.environ["REDIS_URL"])
    try:
        durable_eval = make_openai_batch_eval(
            DurableEvalRedisStore(redis),
            BatchCompletionWebhook(
                get_external_id=lambda submission, _context: submission.batch_id
            ),
        )

        print(await durable_eval.start())
    finally:
        await redis.aclose()


asyncio.run(main())

Handle batch.completed:

# openai_webhook_server.py
import os

from braintrust import BatchCompletionWebhook, DurableEvalRedisStore
from fastapi import FastAPI, Request, Response
from openai import AsyncOpenAI
from redis.asyncio import Redis

from openai_batch_eval import make_openai_batch_eval


openai = AsyncOpenAI()
redis = Redis.from_url(os.environ["REDIS_URL"])
durable_eval = make_openai_batch_eval(
    DurableEvalRedisStore(redis),
    BatchCompletionWebhook(
        get_external_id=lambda submission, _context: submission.batch_id
    ),
)

app = FastAPI()


@app.post("/openai-webhook")
async def openai_webhook(request: Request) -> Response:
    # Signature verification requires the unparsed request body.
    event = openai.webhooks.unwrap(
        await request.body(),
        request.headers,
        secret=os.environ["OPENAI_WEBHOOK_SECRET"],
    )

    if event.type == "batch.completed":
        batch = await openai.batches.retrieve(event.data.id)
        run_id = batch.metadata.get("braintrust_run_id") if batch.metadata else None

        if not run_id:
            raise RuntimeError(f"Batch {batch.id} has no Braintrust run ID")

        # In production, enqueue this work and return 200 immediately.
        result = await durable_eval.process_batch_result(
            run_id,
            external_id=batch.id,
        )

        print(result)

    return Response(status_code=200)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 87a15e5e03

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "Codex (@codex) review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "Codex (@codex) address that feedback".

[chatgpt-codex-connector]

Make action claims recoverable after callback interruption

When a submit, task, scorer, classifier, or finalization callback raises—or the worker stops after this claim but before writing its result—the claim remains permanently stored. Every subsequent poll() sees created=False, skips the incomplete action, and can never satisfy the missing stage record, leaving the run in waiting forever with the memory store and generally until the run itself expires with Redis. Use a recoverable lease/state transition or otherwise permit retrying claims whose corresponding result was not persisted.

Useful? React with 👍 / 👎.

[Luca Forstner (lforst)]

solving this right now is overkill