algonotes_rag/scripts/perf_report.py

# scripts/perf_report.py
# Performance report: parse logs/perf.log, aggregate by call_type,
# and present latency statistics.

import json
import statistics
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from rich.console import Console
from rich.table import Table

console = Console()

DEFAULT_PERF_LOG = Path("logs/perf.log")


@dataclass
class PerfEntry:
    """A single parsed performance log entry."""

    timestamp: str
    call_type: str
    latency_ms: float
    success: bool
    model: str | None = None
    doc_count: int | None = None
    error: str | None = None
    extra: dict[str, Any] = field(default_factory=dict)


def load_entries(path: Path | None = None) -> list[PerfEntry]:
    """Parse ``logs/perf.log`` into a list of :class:`PerfEntry`.

    Args:
        path: Path to perf log file. Defaults to ``logs/perf.log``.

    Returns:
        List of parsed entries (empty if file missing or unparseable).
    """
    path = path or DEFAULT_PERF_LOG
    if not path.exists():
        return []

    entries: list[PerfEntry] = []
    with path.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue

            entries.append(PerfEntry(
                timestamp=obj.get("timestamp", ""),
                call_type=obj.get("call_type", "unknown"),
                latency_ms=float(obj.get("latency_ms", 0)),
                success=bool(obj.get("success", True)),
                model=obj.get("model"),
                doc_count=obj.get("doc_count"),
                error=obj.get("error"),
                extra={k: v for k, v in obj.items()
                       if k not in ("timestamp", "level", "logger", "message",
                                    "call_type", "latency_ms", "success",
                                    "model", "doc_count", "error")},
            ))
    return entries


def aggregate(entries: list[PerfEntry]) -> dict[str, dict[str, Any]]:
    """Aggregate latency statistics grouped by call_type.

    Args:
        entries: Parsed perf log entries.

    Returns:
        Dict keyed by call_type, each value containing:
        count, total_ms, avg_ms, min_ms, max_ms, p50_ms, p95_ms, p99_ms,
        failures, success_rate.
    """
    groups: dict[str, list[float]] = {}
    failures: dict[str, int] = {}
    for e in entries:
        groups.setdefault(e.call_type, []).append(e.latency_ms)
        if not e.success:
            failures[e.call_type] = failures.get(e.call_type, 0) + 1

    result: dict[str, dict[str, Any]] = {}
    for ct, lats in sorted(groups.items()):
        total = len(lats)
        fails = failures.get(ct, 0)
        sorted_lats = sorted(lats)
        result[ct] = {
            "count": total,
            "total_ms": round(sum(lats), 1),
            "avg_ms": round(statistics.mean(lats), 1),
            "min_ms": round(min(lats), 1),
            "max_ms": round(max(lats), 1),
            "p50_ms": round(_percentile(sorted_lats, 50), 1),
            "p95_ms": round(_percentile(sorted_lats, 95), 1),
            "p99_ms": round(_percentile(sorted_lats, 99), 1),
            "failures": fails,
            "success_rate": f"{(total - fails) / total * 100:.1f}%",
        }
    return result


def _percentile(sorted_values: list[float], pct: float) -> float:
    """Compute the pct-th percentile of sorted values (linear interpolation)."""
    if not sorted_values:
        return 0.0
    n = len(sorted_values)
    k = (pct / 100) * (n - 1)
    f = int(k)
    c = k - f
    if f + 1 < n:
        return sorted_values[f] + c * (sorted_values[f + 1] - sorted_values[f])
    return sorted_values[f]


def print_report(
    entries: list[PerfEntry],
    detail: bool = False,
) -> None:
    """Print a latency statistics report to the console.

    Args:
        entries: Parsed perf log entries.
        detail: If True, also print a per-entry table.
    """
    if not entries:
        console.print("[yellow]No performance data found.[/yellow]")
        console.print(
            f"[dim]Run some queries first — logs are stored in "
            f"{DEFAULT_PERF_LOG}[/dim]"
        )
        return

    stats = aggregate(entries)

    # ── Summary table ───────────────────────────────────
    table = Table(title="📈 性能统计 (Perf Log)")
    table.add_column("call_type", style="cyan")
    table.add_column("次数", justify="right")
    table.add_column("avg", justify="right")
    table.add_column("P50", justify="right")
    table.add_column("P95", justify="right")
    table.add_column("P99", justify="right")
    table.add_column("min", justify="right")
    table.add_column("max", justify="right")
    table.add_column("成功率", justify="right")

    for ct, s in stats.items():
        table.add_row(
            ct,
            str(s["count"]),
            f"{s['avg_ms']:.0f}ms",
            f"{s['p50_ms']:.0f}ms",
            f"{s['p95_ms']:.0f}ms",
            f"{s['p99_ms']:.0f}ms",
            f"{s['min_ms']:.0f}ms",
            f"{s['max_ms']:.0f}ms",
            s["success_rate"],
        )

    console.print(table)

    # ── Detail table ────────────────────────────────────
    if detail:
        console.print()
        dtable = Table(title="📋 详细记录")
        dtable.add_column("timestamp", style="dim")
        dtable.add_column("call_type", style="cyan")
        dtable.add_column("latency", justify="right")
        dtable.add_column("status", justify="center")
        dtable.add_column("details")

        for e in entries:
            status = "✅" if e.success else "❌"
            extras = []
            if e.model:
                extras.append(f"model={e.model}")
            if e.doc_count is not None:
                extras.append(f"docs={e.doc_count}")
            if e.error:
                extras.append(f"error={e.error}")
            for k, v in e.extra.items():
                if isinstance(v, (int, float, str)):
                    extras.append(f"{k}={v}")
            dtable.add_row(
                e.timestamp,
                e.call_type,
                f"{e.latency_ms:.0f}ms",
                status,
                ", ".join(extras) if extras else "—",
            )

        console.print(dtable)


def report_to_dict(entries: list[PerfEntry]) -> dict[str, Any]:
    """Return aggregated stats + raw entries as a JSON-serializable dict."""
    return {
        "summary": aggregate(entries),
        "total_entries": len(entries),
        "entries": [
            {
                "timestamp": e.timestamp,
                "call_type": e.call_type,
                "latency_ms": e.latency_ms,
                "success": e.success,
                "model": e.model,
                "doc_count": e.doc_count,
                "error": e.error,
                **e.extra,
            }
            for e in entries
        ],
    }