Nomearod/agentbench

1"""Calibration runner: generate-outputs | run-judges | build-table.2 3Orchestrates Steps A, C, D from the design doc's data flow. Step B4(hand-labeling) is manual — done in a Jupyter notebook reading5results/calibration_v1_system_outputs.json and appending to6measurements/2026-05-04-judge-calibration-labels.jsonl.7 8Examples:9    python scripts/run_calibration.py generate-outputs --concurrency 510    python scripts/run_calibration.py run-judges --row-config=configs/calibration/rows/baseline.yaml11    python scripts/run_calibration.py build-table12    python scripts/run_calibration.py build-table --strict13"""14 15from __future__ import annotations16 17import argparse18import asyncio19import hashlib20import json21from pathlib import Path22 23import structlog24import yaml25 26logger = structlog.get_logger()27 28REPO = Path(__file__).resolve().parents[1]29CALIBRATION_SPEC = REPO / "agent_bench/evaluation/datasets/calibration_v1.json"30SYSTEM_OUTPUTS = REPO / "results/calibration_v1_system_outputs.json"31LABELS_PATH = REPO / "measurements/2026-05-04-judge-calibration-labels.jsonl"32KAPPA_TABLE_OUT = REPO / "docs/_generated/kappa_table.md"33 34 35def _resolve_concurrency(cli_value: int | None) -> int:36    """CLI flag overrides config field; default is 5. Logs the resolved value."""37    if cli_value is not None:38        resolved = cli_value39    else:40        cfg_path = REPO / "configs/default.yaml"41        cfg_concurrency = None42        if cfg_path.exists():43            cfg = yaml.safe_load(cfg_path.read_text()) or {}44            cfg_concurrency = (cfg.get("evaluation", {}) or {}).get(45                "calibration_concurrency"46            )47        resolved = cfg_concurrency if cfg_concurrency is not None else 548    logger.info("calibration_concurrency_resolved", value=resolved)49    return resolved50 51 52# --- Subcommand: generate-outputs (Step A) ---53 54 55def _build_corpus_orchestrator(cfg, corpus_name: str, embedder, provider):56    """Build a per-corpus Orchestrator wired to that corpus's HybridStore.57 58    Mirrors the per-corpus construction in scripts/evaluate.py so calibration59    runs use the same retrieval stack as production evaluation. The embedder60    and provider are shared across corpora — only the store/retriever/61    SearchTool differ.62    """63    from agent_bench.agents.orchestrator import Orchestrator64    from agent_bench.rag.retriever import Retriever65    from agent_bench.rag.store import HybridStore66    from agent_bench.tools.calculator import CalculatorTool67    from agent_bench.tools.registry import ToolRegistry68    from agent_bench.tools.search import SearchTool69 70    corpus_cfg = cfg.corpora[corpus_name]71    store = HybridStore.load(corpus_cfg.store_path, rrf_k=cfg.rag.retrieval.rrf_k)72    reranker = None73    if cfg.rag.reranker.enabled:74        from agent_bench.rag.reranker import CrossEncoderReranker75 76        reranker = CrossEncoderReranker(model_name=cfg.rag.reranker.model_name)77    retriever = Retriever(78        embedder=embedder,79        store=store,80        default_strategy=cfg.rag.retrieval.strategy,81        candidates_per_system=cfg.rag.retrieval.candidates_per_system,82        reranker=reranker,83        reranker_top_k=cfg.rag.reranker.top_k,84    )85    registry = ToolRegistry()86    registry.register(87        SearchTool(88            retriever=retriever,89            default_top_k=cfg.rag.retrieval.top_k,90            refusal_threshold=corpus_cfg.refusal_threshold,91        )92    )93    registry.register(CalculatorTool())94    return Orchestrator(95        provider=provider,96        registry=registry,97        max_iterations=cfg.agent.max_iterations,98        temperature=cfg.agent.temperature,99    )100 101 102async def cmd_generate_outputs(concurrency: int) -> None:103    """Run the orchestrator against the 30 calibration items with a frozen104    configuration; write results/calibration_v1_system_outputs.json.105 106    The calibration spec is mixed-corpus (k8s + fastapi). Each item carries a107    `corpus` field; we build one Orchestrator per corpus and route by that108    field. A KeyError on an unrecognized corpus is preferable to silently109    misrouting an item to the wrong store.110    """111    from agent_bench.core.config import load_config112    from agent_bench.core.provider import AnthropicProvider113    from agent_bench.evaluation.harness import load_golden_dataset114    from agent_bench.rag.embedder import Embedder115 116    spec = json.loads(CALIBRATION_SPEC.read_text())117    target_ids = {i["id"]: i for i in spec["items"]}118 119    fastapi = load_golden_dataset(120        REPO / "agent_bench/evaluation/datasets/tech_docs_golden.json"121    )122    k8s = load_golden_dataset(123        REPO / "agent_bench/evaluation/datasets/k8s_golden.json"124    )125    items = [q for q in (fastapi + k8s) if q.id in target_ids]126    if len(items) != len(target_ids):127        missing = set(target_ids) - {q.id for q in items}128        raise SystemExit(129            f"calibration items not found in goldens: {sorted(missing)}"130        )131 132    cfg = load_config()133    provider = AnthropicProvider(cfg)134    embedder = Embedder(model_name=cfg.embedding.model, cache_dir=cfg.embedding.cache_dir)135 136    item_corpus = {it.id: target_ids[it.id]["corpus"] for it in items}137    unknown: dict[str, list[str]] = {}138    for it_id, corpus in item_corpus.items():139        if corpus not in cfg.corpora:140            unknown.setdefault(corpus, []).append(it_id)141    if unknown:142        examples = "; ".join(143            f"{cor!r}: {sorted(ids)[:3]}" for cor, ids in sorted(unknown.items())144        )145        raise KeyError(146            f"calibration spec references corpora not in cfg.corpora — "147            f"{examples}; configured corpora: {sorted(cfg.corpora)!r}"148        )149 150    corpora_needed = sorted(set(item_corpus.values()))151    orchestrators = {152        name: _build_corpus_orchestrator(cfg, name, embedder, provider)153        for name in corpora_needed154    }155 156    sem = asyncio.Semaphore(concurrency)157 158    async def _run_one(item):159        async with sem:160            response = await orchestrators[item_corpus[item.id]].run(161                question=item.question,162                system_prompt="You are a helpful assistant.",163            )164            answer = response.answer165            sources = sorted(s.source for s in response.sources)166            sys_hash = hashlib.sha256(167                f"{item.id}\x00{answer}\x00{','.join(sources)}".encode("utf-8")168            ).hexdigest()169            return {170                "item_id": item.id,171                "question": item.question,172                "category": item.category,173                "answer": answer,174                "sources": [s.source for s in response.sources],175                "ranked_sources": response.ranked_sources,176                "source_chunks": response.source_chunks,177                "source_snippets": item.source_snippets,178                "reference_answer": item.reference_answer,179                "system_output_hash": sys_hash,180                "stratum": target_ids[item.id]["stratum"],181                "corpus": target_ids[item.id]["corpus"],182            }183 184    records = await asyncio.gather(*[_run_one(it) for it in items])185    SYSTEM_OUTPUTS.parent.mkdir(parents=True, exist_ok=True)186    SYSTEM_OUTPUTS.write_text(json.dumps(records, indent=2) + "\n")187    logger.info(188        "generate_outputs_complete", count=len(records), path=str(SYSTEM_OUTPUTS)189    )190 191 192# --- Subcommand: run-judges (Step C, one row per invocation) ---193 194 195def _make_provider(name: str, cfg, *, model: str | None = None):196    from agent_bench.core.provider import AnthropicProvider, OpenAIProvider197 198    if name == "anthropic":199        return AnthropicProvider(cfg, model=model)200    if name == "openai":201        return OpenAIProvider(cfg, model=model)202    raise ValueError(f"unknown provider: {name}")203 204 205def _make_judge(206    provider_name: str,207    model_id: str,208    dimension: str,209    cfg,210    *,211    use_cot: bool = True,212    use_anchors: bool = True,213    abstain_allowed_override: bool | None = None,214):215    from agent_bench.evaluation.judges.base import Rubric216    from agent_bench.evaluation.judges.citation_faithfulness import (217        CitationFaithfulnessJudge,218    )219    from agent_bench.evaluation.judges.completeness import CompletenessJudge220    from agent_bench.evaluation.judges.groundedness import GroundednessJudge221    from agent_bench.evaluation.judges.relevance import RelevanceJudge222 223    judge_class = {224        "groundedness": GroundednessJudge,225        "relevance": RelevanceJudge,226        "completeness": CompletenessJudge,227        "citation_faithfulness": CitationFaithfulnessJudge,228    }229    rubric_dir = REPO / "agent_bench/evaluation/rubrics"230    rubric = Rubric.from_markdown_file(rubric_dir / f"{dimension}.md")231    if not use_anchors:232        # Strip ### Example sections — body_markdown changes, so233        # ScoreResult.rubric_version naturally distinguishes anchored vs234        # stripped variants when the calibration report buckets results.235        rubric = rubric.strip_anchors()236    return judge_class[dimension](237        judge_provider=_make_provider(provider_name, cfg, model=model_id),238        rubric=rubric,239        model_id=model_id,240        use_cot=use_cot,241        abstain_allowed_override=abstain_allowed_override,242    )243 244 245def _row_judge_options(row: dict) -> dict:246    """Pull `options` from a row config and project to _make_judge kwargs.247 248    Defaults (when keys are missing) match the baseline contract: CoT on,249    anchors on, abstain follows the rubric (no override).250    """251    opts = row.get("options") or {}252    abstain_allowed = opts.get("abstain_allowed")253    return {254        "use_cot": bool(opts.get("use_cot", True)),255        "use_anchors": bool(opts.get("use_anchors", True)),256        # None = follow rubric; explicit True/False = override257        "abstain_allowed_override": (258            None if abstain_allowed is None else bool(abstain_allowed)259        ),260    }261 262 263def _build_item_and_output(rec: dict):264    from agent_bench.agents.orchestrator import AgentResponse, SourceReference265    from agent_bench.core.types import TokenUsage266    from agent_bench.evaluation.harness import GoldenQuestion267 268    item = GoldenQuestion(269        id=rec["item_id"],270        question=rec["question"],271        expected_answer_keywords=[],272        expected_sources=[],273        category=rec["category"],274        difficulty="easy",275        requires_calculator=False,276        source_snippets=rec.get("source_snippets", []),277        reference_answer=rec.get("reference_answer", ""),278    )279    output = AgentResponse(280        answer=rec["answer"],281        sources=[SourceReference(source=s) for s in rec["sources"]],282        ranked_sources=rec.get("ranked_sources", []),283        source_chunks=rec.get("source_chunks", []),284        iterations=1,285        usage=TokenUsage(input_tokens=0, output_tokens=0, estimated_cost_usd=0),286        latency_ms=0,287    )288    return item, output289 290 291async def cmd_run_judges(row_config_path: Path, concurrency: int) -> None:292    """Score the frozen system outputs with the row's judge configuration."""293    from agent_bench.core.config import load_config294    from agent_bench.evaluation.variance.jury import jury295    from agent_bench.evaluation.variance.rubric_permute import rubric_permute296 297    if not SYSTEM_OUTPUTS.exists():298        raise SystemExit(299            f"{SYSTEM_OUTPUTS} not found — run `generate-outputs` first."300        )301    row = yaml.safe_load(row_config_path.read_text())302    outputs = json.loads(SYSTEM_OUTPUTS.read_text())303 304    cfg = load_config()305    sem = asyncio.Semaphore(concurrency)306    all_results: list[dict] = []307    strategy = row["strategy"]308 309    def _skip_oos(rec: dict, dim: str) -> bool:310        return rec["category"] == "out_of_scope" and dim != "relevance"311 312    judge_opts = _row_judge_options(row)313 314    if strategy == "single":315        # Build one judge per dimension up-front, then gather all316        # (dim, item) pairs in a single asyncio.gather call. Previous317        # design serialized across dimensions (each dim awaited fully318        # before the next started), leaving Phase-11 wall-clock on the319        # table when the calibration spend is API-rate-limited.320        judges_by_dim = {321            dim: _make_judge(322                row["provider"], row["model_id"], dim, cfg, **judge_opts323            )324            for dim in row["dimensions"]325        }326 327        async def score_one(rec: dict, dim: str, judge):328            async with sem:329                if _skip_oos(rec, dim):330                    return None331                item, output = _build_item_and_output(rec)332                result = await judge.score(item, output)333                return {"item_id": rec["item_id"], "dimension": dim, **result.model_dump()}334 335        coros = [336            score_one(rec, dim, judge)337            for dim, judge in judges_by_dim.items()338            for rec in outputs339        ]340        gathered = await asyncio.gather(*coros)341        all_results.extend([r for r in gathered if r is not None])342 343    elif strategy == "rubric_permute":344        # Sequential per-item by design: PermutedJudge writes to the345        # sidecar JSONL with append mode and within-call ordering matters346        # for downstream per-permutation analysis (the kappa_table joins347        # by item_id but the sidecar order encodes the permutation seed348        # sequence). Across-dim parallelism is left for v1.1 once the349        # sidecar contract proves stable.350        for dim in row["dimensions"]:351            judge = _make_judge(352                row["provider"], row["model_id"], dim, cfg, **judge_opts353            )354            sidecar = REPO / row.get(355                "sidecar_path", "results/calibration_v1_permute_members.jsonl"356            )357            permuted = rubric_permute(358                judge,359                n=row["options"]["n_permutations"],360                seeds=row["options"]["seeds"],361                sidecar_path=sidecar,362            )363            for rec in outputs:364                if _skip_oos(rec, dim):365                    continue366                item, output = _build_item_and_output(rec)367                result = await permuted.score(item, output)368                all_results.append({"item_id": rec["item_id"], "dimension": dim, **result.model_dump()})369 370    elif strategy == "jury":371        # Same sequential rationale as rubric_permute: jury writes a372        # per-member sidecar and downstream analysis benefits from stable373        # ordering. The asyncio.gather inside Jury.score does parallelize374        # member calls within an item; the across-item / across-dim375        # serialization is the conservative choice.376        for dim in row["dimensions"]:377            members = [378                _make_judge(m["provider"], m["model_id"], dim, cfg, **judge_opts)379                for m in row["members"]380            ]381            sidecar = REPO / row["sidecar_path"]382            weights = (383                _compute_kappa_weights(384                    REPO / row["weights_source"],385                    dim,386                    expected_judge_ids={m.judge_id for m in members},387                )388                if row.get("aggregation") == "kappa_weighted"389                else None390            )391            j = jury(392                judges=members,393                aggregation=row["aggregation"],394                weights=weights,395                quorum=row.get("quorum"),396                sidecar_path=sidecar,397            )398            for rec in outputs:399                if _skip_oos(rec, dim):400                    continue401                item, output = _build_item_and_output(rec)402                result = await j.score(item, output)403                all_results.append({"item_id": rec["item_id"], "dimension": dim, **result.model_dump()})404    else:405        raise SystemExit(f"unknown strategy: {strategy}")406 407    out_path = REPO / row["output_path"]408    out_path.parent.mkdir(parents=True, exist_ok=True)409    out_path.write_text(json.dumps(all_results, indent=2) + "\n")410    logger.info(411        "run_judges_complete",412        row=row["label"],413        count=len(all_results),414        path=str(out_path),415    )416 417 418def _compute_kappa_weights(419    predictions_path: Path,420    dimension: str,421    expected_judge_ids: set[str],422) -> dict[str, float]:423    """Compute per-judge weight = max(0, Cohen's κ vs gold labels) for the424    dimension, from a predictions file (JSON list or JSONL).425 426    v1.1 replaces v1's stub (which returned 1.0 for every judge_id seen,427    causing asymmetric coverage to amplify rather than suppress an428    unweighted member). Hard-errors if `predictions_path` is missing,429    if any `expected_judge_ids` member has no scored (non-abstain)430    predictions for `dimension`, or if no labels are available for the431    dimension.432 433    The κ → weight mapping clips negative κ to 0; a member with κ ≤ 0 on434    a dimension contributes weight 0 (effective exclusion via weighting).435    This is the "soft exclusion" behavior — explicit per-dimension436    exclusion is tracked separately on the v1.2 fix-list.437 438    Pragmatic v1.1: `predictions_path` may point at the same calibration439    set used for κ reporting (circular weighting); this is documented in440    the v1.1 jury-rescue DECISIONS entry. v1.2 will require a held-out441    validation set.442    """443    from agent_bench.evaluation.calibration.metrics import cohen_kappa444 445    if not predictions_path.exists():446        raise FileNotFoundError(447            f"weights source {predictions_path} does not exist; v1.1 "448            f"requires explicit κ-derived weights — no silent fallback"449        )450 451    # Load predictions: JSON list (baseline-style) or JSONL (sidecar-style).452    raw = predictions_path.read_text()453    if predictions_path.suffix == ".jsonl":454        preds = [json.loads(line) for line in raw.splitlines() if line.strip()]455    else:456        preds = json.loads(raw)457 458    if not LABELS_PATH.exists():459        raise FileNotFoundError(460            f"labels file {LABELS_PATH} does not exist; cannot compute "461            f"κ-derived weights"462        )463    labels: dict[str, int] = {}464    for line in LABELS_PATH.read_text().splitlines():465        if not line.strip():466            continue467        rec = json.loads(line)468        if rec.get("dimension") != dimension or rec.get("abstained"):469            continue470        labels[rec["system_output_hash"]] = int(rec["score"])471 472    if not labels:473        raise ValueError(474            f"no gold labels for dimension={dimension!r} in {LABELS_PATH}; "475            f"cannot compute κ-derived weights"476        )477 478    # Group predictions by judge_id, joining to labels by system_output_hash.479    # The sidecar JSONL has one record per (judge × item × dim); the baseline480    # JSON has the same. Both expose `judge_id` of the form `{model}_{dim}`,481    # `system_output_hash`, `score`, and (for the abstain-aware filter) the482    # `Unknown` sentinel.483    by_judge: dict[str, list[tuple[int, int]]] = {}484    for p in preds:485        # JSONL sidecar lacks `dimension` field; we filter by suffix on486        # judge_id instead, which encodes dimension.487        if not p["judge_id"].endswith(f"_{dimension}"):488            continue489        if p["score"] == "Unknown":490            continue491        h = p["system_output_hash"]492        if h not in labels:493            continue494        by_judge.setdefault(p["judge_id"], []).append(495            (labels[h], int(p["score"]))496        )497 498    missing = expected_judge_ids - by_judge.keys()499    if missing:500        raise ValueError(501            f"weights source {predictions_path} has no predictions for "502            f"expected judge_ids {sorted(missing)} on dimension={dimension!r}. "503            f"Source covers {sorted(by_judge.keys())}. v1.1 requires "504            f"symmetric coverage — point weights_source at a predictions "505            f"file containing every jury member's verdicts (e.g. the jury "506            f"sidecar from a prior run)."507        )508 509    weights: dict[str, float] = {}510    for jid in expected_judge_ids:511        pairs = by_judge[jid]512        y_lab = [a for a, _ in pairs]513        y_pred = [b for _, b in pairs]514        kappa = cohen_kappa(y_lab, y_pred)515        weights[jid] = max(0.0, kappa)516        logger.info(517            "kappa_weight_computed",518            judge_id=jid,519            dimension=dimension,520            kappa=kappa,521            weight=weights[jid],522            n=len(pairs),523        )524    return weights525 526 527# --- Subcommand: build-table (Step D) ---528 529 530def cmd_build_table(strict: bool) -> None:531    from agent_bench.evaluation.calibration.report import generate_kappa_table532 533    predictions_glob = str(REPO / "results/calibration_v1_judge_*.json")534    generate_kappa_table(535        predictions_glob=predictions_glob,536        labels_path=str(LABELS_PATH),537        output_path=str(KAPPA_TABLE_OUT),538        strict=strict,539    )540    logger.info("build_table_complete", path=str(KAPPA_TABLE_OUT), strict=strict)541 542 543def main() -> None:544    parser = argparse.ArgumentParser(545        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter546    )547    sub = parser.add_subparsers(dest="cmd", required=True)548 549    p_gen = sub.add_parser(550        "generate-outputs", help="Step A: generate frozen system outputs"551    )552    p_gen.add_argument("--concurrency", type=int, default=None)553 554    p_run = sub.add_parser("run-judges", help="Step C: score one ablation row")555    p_run.add_argument("--row-config", type=Path, required=True)556    p_run.add_argument("--concurrency", type=int, default=None)557 558    p_tab = sub.add_parser(559        "build-table", help="Step D: aggregate predictions into κ table"560    )561    p_tab.add_argument(562        "--strict",563        action="store_true",564        help="Raise on missing predictions/labels (final-artifact path)",565    )566 567    args = parser.parse_args()568    if args.cmd == "generate-outputs":569        asyncio.run(cmd_generate_outputs(_resolve_concurrency(args.concurrency)))570    elif args.cmd == "run-judges":571        asyncio.run(572            cmd_run_judges(args.row_config, _resolve_concurrency(args.concurrency))573        )574    elif args.cmd == "build-table":575        cmd_build_table(strict=args.strict)576 577 578if __name__ == "__main__":579    main()580