PMFlow/scripts/agent_capability_replay.py

#!/usr/bin/env python3
"""Offline replay scorer for dual-agent capability metrics."""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any


@dataclass(slots=True)
class MetricAccumulator:
    count: int = 0
    intent_sum: float = 0.0
    gate_sum: float = 0.0
    exec_sum: float = 0.0
    correction_sum: float = 0.0

    def append(self, *, intent: float, gate: float, exec_success: float, correction: float) -> None:
        self.count += 1
        self.intent_sum += intent
        self.gate_sum += gate
        self.exec_sum += exec_success
        self.correction_sum += correction

    def summary(self) -> dict[str, Any]:
        if self.count <= 0:
            return {
                "count": 0,
                "intent_accuracy": 0.0,
                "gate_precision": 0.0,
                "exec_success": 0.0,
                "correction_success": 0.0,
                "score": 0.0,
            }
        intent_accuracy = self.intent_sum / self.count
        gate_precision = self.gate_sum / self.count
        exec_success = self.exec_sum / self.count
        correction_success = self.correction_sum / self.count
        return {
            "count": self.count,
            "intent_accuracy": round(intent_accuracy, 4),
            "gate_precision": round(gate_precision, 4),
            "exec_success": round(exec_success, 4),
            "correction_success": round(correction_success, 4),
            "score": _compute_capability_score(
                intent_accuracy=intent_accuracy,
                gate_precision=gate_precision,
                exec_success=exec_success,
                correction_success=correction_success,
            ),
        }


def _compute_capability_score(
    *,
    intent_accuracy: float,
    gate_precision: float,
    exec_success: float,
    correction_success: float,
) -> float:
    score = (
        0.30 * max(0.0, min(1.0, float(intent_accuracy or 0.0)))
        + 0.30 * max(0.0, min(1.0, float(gate_precision or 0.0)))
        + 0.25 * max(0.0, min(1.0, float(exec_success or 0.0)))
        + 0.15 * max(0.0, min(1.0, float(correction_success or 0.0)))
    )
    return round(score, 4)


def _to_binary(value: Any) -> float:
    if isinstance(value, bool):
        return 1.0 if value else 0.0
    if isinstance(value, (int, float)):
        return 1.0 if float(value) >= 1.0 else (1.0 if float(value) >= 0.5 else 0.0)
    text = str(value or "").strip().lower()
    return 1.0 if text in {"1", "true", "yes", "y", "ok", "pass", "passed", "success", "succeeded"} else 0.0


def _load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line:
            continue
        parsed = json.loads(line)
        if isinstance(parsed, dict):
            rows.append(parsed)
    return rows


def _bucket_of(row: dict[str, Any]) -> str:
    bucket = str(row.get("bucket") or "").strip().lower()
    if bucket in {"simple", "complex"}:
        return bucket
    trace = row.get("decision_trace") if isinstance(row.get("decision_trace"), dict) else {}
    route_tier = str(trace.get("route_tier") or row.get("route_tier") or "").strip().lower()
    return "complex" if route_tier == "high" else "simple"


def _metric(row: dict[str, Any], key: str, fallback_keys: list[str]) -> float:
    if key in row:
        return _to_binary(row.get(key))
    for name in fallback_keys:
        if name in row:
            return _to_binary(row.get(name))
    return 0.0


def main() -> None:
    parser = argparse.ArgumentParser(description="Replay dual-agent metrics from JSONL samples.")
    parser.add_argument("--input", required=True, help="Path to JSONL file.")
    parser.add_argument("--output", default="", help="Optional output path; default prints to stdout.")
    args = parser.parse_args()

    rows = _load_jsonl(Path(args.input))
    overall = MetricAccumulator()
    simple = MetricAccumulator()
    complex_bucket = MetricAccumulator()

    for row in rows:
        intent = _metric(row, "intent_correct", ["intent_accuracy", "intent_success"])
        gate = _metric(row, "gate_correct", ["gate_precision", "gate_success"])
        exec_success = _metric(row, "exec_success", ["execution_success", "succeeded"])
        correction = _metric(row, "correction_success", ["correction_correct"])
        overall.append(intent=intent, gate=gate, exec_success=exec_success, correction=correction)
        bucket_name = _bucket_of(row)
        if bucket_name == "complex":
            complex_bucket.append(intent=intent, gate=gate, exec_success=exec_success, correction=correction)
        else:
            simple.append(intent=intent, gate=gate, exec_success=exec_success, correction=correction)

    payload = {
        "overall": overall.summary(),
        "simple": simple.summary(),
        "complex": complex_bucket.summary(),
        "thresholds": {
            "overall_score_target": 0.80,
            "bucket_score_target": 0.78,
            "intent_target": 0.80,
            "gate_target": 0.78,
            "correction_target": 0.75,
        },
    }
    text = json.dumps(payload, ensure_ascii=False, indent=2)
    if args.output:
        Path(args.output).write_text(text + "\n", encoding="utf-8")
        return
    print(text)


if __name__ == "__main__":
    main()