rein-aharness/rein_aharness/metrics.py

"""Write per-run records into the target repo's `.kaizen/metrics` tree.

Follows kaizen-agentic ADR-004 conventions so the optimization loop can
observe harness-run agents:

  .kaizen/metrics/<agent>/
    executions.jsonl   # append-only
    summary.json       # regenerated on write
"""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


@dataclass
class ExecutionRecord:
    timestamp: str
    agent: str
    success: bool
    execution_time_s: float = 0.0
    session_id: str | None = None
    quality_score: float | None = None
    primary_metric: dict[str, Any] | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    # Helix / harness correlation (ADR-004 optional fields)
    repo: str | None = None
    tokens: int | None = None
    committed: bool | None = None
    head_after: str | None = None
    reason: str | None = None
    harness: str = "rein-aharness"

    def to_json_line(self) -> str:
        data = asdict(self)
        # Drop Nones for a compact record; required fields always present.
        compact = {k: v for k, v in data.items() if v is not None}
        return json.dumps(compact, sort_keys=True)


def metrics_dir(project_root: Path, agent: str) -> Path:
    return Path(project_root) / ".kaizen" / "metrics" / agent


def _utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace(
        "+00:00", "Z"
    )


def _load_executions(path: Path) -> list[dict[str, Any]]:
    if not path.is_file():
        return []
    records: list[dict[str, Any]] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            records.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    return records


def _trend(values: list[float]) -> str:
    if len(values) < 4:
        return "stable"
    mid = len(values) // 2
    early = sum(values[:mid]) / max(1, mid)
    late = sum(values[mid:]) / max(1, len(values) - mid)
    if late - early > 0.05:
        return "up"
    if early - late > 0.05:
        return "down"
    return "stable"


def regenerate_summary(agent: str, executions: list[dict[str, Any]]) -> dict[str, Any]:
    count = len(executions)
    successes = [e for e in executions if e.get("success")]
    success_rate = (len(successes) / count) if count else 0.0
    times = [float(e["execution_time_s"]) for e in executions if "execution_time_s" in e]
    qualities = [
        float(e["quality_score"])
        for e in executions
        if isinstance(e.get("quality_score"), (int, float))
    ]
    last_ts = executions[-1].get("timestamp") if executions else None
    return {
        "agent": agent,
        "execution_count": count,
        "success_rate": round(success_rate, 3),
        "avg_quality_score": (
            round(sum(qualities) / len(qualities), 3) if qualities else None
        ),
        "avg_execution_time_s": (
            round(sum(times) / len(times), 3) if times else None
        ),
        "last_execution": last_ts,
        "trend": {
            "success_rate": _trend(
                [1.0 if e.get("success") else 0.0 for e in executions]
            ),
            "quality_score": _trend(qualities) if qualities else "stable",
        },
    }


def record_execution(
    project_root: Path | str,
    agent: str,
    *,
    success: bool,
    execution_time_s: float = 0.0,
    tokens: int | None = None,
    committed: bool | None = None,
    head_after: str | None = None,
    reason: str | None = None,
    metadata: dict[str, Any] | None = None,
    session_id: str | None = None,
) -> Path:
    """Append one execution record and regenerate summary.json.

    Returns the path to executions.jsonl. Never raises for missing parent
    dirs (creates them). Callers that must not write should skip this.
    """
    root = Path(project_root)
    directory = metrics_dir(root, agent)
    directory.mkdir(parents=True, exist_ok=True)

    record = ExecutionRecord(
        timestamp=_utc_now(),
        agent=agent,
        success=success,
        execution_time_s=float(execution_time_s),
        session_id=session_id,
        metadata=metadata or {},
        repo=root.name,
        tokens=tokens,
        committed=committed,
        head_after=head_after,
        reason=reason,
    )

    executions_path = directory / "executions.jsonl"
    with executions_path.open("a", encoding="utf-8") as fh:
        fh.write(record.to_json_line() + "\n")

    all_records = _load_executions(executions_path)
    summary = regenerate_summary(agent, all_records)
    (directory / "summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return executions_path