scripts/build_viewer.py

#!/usr/bin/env python3
"""Build the static PhantomWiki 100Q viewer data bundle."""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterator


ROOT = Path(__file__).resolve().parents[1]
DATASET = "phantom_wiki"
UNIVERSE = "depth_20_size_5000_seed_1"
EXPECTED_QIDS = 100
EXPECTED_QID_SHA256 = (
    "69ac438794500cb64239009c3cb3a4d70f68b2073d29f592e246db1edd84cd6c"
)
STRING_LIMIT = 12_000
TRAJECTORY_LIMIT_BYTES = 384 * 1024
INPUT_USD_PER_MILLION = 2.5
OUTPUT_USD_PER_MILLION = 15.0
EXACT_ANSWER_RE = re.compile(
    r"^\s*Exact Answer:\s*(.*?)\s*$",
    re.IGNORECASE | re.MULTILINE,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--qid-source", type=Path, required=True)
    parser.add_argument("--corpus", type=Path, required=True)
    parser.add_argument("--eval", dest="eval_path", type=Path, required=True)
    parser.add_argument("--structures-root", type=Path, required=True)
    parser.add_argument("--run-config", type=Path, required=True)
    parser.add_argument("--out", type=Path, default=ROOT)
    parser.add_argument("--expected-qid-sha256", default=EXPECTED_QID_SHA256)
    return parser.parse_args()


def iter_jsonl(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:
                value = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_number}: {exc}") from exc
            if not isinstance(value, dict):
                raise ValueError(f"{path}:{line_number}: expected an object")
            yield value


def read_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, value: Any, *, indent: int | None = None) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(value, ensure_ascii=False, indent=indent) + ("\n" if indent else ""),
        encoding="utf-8",
    )


def qid_digest(qids: set[str]) -> str:
    return hashlib.sha256("\n".join(sorted(qids)).encode()).hexdigest()


def load_qids(path: Path, expected_sha256: str) -> set[str]:
    qids = {
        str(row["qid"])
        for row in iter_jsonl(path)
        if row.get("dataset") == DATASET and row.get("qid")
    }
    if len(qids) != EXPECTED_QIDS:
        raise ValueError(f"{path}: expected {EXPECTED_QIDS} qids, got {len(qids)}")
    actual = qid_digest(qids)
    if actual != expected_sha256:
        raise ValueError(
            f"{path}: qid SHA-256 mismatch: expected {expected_sha256}, got {actual}"
        )
    return qids


def normalize_answer(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", text.lower())


def extract_exact_answer(text: str) -> str:
    match = EXACT_ANSWER_RE.search(str(text or ""))
    return match.group(1).strip() if match else str(text or "").strip()


def score_prediction(gold: list[Any], prediction: str) -> dict[str, Any]:
    extracted = extract_exact_answer(prediction)
    predicted_items = [item.strip() for item in extracted.split(";") if item.strip()]
    predicted_set = {
        normalize_answer(item) for item in predicted_items if normalize_answer(item)
    }
    gold_set = {
        normalize_answer(str(item)) for item in gold if normalize_answer(str(item))
    }
    hit = len(predicted_set & gold_set)
    precision = hit / len(predicted_set) if predicted_set else 0.0
    recall = hit / len(gold_set) if gold_set else 0.0
    f1 = (
        2 * precision * recall / (precision + recall)
        if precision + recall
        else 0.0
    )
    return {
        "extracted_answer": extracted,
        "exact": bool(gold_set) and predicted_set == gold_set,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "n_predicted": len(predicted_set),
        "n_gold": len(gold_set),
        "n_hit": hit,
    }


def cap_string(value: str, limit: int = STRING_LIMIT) -> str:
    if len(value) <= limit:
        return value
    marker = f"\n... [truncated from {len(value)} chars]"
    return value[: max(0, limit - len(marker))] + marker


def cap_value(value: Any) -> Any:
    if isinstance(value, str):
        return cap_string(value)
    if isinstance(value, list):
        return [cap_value(item) for item in value]
    if isinstance(value, tuple):
        return [cap_value(item) for item in value]
    if isinstance(value, dict):
        return {cap_string(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 []
    compacted: list[dict[str, Any]] = []
    used = 2
    original_bytes = len(
        json.dumps(events, ensure_ascii=False, separators=(",", ":")).encode()
    )
    for index, raw in enumerate(events):
        if isinstance(raw, dict):
            event = {
                key: cap_value(raw[key])
                for key in ("type", "name", "input", "content")
                if raw.get(key) is not None
            }
        else:
            event = {"type": "event", "content": cap_value(raw)}
        encoded = json.dumps(
            event, ensure_ascii=False, separators=(",", ":")
        ).encode()
        if used + len(encoded) + 1 > TRAJECTORY_LIMIT_BYTES:
            compacted.append(
                {
                    "type": "truncated",
                    "content": (
                        f"{len(events) - index} events omitted; original trajectory "
                        f"was {original_bytes:,} bytes and the viewer limit is "
                        f"{TRAJECTORY_LIMIT_BYTES:,} bytes."
                    ),
                }
            )
            break
        compacted.append(event)
        used += len(encoded) + 1
    return compacted


def normalized_tokens(response: dict[str, Any] | None) -> dict[str, int]:
    raw = (response or {}).get("tokens")
    if not isinstance(raw, dict):
        raw = (response or {}).get("usage")
    if not isinstance(raw, dict):
        return {}

    def integer(*keys: str) -> int:
        for key in keys:
            value = raw.get(key)
            if isinstance(value, (int, float)):
                return int(value)
        return 0

    prompt = integer("prompt_tokens", "prompt", "input_tokens", "input")
    completion = integer(
        "completion_tokens", "completion", "output_tokens", "output"
    )
    reasoning = integer("reasoning_tokens", "reasoning")
    total = integer("total_tokens", "total") or prompt + completion
    return {
        "prompt_tokens": prompt,
        "completion_tokens": completion,
        "reasoning_tokens": reasoning,
        "total_tokens": total,
    }


def failure_reason(response: dict[str, Any] | None, answered: bool) -> str | None:
    if answered:
        return None
    if not response:
        return "missing_response"
    if response.get("stop_reason"):
        return str(response["stop_reason"])
    finish_reasons = response.get("finish_reasons")
    if isinstance(finish_reasons, list) and finish_reasons:
        return str(finish_reasons[-1])
    for attempt in reversed(response.get("attempt_log") or []):
        if not isinstance(attempt, dict):
            continue
        if attempt.get("exc"):
            return f"exc:{str(attempt['exc']).split(':', 1)[0]}"
        if attempt.get("outcome") and attempt.get("outcome") != "success":
            return str(attempt["outcome"])
    return "unanswered"


def status_for(answered: bool, exact: bool, f1: float) -> str:
    if not answered:
        return "missing"
    if exact:
        return "correct"
    if f1 > 0:
        return "partial"
    return "incorrect"


def load_predictions(path: Path) -> dict[str, dict[str, Any]]:
    rows: dict[str, dict[str, Any]] = {}
    for row in iter_jsonl(path):
        if row.get("dataset") != DATASET:
            continue
        qid = str(row.get("qid") or "")
        if not qid:
            raise ValueError(f"{path}: PhantomWiki row has no qid")
        if qid in rows:
            raise ValueError(f"{path}: duplicate qid {qid}")
        rows[qid] = row
    return rows


def build_corpus_and_eval(
    *,
    corpus_path: Path,
    eval_path: Path,
    qids: set[str],
    data_root: Path,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    corpus_rows = read_json(corpus_path)
    eval_source = read_json(eval_path)
    if not isinstance(corpus_rows, list) or not isinstance(eval_source, list):
        raise ValueError("corpus and eval sources must both be JSON lists")
    eval_rows = [row for row in eval_source if str(row.get("id")) in qids]
    found = {str(row["id"]) for row in eval_rows}
    if found != qids:
        missing = sorted(qids - found)
        raise ValueError(f"{eval_path}: missing {len(missing)} qids: {missing[:5]}")
    write_json(data_root / "corpus.json", corpus_rows)
    write_json(data_root / "eval.json", eval_rows)
    return corpus_rows, eval_rows


def build_structures(
    *,
    source_root: Path,
    qids: set[str],
    eval_order: dict[str, int],
    data_root: Path,
) -> dict[str, Any]:
    source_index = read_json(source_root / "index.json")
    source_rows = source_index.get("rows") or []
    selected = [row for row in source_rows if str(row.get("qid")) in qids]
    selected.sort(key=lambda row: eval_order[str(row["qid"])])
    selected_qids = {str(row["qid"]) for row in selected}
    if selected_qids != qids:
        missing = sorted(qids - selected_qids)
        raise ValueError(f"{source_root}: missing structure qids: {missing[:5]}")

    destination = data_root / "structures"
    records_root = destination / "records"
    records_root.mkdir(parents=True, exist_ok=True)
    extraction_docs: set[str] = set()
    source_docs: set[str] = set()
    total_structures = 0
    output_rows: list[dict[str, Any]] = []
    for row in selected:
        source_path = source_root / str(row["path"])
        record_name = source_path.name
        shutil.copy2(source_path, records_root / record_name)
        copied = dict(row)
        copied["path"] = f"records/{record_name}"
        output_rows.append(copied)
        extraction_docs.update(str(value) for value in row.get("doc_ids") or [])
        source_docs.update(str(value) for value in row.get("source_doc_ids") or [])
        total_structures += int(row.get("n_structures") or 0)

    meta = {
        key: value
        for key, value in dict(source_index.get("meta") or {}).items()
        if key not in {"scaffolds_dir", "unified", "bundle_manifest"}
    }
    meta.update(
        {
            "label": "E2E Structures v3 - fixed 100Q subset",
            "n_qids": len(output_rows),
            "n_extraction_docs": len(extraction_docs),
            "n_source_docs": len(source_docs),
            "n_structures_total": total_structures,
            "source_scope": UNIVERSE,
            "source": "Shared canonical E2E v3 scaffold artifacts",
            "qid_sha256": qid_digest(qids),
        }
    )
    write_json(destination / "index.json", {"meta": meta, "rows": output_rows})
    return meta


def normalize_run_record(
    *,
    run: dict[str, Any],
    gold: dict[str, Any],
    response: dict[str, Any] | None,
) -> tuple[dict[str, Any], dict[str, Any]]:
    prediction = str((response or {}).get("answer") or "").strip()
    score = score_prediction(list(gold.get("answer") or []), prediction)
    answered = bool(prediction)
    tokens = normalized_tokens(response)
    tool_counts = (response or {}).get("tool_call_counts")
    tool_counts = tool_counts if isinstance(tool_counts, dict) else {}
    tool_calls = sum(
        value for value in tool_counts.values() if isinstance(value, int)
    )
    failure = failure_reason(response, answered)
    status = status_for(answered, score["exact"], score["f1"])
    metadata = gold.get("meta") if isinstance(gold.get("meta"), dict) else {}
    execution = {
        "answered": answered,
        "status": status,
        "stop_reason": (response or {}).get("stop_reason"),
        "finish_reasons": (response or {}).get("finish_reasons") or [],
        "failure_reason": failure,
        "latency_ms": (response or {}).get("latency_ms"),
        "turns": (response or {}).get("turns"),
        "tool_calls": tool_calls,
        "tool_call_counts": tool_counts,
        "tokens": tokens,
        "attempts": (response or {}).get("attempts"),
    }
    provenance = {
        "job": run.get("job"),
        "variant": run.get("variant"),
        "prompt": run.get("prompt"),
        "model": (response or {}).get("model"),
        "mode": (response or {}).get("mode"),
        "reasoning_effort": (response or {}).get("reasoning_effort"),
        "relation_format": (response or {}).get("relation_format"),
        "rendered_prompt_sha256": (response or {}).get("rendered_prompt_sha256"),
        "rendered_prompt_chars": (response or {}).get("rendered_prompt_chars"),
        "system_prompt_file": (response or {}).get("system_prompt_file"),
        "canonical_prompt_cell": (response or {}).get("canonical_prompt_cell"),
    }
    result = cap_value(
        {
            "qid": str(gold["id"]),
            "run_slot": run["slot"],
            "run_label": run["label"],
            "question": gold.get("question") or "",
            "gold": gold.get("answer") or [],
            "metadata": {
                "difficulty": metadata.get("Difficulty"),
                "type": metadata.get("Type"),
                "template": gold.get("template") or "",
                "prolog": gold.get("prolog") or [],
                "prolog_answer": gold.get("prolog_answer") or "",
                "supporting_titles": gold.get("supporting_titles") or [],
            },
            "prediction": prediction,
            "score": score,
            "execution": execution,
            "provenance": provenance,
        }
    )
    trajectory = cap_value(
        {
            "qid": str(gold["id"]),
            "run_slot": run["slot"],
            "run_label": run["label"],
            "question": gold.get("question") or "",
            "gold": gold.get("answer") or [],
            "prediction": prediction,
            "score": score,
            "execution": execution,
            "provenance": provenance,
            "attempt_log": (response or {}).get("attempt_log") or [],
            "events": compact_events((response or {}).get("events")),
        }
    )
    return result, trajectory


def build_run(
    *,
    run: dict[str, Any],
    eval_rows: list[dict[str, Any]],
    qids: set[str],
    data_root: Path,
) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]:
    prediction_path = Path(run["prediction"])
    predictions = load_predictions(prediction_path)
    selected_predictions = set(predictions) & qids
    unexpected = set(predictions) - qids
    if run.get("subset_only") and unexpected:
        raise ValueError(
            f"{run['slot']}: expected only subset qids, found {len(unexpected)} extras"
        )
    if len(selected_predictions) != EXPECTED_QIDS:
        missing = sorted(qids - selected_predictions)
        raise ValueError(
            f"{run['slot']}: expected {EXPECTED_QIDS} prediction records, "
            f"got {len(selected_predictions)}; missing {missing[:5]}"
        )

    run_root = data_root / "runs" / run["slot"]
    result_root = run_root / "results"
    trajectory_root = run_root / "trajectories"
    result_root.mkdir(parents=True, exist_ok=True)
    trajectory_root.mkdir(parents=True, exist_ok=True)
    summaries: list[dict[str, Any]] = []
    records: dict[str, dict[str, Any]] = {}
    precision_total = recall_total = f1_total = 0.0
    answered = exact = 0
    prompt_tokens = completion_tokens = reasoning_tokens = total_tokens = 0

    for gold in eval_rows:
        qid = str(gold["id"])
        result, trajectory = normalize_run_record(
            run=run, gold=gold, response=predictions.get(qid)
        )
        records[qid] = result
        write_json(result_root / f"{qid}.json", result)
        write_json(trajectory_root / f"{qid}.json", trajectory)
        score = result["score"]
        execution = result["execution"]
        tokens = execution["tokens"]
        answered += int(execution["answered"])
        exact += int(score["exact"])
        precision_total += score["precision"]
        recall_total += score["recall"]
        f1_total += score["f1"]
        prompt_tokens += int(tokens.get("prompt_tokens") or 0)
        completion_tokens += int(tokens.get("completion_tokens") or 0)
        reasoning_tokens += int(tokens.get("reasoning_tokens") or 0)
        total_tokens += int(tokens.get("total_tokens") or 0)
        summaries.append(
            {
                "qid": qid,
                "question": result["question"],
                "gold": result["gold"],
                "status": execution["status"],
                "answered": execution["answered"],
                "exact": score["exact"],
                "f1": score["f1"],
                "failure_reason": execution["failure_reason"],
            }
        )

    total = len(eval_rows)
    cost = (
        prompt_tokens * INPUT_USD_PER_MILLION
        + completion_tokens * OUTPUT_USD_PER_MILLION
    ) / 1_000_000
    index = {
        "slot": run["slot"],
        "label": run["label"],
        "variant": run.get("variant"),
        "family": run.get("family", "experiment"),
        "accent": run.get("accent", "#8b5cf6"),
        "job": run.get("job"),
        "prompt": run.get("prompt"),
        "total": total,
        "answered": answered,
        "exact": exact,
        "coverage": answered / total,
        "mean_precision": precision_total / total,
        "mean_recall": recall_total / total,
        "mean_f1": f1_total / total,
        "cost_total_usd": cost,
        "cost_per_question_usd": cost / total,
        "tokens": {
            "prompt": prompt_tokens,
            "completion": completion_tokens,
            "reasoning": reasoning_tokens,
            "total": total_tokens,
        },
        "records": summaries,
    }
    write_json(run_root / "index.json", index)
    return index, records


def compare_projection(record: dict[str, Any]) -> dict[str, Any]:
    return {
        "prediction": record["prediction"],
        "score": record["score"],
        "execution": record["execution"],
        "provenance": record["provenance"],
    }


def build_compare(
    *,
    eval_rows: list[dict[str, Any]],
    run_indexes: list[dict[str, Any]],
    records_by_run: dict[str, dict[str, dict[str, Any]]],
    data_root: Path,
) -> None:
    compare_root = data_root / "compare"
    compare_records_root = compare_root / "records"
    compare_records_root.mkdir(parents=True, exist_ok=True)
    summaries: list[dict[str, Any]] = []
    slots = [index["slot"] for index in run_indexes]
    for gold in eval_rows:
        qid = str(gold["id"])
        records = {slot: records_by_run[slot][qid] for slot in slots}
        f1_values = {slot: record["score"]["f1"] for slot, record in records.items()}
        exact_values = {
            slot: record["score"]["exact"] for slot, record in records.items()
        }
        answered_values = {
            slot: record["execution"]["answered"] for slot, record in records.items()
        }
        best_f1 = max(f1_values.values()) if f1_values else 0.0
        best_slots = [
            slot for slot, value in f1_values.items() if abs(value - best_f1) < 1e-12
        ]
        summary = {
            "qid": qid,
            "question": gold.get("question") or "",
            "f1": f1_values,
            "exact": exact_values,
            "answered": answered_values,
            "best_f1": best_f1,
            "best_slots": best_slots,
            "disagreement": len({round(value, 12) for value in f1_values.values()}) > 1,
            "any_missing": not all(answered_values.values()),
        }
        summaries.append(summary)
        write_json(
            compare_records_root / f"{qid}.json",
            {
                "qid": qid,
                "question": gold.get("question") or "",
                "gold": gold.get("answer") or [],
                "metadata": records[slots[0]]["metadata"],
                "runs": {
                    slot: compare_projection(record)
                    for slot, record in records.items()
                },
            },
        )
    write_json(
        compare_root / "index.json",
        {"total": len(summaries), "slots": slots, "records": summaries},
    )


def main() -> None:
    args = parse_args()
    qids = load_qids(args.qid_source, args.expected_qid_sha256)
    data_root = args.out / "data"
    if data_root.exists():
        shutil.rmtree(data_root)
    data_root.mkdir(parents=True)

    corpus_rows, eval_rows = build_corpus_and_eval(
        corpus_path=args.corpus,
        eval_path=args.eval_path,
        qids=qids,
        data_root=data_root,
    )
    eval_order = {str(row["id"]): index for index, row in enumerate(eval_rows)}
    structures_meta = build_structures(
        source_root=args.structures_root,
        qids=qids,
        eval_order=eval_order,
        data_root=data_root,
    )

    run_config = read_json(args.run_config)
    if not isinstance(run_config, list) or not run_config:
        raise ValueError(f"{args.run_config}: expected a non-empty JSON list")
    slots = [str(run["slot"]) for run in run_config]
    if len(slots) != len(set(slots)):
        raise ValueError(f"{args.run_config}: duplicate run slots")

    run_indexes: list[dict[str, Any]] = []
    records_by_run: dict[str, dict[str, dict[str, Any]]] = {}
    for run in run_config:
        index, records = build_run(
            run=run,
            eval_rows=eval_rows,
            qids=qids,
            data_root=data_root,
        )
        run_indexes.append(index)
        records_by_run[index["slot"]] = records
        print(
            f"{index['slot']:20s} F1={index['mean_f1'] * 100:6.2f} "
            f"EM={index['exact']:3d}/{index['total']} "
            f"answered={index['answered']:3d}/{index['total']} "
            f"cost/Q=${index['cost_per_question_usd']:.4f}"
        )

    build_compare(
        eval_rows=eval_rows,
        run_indexes=run_indexes,
        records_by_run=records_by_run,
        data_root=data_root,
    )
    manifest = {
        "title": "PhantomWiki Prompt Iteration - Fixed 100Q",
        "dataset": DATASET,
        "universe": UNIVERSE,
        "total_qids": len(qids),
        "qid_sha256": qid_digest(qids),
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "files": {
            "corpus": "data/corpus.json",
            "eval": "data/eval.json",
            "structures": "data/structures/index.json",
            "compare": "data/compare/index.json",
        },
        "corpus_documents": len(corpus_rows),
        "structures": structures_meta,
        "runs": [
            {
                key: index[key]
                for key in (
                    "slot",
                    "label",
                    "variant",
                    "family",
                    "accent",
                    "job",
                    "prompt",
                    "total",
                    "answered",
                    "exact",
                    "coverage",
                    "mean_precision",
                    "mean_recall",
                    "mean_f1",
                    "cost_per_question_usd",
                )
            }
            | {
                "index": f"data/runs/{index['slot']}/index.json",
                "results": f"data/runs/{index['slot']}/results",
                "trajectories": f"data/runs/{index['slot']}/trajectories",
            }
            for index in run_indexes
        ],
    }
    write_json(data_root / "manifest.json", manifest, indent=2)
    write_json(data_root / "qids.json", sorted(qids), indent=2)
    print(
        f"built {args.out}: {len(corpus_rows)} corpus docs, {len(eval_rows)} eval "
        f"questions, {len(run_indexes)} runs"
    )


if __name__ == "__main__":
    main()