aggregate.py

"""Aggregate a benchmark JSONL into a (task, transport) summary table.

Reads one or more raw JSONL files, groups by (task, transport), computes
median + IQR for the numeric metrics, and success rate. Prints a table
to stdout and optionally writes a machine-readable JSON summary.

Usage:
    python -m src.aggregate results/raw/20260408T*.jsonl
    python -m src.aggregate --out results/scrubbed/summary.json results/raw/latest.jsonl

Metrics summarised (per (task, transport)):
    - effective_input_tokens  (input + cache_read + cache_creation)
    - output_tokens
    - tool_call_count
    - wall_clock_ms
    - success_rate
    - n
"""

from __future__ import annotations

import argparse
import json
import statistics
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any

NUMERIC_METRICS = (
    "effective_input_tokens",
    "output_tokens",
    "tool_call_count",
    "wall_clock_ms",
)


def load_jsonl(paths: list[Path]) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for p in paths:
        with p.open() as f:
            for line in f:
                line = line.strip()
                if line:
                    records.append(json.loads(line))
    return records


def _pct(values: list[float], q: float) -> float:
    if not values:
        return 0.0
    s = sorted(values)
    if len(s) == 1:
        return s[0]
    k = (len(s) - 1) * q
    lo = int(k)
    hi = min(lo + 1, len(s) - 1)
    return s[lo] + (s[hi] - s[lo]) * (k - lo)


def summarise(records: list[dict[str, Any]]) -> dict[tuple[str, str], dict[str, Any]]:
    grouped: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for r in records:
        grouped[(r["task_id"], r["transport"])].append(r)

    out: dict[tuple[str, str], dict[str, Any]] = {}
    for key, rs in grouped.items():
        row: dict[str, Any] = {"n": len(rs)}
        row["success_rate"] = sum(
            1 for r in rs if r.get("verdict", {}).get("ok")
        ) / len(rs)

        # Collect numeric series. effective_input_tokens is computed from
        # the usage dict; the others can be read directly.
        series: dict[str, list[float]] = {m: [] for m in NUMERIC_METRICS}
        for r in rs:
            u = r.get("usage", {}) or {}
            series["effective_input_tokens"].append(
                float(u.get("effective_input_tokens", 0))
            )
            series["output_tokens"].append(float(u.get("output_tokens", 0)))
            series["tool_call_count"].append(float(r.get("tool_call_count", 0)))
            series["wall_clock_ms"].append(float(r.get("wall_clock_ms", 0)))

        for m, values in series.items():
            row[m] = {
                "median": statistics.median(values),
                "p25": _pct(values, 0.25),
                "p75": _pct(values, 0.75),
                "min": min(values),
                "max": max(values),
                "mean": statistics.mean(values),
            }
        out[key] = row
    return out


def print_table(summary: dict[tuple[str, str], dict[str, Any]]) -> None:
    tasks = sorted({t for t, _ in summary})
    transports_seen: list[str] = []
    for _, t in summary:
        if t not in transports_seen:
            transports_seen.append(t)
    transports = transports_seen or sorted({t for _, t in summary})

    # Header
    col = "{task:<22} {transport:<9} {n:>3} {ok:>5} {in_med:>7} {in_iqr:>13} {out_med:>5} {calls:>5} {wall:>7}"
    print(
        col.format(
            task="task",
            transport="transport",
            n="n",
            ok="ok%",
            in_med="in(med)",
            in_iqr="in(p25-p75)",
            out_med="out",
            calls="calls",
            wall="wall_ms",
        )
    )
    print("-" * 80)
    for task in tasks:
        for transport in transports:
            row = summary.get((task, transport))
            if not row:
                continue
            i = row["effective_input_tokens"]
            o = row["output_tokens"]
            c = row["tool_call_count"]
            w = row["wall_clock_ms"]
            print(
                col.format(
                    task=task,
                    transport=transport,
                    n=row["n"],
                    ok=f"{int(row['success_rate'] * 100):>4}%",
                    in_med=f"{int(i['median']):>7}",
                    in_iqr=f"{int(i['p25'])}-{int(i['p75'])}",
                    out_med=f"{int(o['median']):>5}",
                    calls=f"{c['median']:.1f}",
                    wall=f"{int(w['median']):>7}",
                )
            )
        print()


def to_json(
    summary: dict[tuple[str, str], dict[str, Any]],
) -> list[dict[str, Any]]:
    return [
        {"task": task, "transport": transport, **row}
        for (task, transport), row in sorted(summary.items())
    ]


def main() -> int:
    parser = argparse.ArgumentParser(description="Aggregate raw benchmark JSONL")
    parser.add_argument("inputs", nargs="+", help="raw JSONL file(s)")
    parser.add_argument("--out", default="", help="optional JSON summary path")
    args = parser.parse_args()

    paths = [Path(p) for p in args.inputs]
    records = load_jsonl(paths)
    if not records:
        print("no records found", file=sys.stderr)
        return 1
    summary = summarise(records)
    print_table(summary)

    if args.out:
        out_path = Path(args.out)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(json.dumps(to_json(summary), indent=2))
        print(f"\nwrote {out_path}")
    return 0


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