core/kpi.py

from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, Any, List


@dataclass
class KPIReport:
    steps: int
    simulated_seconds: float
    wall_seconds: float
    total_cost_usd: float
    total_tokens_in: int
    total_tokens_out: int
    total_actions: int
    per_agent: Dict[str, Dict[str, Any]]


def compute_kpis(run_log: List[Dict[str, Any]]) -> KPIReport:
    if not run_log:
        return KPIReport(
            steps=0,
            simulated_seconds=0.0,
            wall_seconds=0.0,
            total_cost_usd=0.0,
            total_tokens_in=0,
            total_tokens_out=0,
            total_actions=0,
            per_agent={},
        )

    steps = int(max(e.get("t", 0) for e in run_log))
    simulated_seconds = float(sum(e.get("sim_seconds", 0.0) for e in run_log))
    wall_seconds = float(sum(e.get("wall_seconds", 0.0) for e in run_log))
    total_cost_usd = float(sum(e.get("cost_usd", 0.0) for e in run_log))
    total_tokens_in = int(sum(e.get("tokens_in", 0) for e in run_log))
    total_tokens_out = int(sum(e.get("tokens_out", 0) for e in run_log))
    total_actions = int(sum(1 for e in run_log if e.get("action")))

    per_agent: Dict[str, Dict[str, Any]] = {}
    for e in run_log:
        agent = e.get("agent")
        if not agent:
            continue
        d = per_agent.setdefault(
            agent,
            {"actions": 0, "cost_usd": 0.0, "tokens_in": 0, "tokens_out": 0, "wall_seconds": 0.0, "sim_seconds": 0.0},
        )
        if e.get("action"):
            d["actions"] += 1
        d["cost_usd"] += float(e.get("cost_usd", 0.0))
        d["tokens_in"] += int(e.get("tokens_in", 0))
        d["tokens_out"] += int(e.get("tokens_out", 0))
        d["wall_seconds"] += float(e.get("wall_seconds", 0.0))
        d["sim_seconds"] += float(e.get("sim_seconds", 0.0))

    return KPIReport(
        steps=steps,
        simulated_seconds=simulated_seconds,
        wall_seconds=wall_seconds,
        total_cost_usd=total_cost_usd,
        total_tokens_in=total_tokens_in,
        total_tokens_out=total_tokens_out,
        total_actions=total_actions,
        per_agent=per_agent,
    )