scripts/build_trajectories.py

#!/usr/bin/env python3
"""Build open-wikitable `trajectories/` (or `trajectories_corpus/`) shards.

Companion to ``build_trajectories_e2e.py`` but for the canonical *standalone*
agentic_answer responses (not the e2e pipeline). Reads one
``response.jsonl`` from ``information-scaffolds/outputs/agentic_wiki_opentable/``
(each row carries the full agent loop in ``events``), joins against the
wiki_opentable gold (``answer_list`` in ``test_with_chunks.jsonl``), scores
each prediction with the canonical set-based F1 / EM (semicolon parser),
and writes one shard per qid plus an ``index.json`` with summary stats so
the sidebar filter/sort can work without paying for every shard.

Default output is ``./trajectories``. Pass ``--out trajectories_corpus`` for
the flat-corpus (Baseline B) variant.

Per-event tool-result content is truncated at 8 KB to keep each shard
browser-friendly (same cap as build_trajectories_e2e.py).

Usage:
    python scripts/build_trajectories.py \\
        --predictions <cell4_agentic_a.response.jsonl> \\
        --gold        $DATA_ROOT/eval/wiki_opentable/raw/test_with_chunks.jsonl \\
        --out         trajectories \\
        --label       "Baseline A — per-qid scaffolds (cell4)"
"""
from __future__ import annotations

import argparse
import json
import re
import shutil
import sys
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple


# ─── Parsing the model's Exact-Answer block (mirrors _parse_exact_answer_semicolon.py) ──

_EXACT_ANSWER_RE = re.compile(
    r"Exact\s*Answer\s*:\s*(.*?)(?:\n\s*Confidence\s*:|\Z)",
    re.IGNORECASE | re.DOTALL,
)


def _strip_uncertainty(s: str) -> str:
    s = s.strip()
    if len(s) >= 2 and s[0] in '"\u201c\u201d\'' and s[-1] in '"\u201c\u201d\'':
        s = s[1:-1].strip()
    while s.endswith("?"):
        s = s[:-1].rstrip()
    return s


def extract_answer_payload(model_answer: str) -> str:
    if not model_answer:
        return ""
    m = _EXACT_ANSWER_RE.search(model_answer)
    if not m:
        return model_answer.strip()
    return m.group(1).strip()


def extract_answer_items(model_answer: str) -> List[str]:
    payload = extract_answer_payload(model_answer)
    if not payload:
        return []
    stripped = payload.lstrip()
    if stripped.startswith("["):
        end = stripped.rfind("]")
        if end > 0:
            candidate = stripped[: end + 1]
            try:
                parsed = json.loads(candidate)
                if isinstance(parsed, list):
                    items = [
                        _strip_uncertainty(str(x))
                        for x in parsed
                        if x is not None and str(x).strip()
                    ]
                    return [x for x in items if x]
            except json.JSONDecodeError:
                pass
    one_line = payload.splitlines()[0].strip()
    if ";" in one_line:
        items = [_strip_uncertainty(p) for p in one_line.split(";")]
        items = [x for x in items if x]
        if len(items) >= 2:
            return items
    s = _strip_uncertainty(payload)
    return [s] if s else []


# ─── Set-based scoring (mirrors wiki_opentable_adapter.py) ────────────────────

_PUNCT_TRAIL = ".,;:!?"
_QUOTE_CHARS = "\"'\u201c\u201d\u2018\u2019"
_WS_RE = re.compile(r"\s+")
_NUM_RE = re.compile(r"^-?\d{1,3}(?:,\d{3})*(?:\.\d+)?$|^-?\d+(?:\.\d+)?$")


def _normalize(s: str) -> str:
    if s is None:
        return ""
    t = str(s).strip()
    if len(t) >= 2 and t[0] in _QUOTE_CHARS and t[-1] in _QUOTE_CHARS:
        t = t[1:-1].strip()
    while t and t[-1] in _PUNCT_TRAIL:
        t = t[:-1].rstrip()
    t = _WS_RE.sub(" ", t).strip().lower()
    if _NUM_RE.match(t):
        t2 = t.replace(",", "")
        try:
            f = float(t2)
            if f.is_integer():
                return str(int(f))
            return str(f)
        except ValueError:
            return t2
    return t


def score_one(pred_items: List[str], gold_items: List[str]) -> Tuple[float, float, float, int]:
    P = {_normalize(x) for x in pred_items if _normalize(x)}
    G = {_normalize(x) for x in gold_items if _normalize(x)}
    if not P and not G:
        return 1.0, 1.0, 1.0, 1
    if not P:
        return 0.0, 0.0, 0.0, 0
    hit = len(P & G)
    prec = hit / len(P)
    rec = hit / len(G) if G else 0.0
    f1 = (2 * prec * rec / (prec + rec)) if (prec + rec) else 0.0
    em = 1 if P == G else 0
    return prec, rec, f1, em


# ─── Event truncation ─────────────────────────────────────────────────────────

_CONTENT_CAP = 8000


def _truncate_content(s: str) -> Tuple[str, bool]:
    if s is None:
        return "", False
    if len(s) <= _CONTENT_CAP:
        return s, False
    head = s[: _CONTENT_CAP // 2]
    tail = s[-_CONTENT_CAP // 2 :]
    return f"{head}\n\n… [truncated {len(s) - _CONTENT_CAP} chars] …\n\n{tail}", True


def project_event(ev: Dict[str, Any]) -> Dict[str, Any]:
    content, trunc = _truncate_content(ev.get("content"))
    out = {
        "type": ev.get("type"),
        "name": ev.get("name"),
        "input": ev.get("input"),
        "content": content,
    }
    if trunc:
        out["truncated"] = True
    return out


# ─── I/O ──────────────────────────────────────────────────────────────────────


def load_predictions(path: Path) -> Dict[str, Dict[str, Any]]:
    out: Dict[str, Dict[str, Any]] = {}
    with path.open() as f:
        for line in f:
            d = json.loads(line)
            if d.get("dataset") != "wiki_opentable":
                continue
            out[str(d["qid"])] = d
    return out


def load_gold(path: Path) -> Dict[str, Dict[str, Any]]:
    out: Dict[str, Dict[str, Any]] = {}
    with path.open() as f:
        for line in f:
            d = json.loads(line)
            out[str(d["qid"])] = d
    return out


# ─── Build ────────────────────────────────────────────────────────────────────


def build_record(
    qid: str,
    pred: Dict[str, Any],
    gold_row: Dict[str, Any],
) -> Dict[str, Any]:
    gold_items = list(gold_row.get("answer_list", []))
    pred_items = extract_answer_items(pred.get("answer", ""))
    prec, rec, f1, em = score_one(pred_items, gold_items)

    events = [project_event(ev) for ev in (pred.get("events") or [])]

    return {
        "qid": qid,
        "dataset": "wiki_opentable",
        "dataset_origin": gold_row.get("dataset_origin"),
        "original_table_id": gold_row.get("original_table_id"),
        "question": pred.get("question") or gold_row.get("question_text"),
        "gold_answers": gold_items,
        "sql": gold_row.get("sql"),
        "model": pred.get("model"),
        "mode": pred.get("mode"),
        "system_prompt_file": pred.get("system_prompt_file"),
        "max_turns": pred.get("max_turns"),
        "max_completion_tokens": pred.get("max_completion_tokens"),
        "stop_reason": pred.get("stop_reason"),
        "error": pred.get("error"),
        "finish_reasons": pred.get("finish_reasons"),
        "n_turns": pred.get("turns"),
        "tool_call_counts": pred.get("tool_call_counts"),
        "tokens": pred.get("tokens"),
        "latency_ms": pred.get("latency_ms"),
        "attempts": pred.get("attempts"),
        "timeout_retries": pred.get("timeout_retries"),
        "model_answer": pred.get("answer"),
        "pred_items": pred_items,
        "metrics": {
            "precision": round(prec, 4),
            "recall": round(rec, 4),
            "f1": round(f1, 4),
            "em": em,
        },
        "events": events,
    }


def build_index_row(rec: Dict[str, Any]) -> Dict[str, Any]:
    m = rec["metrics"]
    return {
        "qid": rec["qid"],
        "question": rec["question"],
        "dataset_origin": rec["dataset_origin"],
        "n_turns": rec["n_turns"],
        "stop_reason": rec["stop_reason"],
        "has_error": (
            (rec["error"] is not None and rec["error"] != "")
            or rec["stop_reason"] == "error"
        ),
        "f1": m["f1"],
        "em": m["em"],
        "n_pred": len(rec["pred_items"]),
        "n_gold": len(rec["gold_answers"]),
    }


def _counter(it: Iterable[Any]) -> Dict[str, int]:
    out: Dict[str, int] = {}
    for x in it:
        k = str(x)
        out[k] = out.get(k, 0) + 1
    return dict(sorted(out.items(), key=lambda kv: -kv[1]))


def main() -> int:
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--predictions", required=True, type=Path)
    ap.add_argument(
        "--gold",
        type=Path,
        default=Path("/mnt/ramdisk/blobstore/timchen0618/data/eval/wiki_opentable/raw/test_with_chunks.jsonl"),
    )
    ap.add_argument("--out", required=True, type=Path,
                    help="Output dir (e.g. trajectories/ or trajectories_corpus/)")
    ap.add_argument("--label", default="standalone agentic run",
                    help="Human label for this run; written to index meta")
    args = ap.parse_args()

    preds = load_predictions(args.predictions)
    gold = load_gold(args.gold)

    common = sorted(set(preds) & set(gold))
    print(f"qids: predictions={len(preds)}, gold={len(gold)}, common={len(common)}")
    if not common:
        print("ERROR: no overlap", file=sys.stderr)
        return 1

    out_dir = args.out
    out_dir.mkdir(parents=True, exist_ok=True)
    rec_dir = out_dir / "records"
    if rec_dir.exists():
        shutil.rmtree(rec_dir)
    rec_dir.mkdir(parents=True)

    index_rows: List[Dict[str, Any]] = []
    sum_f1 = sum_em = 0.0
    for qid in common:
        rec = build_record(qid, preds[qid], gold[qid])
        (rec_dir / f"{qid}.json").write_text(json.dumps(rec, ensure_ascii=False))
        index_rows.append(build_index_row(rec))
        sum_f1 += rec["metrics"]["f1"]
        sum_em += rec["metrics"]["em"]

    n = len(common)
    summary = {
        "label": args.label,
        "n": n,
        "mean_f1": round(sum_f1 / n, 4),
        "mean_em": round(sum_em / n, 4),
        "fraction_error": round(
            sum(1 for r in index_rows if r["has_error"]) / n, 4
        ),
        "stop_reason_counts": _counter(r["stop_reason"] for r in index_rows),
    }
    index_payload = {"meta": summary, "rows": index_rows}
    (out_dir / "index.json").write_text(json.dumps(index_payload, ensure_ascii=False))

    print(f"\n✓ Wrote {out_dir}/index.json + {n} record shards")
    print(f"   mean F1 = {summary['mean_f1']*100:.2f}")
    print(f"   mean EM = {summary['mean_em']*100:.2f}")
    print(f"   error rate = {summary['fraction_error']*100:.2f}")
    return 0


if __name__ == "__main__":
    sys.exit(main())