diff options
Diffstat (limited to 'scripts')
| -rw-r--r-- | scripts/analyze_pi_sessions.py | 458 |
1 files changed, 0 insertions, 458 deletions
diff --git a/scripts/analyze_pi_sessions.py b/scripts/analyze_pi_sessions.py deleted file mode 100644 index ffa9f00..0000000 --- a/scripts/analyze_pi_sessions.py +++ /dev/null @@ -1,458 +0,0 @@ -#!/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()) |