Building a golden dataset to test agent behavior

A golden dataset for testing AI agents is the difference between shipping blind and knowing your agent breaks on edge cases. This guide walks through extracting real trajectories, codifying expected behavior, and wiring the set into continuous integration so every change gets a behavioral diff.

Step 1: Capture representative trajectories from production

Start with what your agent actually did. Pull logged request/response cycles from your orchestrator. Aim for coverage across user intent clusters, not a random sample.

If you log to JSONL, a minimal extractor looks like this:

import json
from pathlib import Path

def extract_trajectories(log_path: Path, out_path: Path):
    kept = []
    for line in log_path.read_text().splitlines():
        rec = json.loads(line)
        if rec.get("agent_version") == "1.2.0" and rec.get("status") == "success":
            kept.append({
                "input": rec["user_msg"],
                "steps": rec["tool_calls"],
                "final": rec["final_answer"]
            })
    out_path.write_text(json.dumps(kept, indent=2))
    return len(kept)

count = extract_trajectories(Path("logs/prod.jsonl"), Path("data/raw.json"))
print(f"captured {count} trajectories")

Run it:

python extract.py

You want at least 50 examples per major intent before labeling. Less than that and you are guessing.

Do not rely on random sampling. Cluster inputs by embedding similarity and pick centroids plus outliers. That guarantees you test the long tail where agents fail. A simple KMeans on stored embeddings works:

from sklearn.cluster import KMeans
import numpy as np

embs = np.load("embeddings.npy")  # from your logging
km = KMeans(n_clusters=20).fit(embs)
labels = km.labels_
# pick one example nearest each centroid
samples = []
for c in range(20):
    idx = np.where(labels == c)[0]
    samples.append(int(idx[np.argmin(np.linalg.norm(embs[idx] - km.cluster_centers_[c], axis=1))]))

Use those indices to pull diverse cases into your raw set.

Step 2: Define expected outcomes as executable specs

A golden dataset for testing AI agents is only useful if it encodes what “correct” means. For agents, correctness is a sequence: which tools get called, in what order, with what arguments, and what the final response asserts.

Use a schema that separates invariant checks from fuzzy ones:

{
  "id": "req-refund-001",
  "input": "I was charged twice for order 8821",
  "expect": {
    "tool_calls": [
      {"tool": "lookup_order", "args_contain": {"order_id": "8821"}},
      {"tool": "issue_refund", "args_contain": {"amount": "duplicate"}}
    ],
    "final_contains": ["refund", "24 hours"],
    "forbidden": ["cancel subscription"]
  }
}

Label by hand for the first pass. Two engineers reviewing each case catches ambiguous specs. Store these as JSONL, one object per line.

Why not just check the final text?

Agents fail silently when they call the wrong tool but still produce plausible text. The tool sequence is your strongest signal. A golden dataset for testing AI agents must assert on actions, not just language. If you only grep the final answer, you will miss the class of bug where the agent emails the wrong person but says “done”.

Review labels with a diff tool. When a teammate disputes an expect block, that disagreement is a spec gap—resolve it before merging the case.

Step 3: Normalize and version the set

Treat the dataset like code. Commit it to the repo under tests/golden/. Use a hash of the content as the version tag so CI can cache runs.

import hashlib, json

def version_dataset(path: str) -> str:
    data = open(path, "rb").read()
    return hashlib.sha256(data).hexdigest()[:12]

print(version_dataset("tests/golden/v1.jsonl"))

When you add cases, bump a human-readable major number in the filename: v1.jsonl → v2.jsonl. Never mutate existing cases without renaming; old evaluations must stay reproducible. If a case is wrong, deprecate it with a _deprecated suffix and add a corrected one.

Keep the file line-delimited. JSONL streams well and avoids the merge conflicts you get with pretty-printed JSON arrays.

Step 4: Generate constrained variations

Real traffic is sparse on edge cases. Expand your golden dataset for testing AI agents with synthetic mutations that preserve the expected behavior contract.

Write a small generator that perturbs inputs but keeps the same expect block:

import json, itertools

base = [json.loads(l) for l in open("tests/golden/v1.jsonl")]
mutations = ["urgent: ", "hi, ", "i think ", "can you help? "]

with open("tests/golden/v1_synth.jsonl", "w") as f:
    for case in base:
        for prefix in mutations:
            variant = dict(case)
            variant["id"] = case["id"] + "-mut-" + str(len(prefix))
            variant["input"] = prefix + case["input"]
            f.write(json.dumps(variant) + "\n")

This yields four extra cases per base without new labeling. Keep mutations conservative; you are testing robustness to phrasing, not new intents.

Add entity swaps for slots like order IDs or usernames. Replace with out-of-distribution values to test validation:

import re
def swap_entities(text: str) -> str:
    return re.sub(r"order \d+", "order 999999", text)

A golden dataset for testing AI agents should include a slice where the entity is invalid. The expect block then asserts the agent calls a reject tool rather than proceeding.

Step 5: Execute the dataset as regression tests

Wire the set into pytest. Each case becomes a test that runs the agent and asserts against expect. If you route agent LLM calls through n4n.ai, its OpenAI-compatible endpoint lets you replay the suite against 240+ models with automatic fallback, exposing model-specific regressions in one job.

A minimal harness:

import json, pytest
from my_agent import run_agent  # your entrypoint

cases = [json.loads(l) for l in open("tests/golden/v1.jsonl")]
synth = [json.loads(l) for l in open("tests/golden/v1_synth.jsonl")]
all_cases = cases + synth

@pytest.mark.parametrize("case", all_cases, ids=lambda c: c["id"])
def test_golden(case):
    result = run_agent(case["input"])
    for exp_call, actual in zip(case["expect"]["tool_calls"], result["steps"]):
        assert actual["tool"] == exp_call["tool"]
        for k, v in exp_call["args_contain"].items():
            assert v in actual["args"].get(k, "")
    for token in case["expect"]["final_contains"]:
        assert token.lower() in result["final"].lower()
    for forbidden in case["expect"].get("forbidden", []):
        assert forbidden.lower() not in result["final"].lower()

Run with:

pytest tests/test_golden.py -q

If a test fails, the diff shows exactly which tool call or phrase broke. That is your behavioral regression signal.

Speed up with parallel execution

Agent runs are I/O bound. Use pytest-xdist to split cases across workers:

pytest tests/test_golden.py -n 8

Cache model responses for unchanged cases by hashing the agent code version plus case ID. A simple SQLite cache avoids re-paying for tokens on every run.

Semantic checks for loose outputs

Some final answers can vary. Use an embedding similarity threshold or a secondary LLM grader only after exact checks pass. Keep it isolated so flakiness doesn’t mask real breaks.

def semantic_ok(expected: str, actual: str, threshold=0.85) -> bool:
    import numpy as np
    a, b = embed(expected), embed(actual)
    return float(np.dot(a, b)) >= threshold

Step 6: Verify success and maintain the loop

Success means three things: (1) the suite runs in under 10 minutes in CI, (2) every PR that changes agent logic shows a green or explicit diff, (3) new production incidents produce a new golden case within a sprint.

Add a guard job that fails if coverage drops:

pytest tests/test_golden.py --cov=my_agent --cov-fail-under=80

When an agent bug hits production, write the repro as a golden case before fixing. That converts every incident into a permanent guard.

A golden dataset for testing AI agents is never finished. Schedule a monthly review to prune obsolete cases and add clusters from support tickets. The set is a living contract between your code and your users’ intentions.

What good looks like

After two months, you should have 200+ hand-labeled base cases and 800+ synthetic variants. CI time stays bounded because you cache model responses for unchanged cases. The dataset becomes the fastest way to answer “did this refactor change behavior?” — without reading diffs.

Maintain a short README in tests/golden/ documenting the intent clusters covered and the mutation rules applied. New engineers should be able to add a case in five minutes.

That is the whole loop. Capture, specify, version, expand, execute, maintain. The discipline pays off the first time a silent tool-call regression gets caught before deploy.