| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| from pathlib import Path |
|
|
|
|
| TASK_ORDER = [ |
| "alzheimer-mouse", |
| "comparative-genomics", |
| "cystic-fibrosis", |
| "deseq", |
| "evolution", |
| "giab", |
| "metagenomics", |
| "single-cell", |
| "transcript-quant", |
| "viral-metagenomics", |
| ] |
|
|
|
|
| def load_json(path: Path, default=None): |
| if not path.exists(): |
| return default |
| return json.loads(path.read_text(encoding="utf-8")) |
|
|
|
|
| def as_int(value) -> int: |
| if isinstance(value, bool): |
| return int(value) |
| if isinstance(value, int): |
| return value |
| if isinstance(value, float): |
| return int(value) |
| if isinstance(value, str) and value.strip(): |
| try: |
| return int(float(value)) |
| except ValueError: |
| return 0 |
| return 0 |
|
|
|
|
| def latest_run_dir(runs_root: Path, task_id: str) -> Path | None: |
| pattern = re.compile(rf"^{re.escape(task_id)}_\d{{8}}_\d{{6}}$") |
| candidates = sorted(run_dir for run_dir in runs_root.iterdir() if run_dir.is_dir() and pattern.match(run_dir.name)) |
| return candidates[-1] if candidates else None |
|
|
|
|
| def summarize_run(task_id: str, run_dir: Path | None) -> dict: |
| row = { |
| "task": task_id, |
| "run_dir": str(run_dir) if run_dir else None, |
| "cases": 1 if run_dir else 0, |
| "cases_with_usage": 0, |
| "total_prompt_tokens": 0, |
| "total_completion_tokens": 0, |
| "total_tokens": 0, |
| "total_planning_context_tokens": 0, |
| "avg_prompt_tokens_per_case": 0.0, |
| "avg_completion_tokens_per_case": 0.0, |
| "avg_total_tokens_per_case": 0.0, |
| "avg_planning_context_tokens_per_case": 0.0, |
| "avg_llm_calls_per_case": 0.0, |
| "planning_latency_seconds": None, |
| "total_runtime_seconds": None, |
| } |
| if run_dir is None: |
| return row |
|
|
| summary = load_json(run_dir / "run_summary.json", default={}) or {} |
| retrieval = load_json(run_dir / "retrieval_plan.json", default={}) or {} |
| token_usage = retrieval.get("token_usage") |
| if not isinstance(token_usage, dict): |
| token_usage = summary |
|
|
| prompt_tokens = as_int(token_usage.get("prompt_tokens")) |
| completion_tokens = as_int(token_usage.get("completion_tokens")) |
| total_tokens = as_int(token_usage.get("total_tokens")) |
| llm_call_count = as_int(token_usage.get("llm_call_count")) |
| planning_context_tokens = as_int(retrieval.get("planning_context_tokens", summary.get("planning_context_tokens"))) |
|
|
| row.update( |
| { |
| "cases_with_usage": 1 if any(x > 0 for x in (prompt_tokens, completion_tokens, total_tokens)) else 0, |
| "total_prompt_tokens": prompt_tokens, |
| "total_completion_tokens": completion_tokens, |
| "total_tokens": total_tokens, |
| "total_planning_context_tokens": planning_context_tokens, |
| "avg_prompt_tokens_per_case": round(float(prompt_tokens), 2), |
| "avg_completion_tokens_per_case": round(float(completion_tokens), 2), |
| "avg_total_tokens_per_case": round(float(total_tokens), 2), |
| "avg_planning_context_tokens_per_case": round(float(planning_context_tokens), 2), |
| "avg_llm_calls_per_case": round(float(llm_call_count), 2), |
| "planning_latency_seconds": retrieval.get("planning_latency_seconds", summary.get("planning_latency_seconds")), |
| "total_runtime_seconds": retrieval.get("total_runtime_seconds", summary.get("total_runtime_seconds")), |
| } |
| ) |
| return row |
|
|
|
|
| def summarize_overall(rows: list[dict]) -> dict: |
| present = [row for row in rows if row.get("cases")] |
| n = len(present) |
| if not n: |
| return { |
| "task": "__overall__", |
| "cases": 0, |
| "cases_with_usage": 0, |
| "total_prompt_tokens": 0, |
| "total_completion_tokens": 0, |
| "total_tokens": 0, |
| "total_planning_context_tokens": 0, |
| "avg_prompt_tokens_per_case": 0.0, |
| "avg_completion_tokens_per_case": 0.0, |
| "avg_total_tokens_per_case": 0.0, |
| "avg_planning_context_tokens_per_case": 0.0, |
| "avg_llm_calls_per_case": 0.0, |
| } |
| total_prompt = sum(as_int(row.get("total_prompt_tokens")) for row in present) |
| total_completion = sum(as_int(row.get("total_completion_tokens")) for row in present) |
| total_tokens = sum(as_int(row.get("total_tokens")) for row in present) |
| total_context = sum(as_int(row.get("total_planning_context_tokens")) for row in present) |
| total_calls = sum(as_int(row.get("avg_llm_calls_per_case")) for row in present) |
| return { |
| "task": "__overall__", |
| "cases": n, |
| "cases_with_usage": sum(as_int(row.get("cases_with_usage")) for row in present), |
| "total_prompt_tokens": total_prompt, |
| "total_completion_tokens": total_completion, |
| "total_tokens": total_tokens, |
| "total_planning_context_tokens": total_context, |
| "avg_prompt_tokens_per_case": round(total_prompt / n, 2), |
| "avg_completion_tokens_per_case": round(total_completion / n, 2), |
| "avg_total_tokens_per_case": round(total_tokens / n, 2), |
| "avg_planning_context_tokens_per_case": round(total_context / n, 2), |
| "avg_llm_calls_per_case": round(total_calls / n, 2), |
| } |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description="Summarize BioAgentBench real token usage and planning-context tokens.") |
| parser.add_argument("--results-dir", type=Path, required=True) |
| parser.add_argument("--out-json", type=Path, default=None) |
| args = parser.parse_args() |
|
|
| rows = [summarize_run(task_id, latest_run_dir(args.results_dir, task_id)) for task_id in TASK_ORDER] |
| rows.append(summarize_overall(rows)) |
| text = json.dumps(rows, indent=2, ensure_ascii=False) |
| print(text) |
| if args.out_json: |
| args.out_json.parent.mkdir(parents=True, exist_ok=True) |
| args.out_json.write_text(text + "\n", encoding="utf-8") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
