diff options
Diffstat (limited to 'scripts')
| -rw-r--r-- | scripts/analyze_pi_sessions.py | 458 |
1 files changed, 458 insertions, 0 deletions
diff --git a/scripts/analyze_pi_sessions.py b/scripts/analyze_pi_sessions.py new file mode 100644 index 0000000..ffa9f00 --- /dev/null +++ b/scripts/analyze_pi_sessions.py @@ -0,0 +1,458 @@ +#!/usr/bin/env python3 + +import argparse +import json +import math +from collections import Counter +from pathlib import Path +from statistics import mean, median +from typing import Any + + +BUCKETS = ("input", "output", "cache_read", "cache_write") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Analyze Pi agent session logs for token mix and cache behavior." + ) + parser.add_argument( + "root", + nargs="?", + default="~/.pi/agent/sessions", + help="Directory containing Pi session .jsonl files.", + ) + parser.add_argument( + "--json-out", + help="Optional path to write the full report as JSON.", + ) + return parser.parse_args() + + +def usage_value(usage: dict[str, Any], *keys: str) -> int: + for key in keys: + value = usage.get(key) + if value is None: + continue + if isinstance(value, bool): + return int(value) + if isinstance(value, (int, float)): + return int(value) + return 0 + + +def safe_div(numerator: float, denominator: float) -> float | None: + if denominator == 0: + return None + return numerator / denominator + + +def percentile(values: list[float], p: float) -> float | None: + if not values: + return None + if len(values) == 1: + return values[0] + ordered = sorted(values) + rank = (len(ordered) - 1) * p + low = math.floor(rank) + high = math.ceil(rank) + if low == high: + return ordered[low] + weight = rank - low + return ordered[low] * (1 - weight) + ordered[high] * weight + + +def summarize(values: list[float]) -> dict[str, float | None]: + return { + "count": len(values), + "mean": mean(values) if values else None, + "median": median(values) if values else None, + "p10": percentile(values, 0.10), + "p90": percentile(values, 0.90), + "min": min(values) if values else None, + "max": max(values) if values else None, + } + + +def session_primary(counter: Counter[str]) -> str | None: + if not counter: + return None + return max(counter.items(), key=lambda item: (item[1], item[0]))[0] + + +def format_number(value: float | None, digits: int = 3) -> str: + if value is None: + return "n/a" + return f"{value:,.{digits}f}" + + +def format_pct(value: float | None, digits: int = 1) -> str: + if value is None: + return "n/a" + return f"{value * 100:.{digits}f}%" + + +def analyze_session(path: Path) -> dict[str, Any]: + totals = {bucket: 0 for bucket in BUCKETS} + usage_messages = 0 + assistant_messages = 0 + models = Counter() + providers = Counter() + apis = Counter() + line_count = 0 + parse_errors = 0 + + with path.open("r", encoding="utf-8") as handle: + for raw_line in handle: + line_count += 1 + line = raw_line.strip() + if not line: + continue + try: + entry = json.loads(line) + except json.JSONDecodeError: + parse_errors += 1 + continue + + message = entry.get("message") + if not isinstance(message, dict): + continue + + role = message.get("role") + if role == "assistant": + assistant_messages += 1 + + usage = entry.get("usage") or message.get("usage") + if role != "assistant" or not isinstance(usage, dict): + continue + + usage_messages += 1 + bucket_values = { + "input": usage_value(usage, "input", "input_tokens"), + "output": usage_value(usage, "output", "output_tokens"), + "cache_read": usage_value( + usage, + "cacheRead", + "cache_read", + "cache_read_input_tokens", + ), + "cache_write": usage_value( + usage, + "cacheWrite", + "cache_write", + "cache_creation_input_tokens", + ), + } + for bucket, value in bucket_values.items(): + totals[bucket] += value + + billable = sum(bucket_values.values()) + if billable <= 0: + continue + + model = ( + entry.get("responseModel") + or entry.get("model") + or message.get("responseModel") + or message.get("model") + ) + provider = entry.get("provider") or message.get("provider") + api = entry.get("api") or message.get("api") + if isinstance(model, str): + models[model] += billable + if isinstance(provider, str): + providers[provider] += billable + if isinstance(api, str): + apis[api] += billable + + prompt_tokens = totals["input"] + totals["cache_read"] + totals["cache_write"] + billable_tokens = prompt_tokens + totals["output"] + cache_activity = totals["cache_read"] + totals["cache_write"] + + shares = { + bucket: safe_div(totals[bucket], billable_tokens) for bucket in BUCKETS + } + + return { + "path": str(path), + "lines": line_count, + "parse_errors": parse_errors, + "assistant_messages": assistant_messages, + "usage_messages": usage_messages, + "totals": totals, + "prompt_tokens": prompt_tokens, + "billable_tokens": billable_tokens, + "cache_activity_tokens": cache_activity, + "primary_model": session_primary(models), + "primary_provider": session_primary(providers), + "primary_api": session_primary(apis), + "metrics": { + "prompt_to_output_ratio": safe_div(prompt_tokens, totals["output"]), + "input_to_output_ratio": safe_div(totals["input"], totals["output"]), + "cache_hit_rate": safe_div(totals["cache_read"], cache_activity), + "cache_coverage": safe_div(cache_activity, prompt_tokens), + "prompt_reuse_rate": safe_div(totals["cache_read"], prompt_tokens), + "output_share": safe_div(totals["output"], billable_tokens), + "prompt_share": safe_div(prompt_tokens, billable_tokens), + **{f"{bucket}_share": value for bucket, value in shares.items()}, + }, + } + + +def build_report(root: Path) -> dict[str, Any]: + session_files = sorted(root.rglob("*.jsonl")) + sessions = [analyze_session(path) for path in session_files] + analyzed_sessions = [s for s in sessions if s["billable_tokens"] > 0] + + pooled_totals = {bucket: 0 for bucket in BUCKETS} + total_prompt = 0 + total_billable = 0 + total_usage_messages = 0 + total_assistant_messages = 0 + total_parse_errors = 0 + provider_sessions = Counter() + model_sessions = Counter() + + per_session_metric_values: dict[str, list[float]] = {} + cache_active_metric_values: dict[str, list[float]] = {} + cache_active_sessions = 0 + + for session in analyzed_sessions: + for bucket in BUCKETS: + pooled_totals[bucket] += session["totals"][bucket] + total_prompt += session["prompt_tokens"] + total_billable += session["billable_tokens"] + total_usage_messages += session["usage_messages"] + total_assistant_messages += session["assistant_messages"] + total_parse_errors += session["parse_errors"] + if session["primary_provider"]: + provider_sessions[session["primary_provider"]] += 1 + if session["primary_model"]: + model_sessions[session["primary_model"]] += 1 + for name, value in session["metrics"].items(): + if value is None: + continue + per_session_metric_values.setdefault(name, []).append(value) + if session["cache_activity_tokens"] > 0 and name in { + "cache_hit_rate", + "cache_coverage", + "prompt_reuse_rate", + "cache_read_share", + "cache_write_share", + }: + cache_active_metric_values.setdefault(name, []).append(value) + per_session_metric_values.setdefault("billable_tokens", []).append( + session["billable_tokens"] + ) + per_session_metric_values.setdefault("prompt_tokens", []).append( + session["prompt_tokens"] + ) + per_session_metric_values.setdefault("output_tokens", []).append( + session["totals"]["output"] + ) + if session["cache_activity_tokens"] > 0: + cache_active_sessions += 1 + + pooled_cache_activity = pooled_totals["cache_read"] + pooled_totals["cache_write"] + pooled_weights = { + bucket: safe_div(pooled_totals[bucket], total_billable) for bucket in BUCKETS + } + + mean_session_weights = { + bucket: mean(per_session_metric_values.get(f"{bucket}_share", [])) + if per_session_metric_values.get(f"{bucket}_share") + else None + for bucket in BUCKETS + } + median_session_weights = { + bucket: median(per_session_metric_values.get(f"{bucket}_share", [])) + if per_session_metric_values.get(f"{bucket}_share") + else None + for bucket in BUCKETS + } + + return { + "root": str(root), + "sessions_scanned": len(session_files), + "sessions_with_usage": len(analyzed_sessions), + "sessions_without_usage": len(session_files) - len(analyzed_sessions), + "assistant_messages": total_assistant_messages, + "assistant_messages_with_usage": total_usage_messages, + "parse_errors": total_parse_errors, + "cache_active_sessions": cache_active_sessions, + "pooled_totals": pooled_totals, + "pooled_metrics": { + "prompt_tokens": total_prompt, + "billable_tokens": total_billable, + "prompt_to_output_ratio": safe_div(total_prompt, pooled_totals["output"]), + "input_to_output_ratio": safe_div( + pooled_totals["input"], pooled_totals["output"] + ), + "cache_hit_rate": safe_div( + pooled_totals["cache_read"], pooled_cache_activity + ), + "cache_coverage": safe_div(pooled_cache_activity, total_prompt), + "prompt_reuse_rate": safe_div(pooled_totals["cache_read"], total_prompt), + "output_share": safe_div(pooled_totals["output"], total_billable), + "prompt_share": safe_div(total_prompt, total_billable), + **{f"{bucket}_share": value for bucket, value in pooled_weights.items()}, + }, + "recommended_weights": { + "pooled_token_mix": pooled_weights, + "mean_session_mix": mean_session_weights, + "median_session_mix": median_session_weights, + }, + "per_session_distributions": { + name: summarize(values) + for name, values in sorted(per_session_metric_values.items()) + }, + "cache_active_session_distributions": { + name: summarize(values) + for name, values in sorted(cache_active_metric_values.items()) + }, + "top_primary_providers": provider_sessions.most_common(10), + "top_primary_models": model_sessions.most_common(10), + } + + +def print_report(report: dict[str, Any]) -> None: + print("Pi Session Cost Analysis") + print("========================") + print(f"Root: {report['root']}") + print( + f"Sessions scanned: {report['sessions_scanned']} " + f"({report['sessions_with_usage']} with usage)" + ) + print( + f"Assistant messages with usage: {report['assistant_messages_with_usage']} " + f"of {report['assistant_messages']}" + ) + print(f"Cache-active sessions: {report['cache_active_sessions']}") + print() + + pooled = report["pooled_metrics"] + totals = report["pooled_totals"] + weights = report["recommended_weights"] + + print("Pooled Totals") + print("-------------") + print(f"Input: {totals['input']:,}") + print(f"Output: {totals['output']:,}") + print(f"Cache read: {totals['cache_read']:,}") + print(f"Cache write: {totals['cache_write']:,}") + print(f"Billable: {pooled['billable_tokens']:,}") + print() + + print("Recommended Weight Vectors") + print("--------------------------") + for label, vector in ( + ("Pooled token mix", weights["pooled_token_mix"]), + ("Mean session mix", weights["mean_session_mix"]), + ("Median session mix", weights["median_session_mix"]), + ): + print( + f"{label:18} " + f"input={format_pct(vector['input'])} " + f"output={format_pct(vector['output'])} " + f"cache_read={format_pct(vector['cache_read'])} " + f"cache_write={format_pct(vector['cache_write'])}" + ) + print() + + print("Key Corpus Metrics") + print("------------------") + print( + f"Prompt/output ratio: {format_number(pooled['prompt_to_output_ratio'])}x " + f"(prompt = input + cache_read + cache_write)" + ) + print( + f"Uncached input/output ratio: {format_number(pooled['input_to_output_ratio'])}x" + ) + print(f"Cache hit rate: {format_pct(pooled['cache_hit_rate'])}") + print(f"Cache coverage of prompt side: {format_pct(pooled['cache_coverage'])}") + print(f"Prompt reuse rate: {format_pct(pooled['prompt_reuse_rate'])}") + print() + + print("Per-Session Distributions") + print("-------------------------") + metric_names = [ + "billable_tokens", + "prompt_to_output_ratio", + "input_to_output_ratio", + "cache_hit_rate", + "cache_coverage", + "prompt_reuse_rate", + "input_share", + "output_share", + "cache_read_share", + "cache_write_share", + ] + for name in metric_names: + stats = report["per_session_distributions"].get(name) + if not stats: + continue + unit = "%" if ( + name.endswith("_share") + or "rate" in name + or "coverage" in name + ) else "" + value_fmt = format_pct if unit == "%" else format_number + print( + f"{name:22} " + f"mean={value_fmt(stats['mean'])} " + f"median={value_fmt(stats['median'])} " + f"p10={value_fmt(stats['p10'])} " + f"p90={value_fmt(stats['p90'])}" + ) + print() + + if report["cache_active_session_distributions"]: + print("Cache-Active Session Distributions") + print("----------------------------------") + for name in [ + "cache_hit_rate", + "cache_coverage", + "prompt_reuse_rate", + "cache_read_share", + "cache_write_share", + ]: + stats = report["cache_active_session_distributions"].get(name) + if not stats: + continue + print( + f"{name:22} " + f"mean={format_pct(stats['mean'])} " + f"median={format_pct(stats['median'])} " + f"p10={format_pct(stats['p10'])} " + f"p90={format_pct(stats['p90'])}" + ) + print() + + print("Top Primary Providers") + print("---------------------") + for provider, count in report["top_primary_providers"]: + print(f"{provider:20} {count}") + print() + + print("Top Primary Models") + print("------------------") + for model, count in report["top_primary_models"]: + print(f"{model:35} {count}") + + +def main() -> int: + args = parse_args() + root = Path(args.root).expanduser().resolve() + report = build_report(root) + + if args.json_out: + out_path = Path(args.json_out).expanduser().resolve() + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(json.dumps(report, indent=2), encoding="utf-8") + + print_report(report) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) |