scripts/build_runs.py

#!/usr/bin/env python3
"""Build normalized HERB full-run and comparison shards.

The source response files can be hundreds of megabytes, so they are read one
JSONL row at a time. Repeated prompts are never retained or written.

Run from the viewer repository root:
    python scripts/build_runs.py
"""

from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path
from typing import Any, Iterator


ROOT = Path(__file__).resolve().parent.parent
TYPE_AWARE_JUDGE_ROOT = Path(
    "/home/azureuser/projects/information-scaffolds/outputs/"
    "herb_type_aware_judge_20260715"
)
MAX_STRING_CHARS = 8192
MAX_EVENTS_BYTES = 256 * 1024
RETAINED_SOURCE_FIELDS = {
    "qid",
    "dataset",
    "answer",
    "usage",
    "tokens",
    "finish_reason",
    "stop_reason",
    "finish_reasons",
    "turns",
    "tool_call_counts",
    "events",
    "parsed",
    "judge_text",
}

RUNS = [
    {
        "slot": "c1",
        "label": "c1 Closed-book",
        "description": "Closed-book baseline without corpus documents.",
        "response": "/home/azureuser/projects/information-scaffolds/outputs/herb_phantom_full/full_closedbook/named-outputs/response/response",
        "recovery_response": "/home/azureuser/projects/information-scaffolds/outputs/herb_phantom_recover/cb_len4x/named-outputs/response/response",
        "judge": str(TYPE_AWARE_JUDGE_ROOT / "c1/named-outputs/judged/judged"),
        "score_mode": "mean_judge_score",
    },
    {
        "slot": "c2",
        "label": "c2 With-docs",
        "description": "Open-book baseline with retrieved documents in the prompt.",
        "response": "/home/azureuser/projects/information-scaffolds/outputs/herb_phantom_full/full_openbook/named-outputs/response/response",
        "recovery_response": "/home/azureuser/projects/information-scaffolds/outputs/herb_phantom_recover/ob_len4x/named-outputs/response/response",
        "judge": str(TYPE_AWARE_JUDGE_ROOT / "c2/named-outputs/judged/judged"),
        "score_mode": "mean_judge_score",
    },
    {
        "slot": "c6",
        "label": "c6 Agentic-DCI",
        "description": "Agentic DCI baseline over corpus scaffolds.",
        "response": "/home/azureuser/projects/information-scaffolds/outputs/herb_phantom_full/full_dci/named-outputs/response/response",
        "judge": str(TYPE_AWARE_JUDGE_ROOT / "c6/named-outputs/judged/judged"),
        "agentic": True,
        "score_mode": "mean_judge_score",
    },
    {
        "slot": "naive",
        "label": "Naive-search",
        "description": "Agentic baseline using naive corpus search.",
        "response": "/home/azureuser/projects/information-scaffolds/outputs/herb_phantom_full/full_naive_herb/named-outputs/response/response",
        "judge": str(TYPE_AWARE_JUDGE_ROOT / "naive/named-outputs/judged/judged"),
        "agentic": True,
        "score_mode": "mean_judge_score",
    },
    {
        "slot": "e2e_raw",
        "label": "E2E raw",
        "description": "End-to-end agent over raw HERB corpus structures.",
        "response": "/home/azureuser/projects/information-scaffolds/outputs/e2e_runs/new-datasets-full-20260711/herb_raw/named-outputs/predictions/predictions",
        "judge": str(TYPE_AWARE_JUDGE_ROOT / "e2e_raw/named-outputs/judged/judged"),
        "agentic": True,
        "score_mode": "mean_judge_score",
    },
    {
        "slot": "e2e_combined",
        "label": "E2E combined",
        "description": "End-to-end agent over the combined HERB structures.",
        "response": "/home/azureuser/projects/information-scaffolds/outputs/e2e_runs/new-datasets-full-20260711/herb_combined/named-outputs/predictions/predictions",
        "judge": str(
            TYPE_AWARE_JUDGE_ROOT / "e2e_combined/named-outputs/judged/judged"
        ),
        "agentic": True,
        "score_mode": "mean_judge_score",
    },
    {
        "slot": "e2e_rawtext",
        "label": "E2E v3 + rawtext",
        "description": "Native E2E-v3 structures with the raw HERB corpus overlay.",
        "response": (
            "/mnt/ramdisk/blobstore/timchen0618/data/eval/herb/viewer_inputs/"
            "e2e_rawtext/predictions"
        ),
        "judge": str(
            TYPE_AWARE_JUDGE_ROOT
            / "e2e_rawtext/named-outputs/canonical_evaluated/evaluated"
        ),
        "agentic": True,
        "score_mode": "mean_judge_score",
    },
]


def jsonl_rows(path: Path) -> Iterator[dict[str, Any]]:
    with path.open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, 1):
            if not line.strip():
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_number}: {exc}") from exc


def iter_herb_rows(
    path: Path, *, include_events: bool
) -> Iterator[tuple[str, dict[str, Any]]]:
    seen: set[str] = set()
    for row in jsonl_rows(path):
        if row.get("dataset") != "herb":
            continue
        qid = row.get("qid")
        if not isinstance(qid, str) or not qid:
            raise ValueError(f"HERB row in {path} has no qid")
        if qid in seen:
            raise ValueError(f"duplicate HERB qid {qid!r} in {path}")
        seen.add(qid)
        retained = {
            key: cap_value(value)
            for key, value in row.items()
            if key in RETAINED_SOURCE_FIELDS and key != "events"
        }
        if include_events:
            retained["events"] = compact_events(row.get("events"))
        yield qid, retained


def load_herb_rows(
    path: Path, *, include_events: bool
) -> dict[str, dict[str, Any]]:
    return dict(iter_herb_rows(path, include_events=include_events))


def cap_value(value: Any) -> Any:
    if isinstance(value, str):
        if len(value) <= MAX_STRING_CHARS:
            return value
        suffix = ""
        for _ in range(2):
            kept = MAX_STRING_CHARS - len(suffix)
            suffix = f"\n… [truncated {len(value) - kept:,} chars]"
        return value[: MAX_STRING_CHARS - len(suffix)] + suffix
    if isinstance(value, list):
        return [cap_value(item) for item in value]
    if isinstance(value, dict):
        return {str(key): cap_value(item) for key, item in value.items()}
    return value


def compact_events(events: Any) -> list[dict[str, Any]]:
    if not isinstance(events, list):
        return []
    compact: list[dict[str, Any]] = []
    used = 2
    for index, raw_event in enumerate(events):
        if not isinstance(raw_event, dict):
            event = {"type": "event", "content": cap_value(raw_event)}
        else:
            event = cap_value(
                {
                    key: raw_event.get(key)
                    for key in ("type", "name", "input", "content")
                    if raw_event.get(key) is not None
                }
            )
        encoded_size = len(
            json.dumps(event, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        ) + (1 if compact else 0)
        if used + encoded_size > MAX_EVENTS_BYTES:
            compact.append(
                {
                    "type": "truncated",
                    "content": (
                        f"Event payload capped near {MAX_EVENTS_BYTES // 1024} KiB; "
                        f"{len(events) - index:,} event(s) omitted."
                    ),
                }
            )
            break
        compact.append(event)
        used += encoded_size
    return compact


def response_tokens(response: dict[str, Any] | None) -> dict[str, Any]:
    if not response:
        return {}
    tokens = response.get("tokens")
    if not isinstance(tokens, dict):
        tokens = response.get("usage")
    return cap_value(tokens) if isinstance(tokens, dict) else {}


def failure_reason(
    response: dict[str, Any] | None, judge: dict[str, Any] | None
) -> str | None:
    if response is None:
        return "missing_response"
    if not str(response.get("answer") or "").strip():
        stop = response.get("stop_reason") or response.get("finish_reason")
        return str(stop or "empty_answer")
    if judge is None:
        return "missing_judge"
    parsed = judge.get("parsed")
    if not isinstance(parsed, dict):
        return "missing_judge_result"
    if parsed.get("parse_error"):
        return "judge_parse_error"
    if not isinstance(parsed.get("correct"), bool) and not isinstance(
        parsed.get("judge_score"), (int, float)
    ):
        return "missing_judge_verdict"
    return None


def normalized_record(
    gold: dict[str, Any],
    response: dict[str, Any] | None,
    judge: dict[str, Any] | None,
) -> dict[str, Any]:
    parsed = judge.get("parsed") if judge else {}
    if not isinstance(parsed, dict):
        parsed = {}
    prediction = response.get("answer") if response else None
    answered = bool(isinstance(prediction, str) and prediction.strip())
    tool_counts = response.get("tool_call_counts") if response else {}
    if not isinstance(tool_counts, dict):
        tool_counts = {}
    finish_reasons = response.get("finish_reasons") if response else []
    if not isinstance(finish_reasons, list):
        finish_reasons = []
    failure = failure_reason(response, judge)
    score = parsed.get("judge_score")
    if not isinstance(score, (int, float)):
        score = None
    correct = bool(
        answered
        and failure is None
        and (
            score == 1.0 if score is not None else parsed.get("correct") is True
        )
    )
    return cap_value(
        {
            "qid": gold["qid"],
            "gid": gold["gid"],
            "product": gold["product"],
            "type": gold.get("type", ""),
            "question": gold.get("question", ""),
            "gold": gold.get("ground_truth"),
            "citations": gold.get("citations", []),
            "prediction": prediction,
            "extracted_answer": (
                parsed.get("extracted_final_answer")
                if parsed.get("extracted_final_answer") is not None
                else parsed.get("extracted_answers")
            ),
            "answered": answered,
            "correct": correct,
            "score": score,
            "score_kind": parsed.get("score_kind"),
            "score_details": {
                key: parsed.get(key)
                for key in ("precision", "recall", "f1", "judge_score")
                if parsed.get(key) is not None
            },
            "judge_text": judge.get("judge_text") if judge else None,
            "confidence": parsed.get("confidence") if answered else None,
            "stop_reason": response.get("stop_reason") if response else None,
            "finish_reason": response.get("finish_reason") if response else None,
            "finish_reasons": finish_reasons,
            "failure": failure,
            "tokens": response_tokens(response),
            "turns": response.get("turns") if response else None,
            "tool_counts": tool_counts,
            "events": response.get("events") if response else [],
        }
    )


def index_item(record: dict[str, Any]) -> dict[str, Any]:
    return {
        "qid": record["qid"],
        "gid": record["gid"],
        "product": record["product"],
        "type": record["type"],
        "question": record["question"],
        "prediction": (record["prediction"] or "")[:1000],
        "extracted_answer": (record["extracted_answer"] or "")[:1000],
        "answered": record["answered"],
        "correct": record["correct"],
        "score": record.get("score"),
        "score_kind": record.get("score_kind"),
        "failure": record["failure"],
    }


def write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, ensure_ascii=False, separators=(",", ":"))
        handle.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--eval", type=Path, default=ROOT / "eval.json")
    parser.add_argument("--out", type=Path, default=ROOT)
    args = parser.parse_args()

    eval_rows = json.loads(args.eval.read_text(encoding="utf-8"))
    canonical: list[dict[str, Any]] = []
    for row in eval_rows:
        if row.get("kind") != "answerable":
            continue
        # HERB eval gids contain exactly one separator; run outputs use "_".
        qid = row["gid"].replace("#", "_", 1)
        canonical.append({**row, "qid": qid})
    if len(canonical) != 815 or len({row["qid"] for row in canonical}) != 815:
        raise ValueError(
            f"expected 815 unique answerable HERB qids, found {len(canonical)}"
        )

    out_root = args.out.resolve()
    runs_root = out_root / "runs"
    compare_root = out_root / "compare"
    for generated in (runs_root, compare_root):
        if generated.exists():
            shutil.rmtree(generated)

    manifest_runs: list[dict[str, Any]] = []
    compare_by_qid: dict[str, dict[str, Any]] = {
        gold["qid"]: {
            "qid": gold["qid"],
            "gid": gold["gid"],
            "product": gold["product"],
            "type": gold.get("type", ""),
            "question": gold.get("question", ""),
            "gold": gold.get("ground_truth"),
            "citations": gold.get("citations", []),
            "runs": {},
        }
        for gold in canonical
    }
    canonical_by_qid = {gold["qid"]: gold for gold in canonical}

    for config in RUNS:
        slot = config["slot"]
        response_path = Path(config["response"])
        recovery_response = config.get("recovery_response")
        recovery_responses = (
            load_herb_rows(Path(recovery_response), include_events=True)
            if recovery_response
            else {}
        )
        judges = load_herb_rows(Path(config["judge"]), include_events=False)
        unknown = (set(recovery_responses) | set(judges)) - set(compare_by_qid)
        if unknown:
            raise ValueError(f"{slot}: {len(unknown)} non-canonical HERB qid(s)")

        index_items: dict[str, dict[str, Any]] = {}
        emitted: set[str] = set()
        answered = 0
        correct = 0
        score_total = 0.0
        f1_scores: list[float] = []

        def emit(qid: str, response: dict[str, Any] | None) -> None:
            nonlocal answered, correct, score_total
            if qid in emitted:
                raise ValueError(f"{slot}: duplicate emitted HERB qid {qid!r}")
            emitted.add(qid)
            record = normalized_record(
                canonical_by_qid[qid], response, judges.get(qid)
            )
            answered += int(record["answered"])
            correct += int(record["correct"])
            if isinstance(record.get("score"), (int, float)):
                score_total += float(record["score"])
            if record.get("score_kind") == "answer_f1" and isinstance(
                record.get("score_details", {}).get("f1"), (int, float)
            ):
                f1_scores.append(float(record["score_details"]["f1"]))
            index_items[qid] = index_item(record)
            write_json(runs_root / slot / "records" / f"{qid}.json", record)
            compare_by_qid[qid]["runs"][slot] = {
                key: record[key]
                for key in (
                    "prediction",
                    "extracted_answer",
                    "answered",
                    "correct",
                    "score",
                    "score_kind",
                    "score_details",
                    "judge_text",
                    "confidence",
                    "stop_reason",
                    "finish_reason",
                    "finish_reasons",
                    "failure",
                    "tokens",
                    "turns",
                    "tool_counts",
                )
            }

        for qid, response in iter_herb_rows(response_path, include_events=True):
            if qid not in canonical_by_qid:
                raise ValueError(f"{slot}: non-canonical HERB qid {qid!r}")
            if qid not in recovery_responses:
                emit(qid, response)
        for qid, response in recovery_responses.items():
            emit(qid, response)
        for gold in canonical:
            if gold["qid"] not in emitted:
                emit(gold["qid"], None)

        score_mode = config.get("score_mode")
        score_pct = (
            round(score_total * 100 / len(canonical), 2)
            if score_mode == "mean_judge_score"
            else round(correct * 100 / len(canonical), 2)
        )
        score_detail = (
            f"{correct} perfect · "
            f"{(sum(f1_scores) * 100 / len(f1_scores)):.2f}% answer F1"
            if score_mode == "mean_judge_score" and f1_scores
            else f"{correct} / {len(canonical)} correct"
        )
        summary = {
            "slot": slot,
            "label": config["label"],
            "description": config["description"],
            "agentic": bool(config.get("agentic")),
            "scope": len(canonical),
            "answered": answered,
            "correct": correct,
            "score_pct": score_pct,
            "score_label": (
                "Type-aware judge" if score_mode == "mean_judge_score" else "Score"
            ),
            "score_detail": score_detail,
            "coverage_pct": round(answered * 100 / len(canonical), 2),
            "response_source": str(response_path),
            "recovery_response_source": recovery_response,
            "judge_source": config["judge"],
            "items": [index_items[gold["qid"]] for gold in canonical],
        }
        write_json(runs_root / slot / "index.json", summary)
        manifest_runs.append({key: summary[key] for key in summary if key != "items"})
        print(
            f"{slot}: score={summary['score_pct']:.2f}%, perfect/correct "
            f"{correct}/{len(canonical)}, "
            f"answered {answered}/{len(canonical)}={summary['coverage_pct']:.2f}%"
        )

    compare_items: list[dict[str, Any]] = []
    slots = [config["slot"] for config in RUNS]
    for gold in canonical:
        record = compare_by_qid[gold["qid"]]
        statuses = [
            (record["runs"][slot]["answered"], record["runs"][slot]["correct"])
            for slot in slots
        ]
        record["any_missing"] = any(not answered for answered, _ in statuses)
        record["disagreement"] = len({correct for _, correct in statuses}) > 1
        record["only_e2e_combined_correct"] = (
            record["runs"]["e2e_combined"]["correct"]
            and all(
                not record["runs"][slot]["correct"]
                for slot in slots
                if slot != "e2e_combined"
            )
        )
        write_json(compare_root / "records" / f"{gold['qid']}.json", record)
        compare_items.append(
            {
                "qid": record["qid"],
                "gid": record["gid"],
                "product": record["product"],
                "type": record["type"],
                "question": record["question"],
                "any_missing": record["any_missing"],
                "disagreement": record["disagreement"],
                "only_e2e_combined_correct": record[
                    "only_e2e_combined_correct"
                ],
            }
        )

    write_json(
        compare_root / "index.json",
        {
            "scope": len(canonical),
            "slots": slots,
            "runs": manifest_runs,
            "items": compare_items,
        },
    )
    write_json(
        runs_root / "manifest.json",
        {
            "scope": len(canonical),
            "scope_label": "815 answerable HERB questions",
            "scoring": (
                "Run-specific canonical score over all 815 questions; "
                "missing or unanswered is zero"
            ),
            "runs": manifest_runs,
        },
    )


if __name__ == "__main__":
    main()