WalkXR-AI / eval / aggregate_eval_run.py
aggregate_eval_run.py
Raw
# eval/aggregate_eval_run.py
import json, argparse
from pathlib import Path
from collections import defaultdict
import csv

def load_jsonl(path):
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                yield json.loads(line)

def safe_mean(xs):
    xs = [x for x in xs if x is not None]
    return sum(xs) / len(xs) if xs else None

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--raw_dir", required=True, help="Path to eval/runs/<eval_id>/raw")
    ap.add_argument("--out_dir", help="Defaults to parent/summary")
    # NEW: filtering flags
    ap.add_argument("--include_run_id", nargs="*", help="Only include these run IDs (e.g., 7e30eb 42af10)")
    ap.add_argument("--exclude_run_id", nargs="*", help="Exclude these run IDs")
    args = ap.parse_args()

    raw_dir = Path(args.raw_dir)
    out_dir = Path(args.out_dir) if args.out_dir else raw_dir.parent / "summary"
    out_dir.mkdir(parents=True, exist_ok=True)

    by_scn = defaultdict(lambda: {
        "latencies": [], "refusals": [], "persona_hits": [], "rag_hits": [],
        "memory_recall_hits": [], "continuity_breaks": [],
        "turns": 0, "errors": 0,
    })
    # Collect detailed recall events for the memory summary report
    memory_recall_events = []
    meta = {"eval_id": None, "git_sha": None, "model_ids": set()}

    def run_id_from_filename(p: Path) -> str:
        # expects ..._<runid>.jsonl
        stem = p.stem
        return stem.rsplit("_", 1)[-1] if "_" in stem else "unknown"

    for p in raw_dir.glob("*.jsonl"):
        rid = run_id_from_filename(p)

        # Apply filters
        if args.include_run_id and rid not in args.include_run_id:
            continue
        if args.exclude_run_id and rid in args.exclude_run_id:
            continue

        for row in load_jsonl(p):
            scn = row.get("scenario_id", "unknown")
            by_scn[scn]["turns"] += 1
            by_scn[scn]["latencies"].append(row.get("latency_ms"))
            m = row.get("metrics", {})
            by_scn[scn]["refusals"].append(m.get("refusal"))
            by_scn[scn]["persona_hits"].append(m.get("persona_hit"))
            by_scn[scn]["rag_hits"].append(m.get("rag_hit"))
            # Memory recall metrics (T3.6 Part 2)
            if m.get("memory_is_recall_turn"):
                by_scn[scn]["memory_recall_hits"].append(m.get("memory_recall_hit"))
                memory_recall_events.append({
                    "scenario_id": scn,
                    "turn": row.get("turn"),
                    "expected": m.get("memory_expected"),
                    "recall_hit": m.get("memory_recall_hit"),
                    "qa_check_type": m.get("memory_qa_check_type"),
                    "user_input": row.get("user_input", "")[:200],
                    "agent_response": row.get("agent_response", "")[:300],
                    "run_id": row.get("run_id"),
                    "repeat": row.get("repeat"),
                })
            if m.get("continuity_break") is not None:
                by_scn[scn]["continuity_breaks"].append(m.get("continuity_break"))
            if row.get("status") != "ok":
                by_scn[scn]["errors"] += 1

            if meta["eval_id"] is None:
                meta["eval_id"] = row.get("eval_id")
                meta["git_sha"] = row.get("git_sha")
            if row.get("model_id"):
                meta["model_ids"].add(row["model_id"])

    # write per-scenario CSV
    csv_path = out_dir / "metrics_by_scenario.csv"
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["scenario_id", "turns", "mean_latency_ms", "refusal_rate",
                     "persona_hit_rate", "rag_hit_rate", "memory_recall_rate",
                     "continuity_break_rate", "errors"])
        for scn, d in sorted(by_scn.items()):
            mean_lat = safe_mean(d["latencies"])
            refusal_rate = safe_mean(d["refusals"])
            persona_rate = safe_mean(d["persona_hits"])
            rag_rate = safe_mean(d["rag_hits"])
            mem_rate = safe_mean(d["memory_recall_hits"])
            cont_rate = safe_mean(d["continuity_breaks"])
            w.writerow([scn, d["turns"],
                        round(mean_lat, 1) if mean_lat is not None else None,
                        round(refusal_rate, 3) if refusal_rate is not None else None,
                        round(persona_rate, 3) if persona_rate is not None else None,
                        round(rag_rate, 3) if rag_rate is not None else None,
                        round(mem_rate, 3) if mem_rate is not None else None,
                        round(cont_rate, 3) if cont_rate is not None else None,
                        d["errors"]])

    # write overall summary
    all_lat = [x for d in by_scn.values() for x in d["latencies"] if x is not None]
    all_ref = [x for d in by_scn.values() for x in d["refusals"] if x is not None]
    all_per = [x for d in by_scn.values() for x in d["persona_hits"] if x is not None]
    all_rag = [x for d in by_scn.values() for x in d["rag_hits"] if x is not None]
    all_mem = [x for d in by_scn.values() for x in d["memory_recall_hits"] if x is not None]
    all_cont = [x for d in by_scn.values() for x in d["continuity_breaks"] if x is not None]

    overall = {
        "eval_id": meta["eval_id"],
        "git_sha": meta["git_sha"],
        "model_ids": sorted(meta["model_ids"]),
        "scenarios": len(by_scn),
        "turns": sum(d["turns"] for d in by_scn.values()),
        "mean_latency_ms": round(safe_mean(all_lat), 1) if all_lat else None,
        "refusal_rate": round(safe_mean(all_ref), 3) if all_ref else None,
        "persona_hit_rate": round(safe_mean(all_per), 3) if all_per else None,
        "rag_hit_rate": round(safe_mean(all_rag), 3) if all_rag else None,
        "memory_recall_rate": round(safe_mean(all_mem), 3) if all_mem else None,
        "memory_recall_events": len(all_mem),
        "continuity_break_rate": round(safe_mean(all_cont), 3) if all_cont else None,
        "total_errors": sum(d["errors"] for d in by_scn.values()),
    }
    (out_dir / "overall.json").write_text(json.dumps(overall, indent=2), encoding="utf-8")
    print(f"✔ Wrote {csv_path}")
    print(f"✔ Wrote {out_dir / 'overall.json'}")

    # write memory recall detail report (T3.6 Part 2)
    if memory_recall_events:
        mem_csv_path = out_dir / "memory_recall_summary.csv"
        with open(mem_csv_path, "w", newline="", encoding="utf-8") as f:
            w = csv.writer(f)
            w.writerow(["scenario_id", "run_id", "repeat", "turn", "expected",
                         "recall_hit", "qa_check_type", "user_input", "agent_response"])
            for evt in memory_recall_events:
                w.writerow([
                    evt["scenario_id"], evt["run_id"], evt["repeat"], evt["turn"],
                    evt["expected"], evt["recall_hit"], evt["qa_check_type"],
                    evt["user_input"], evt["agent_response"],
                ])
        print(f"✔ Wrote {mem_csv_path} ({len(memory_recall_events)} recall events)")

if __name__ == "__main__":
    main()