#!/usr/bin/env python3 """Summarize Biomni task-level scaling metrics into a table-friendly JSON/CSV.""" from __future__ import annotations import argparse import csv import json import re from pathlib import Path from typing import Any TASK_ORDER = [ "alzheimer-mouse", "comparative-genomics", "cystic-fibrosis", "deseq", "evolution", "giab", "metagenomics", "single-cell", "transcript-quant", "viral-metagenomics", ] def load_json(path: Path, default: Any = None) -> Any: if not path.exists(): return default return json.loads(path.read_text(encoding="utf-8")) def load_yaml_if_available(path: Path, default: Any = None) -> Any: if not path.exists(): return default try: import yaml except ModuleNotFoundError: return default return yaml.safe_load(path.read_text(encoding="utf-8")) or default def rough_token_count(text: str) -> int: return max(1, len(re.findall(r"\S+", text))) if text.strip() else 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 selected_tool_names(retrieval_plan: dict[str, Any]) -> list[str]: names = retrieval_plan.get("selected_resource_names", {}).get("tools", []) cleaned = [str(name) for name in names if name] if cleaned: return cleaned tool_items = retrieval_plan.get("selected_resources", {}).get("tools", []) return [item.get("name") for item in tool_items if isinstance(item, dict) and item.get("name")] def infer_tool_universe_size(retrieval_plan: dict[str, Any], mcp_config: Path | None) -> int | None: registered_tool_count = retrieval_plan.get("registered_tool_count") if isinstance(registered_tool_count, int) and registered_tool_count > 0: return registered_tool_count registered_tool_names = retrieval_plan.get("registered_tool_names", []) if isinstance(registered_tool_names, list) and registered_tool_names: return len([name for name in registered_tool_names if name]) if mcp_config is None: return None cfg = load_yaml_if_available(mcp_config, default={}) if not isinstance(cfg, dict): return None total = 0 for server_meta in (cfg.get("mcp_servers") or {}).values(): if not isinstance(server_meta, dict): continue total += len([tool for tool in server_meta.get("tools", []) if isinstance(tool, dict) and tool.get("name")]) return total or None def build_task_row( task_id: str, run_dir: Path | None, evaluation_index: dict[str, dict[str, Any]], gold_index: dict[str, dict[str, Any]], mcp_config: Path | None, scale_label: str, ) -> dict[str, Any]: try: scale_tool_count = int(scale_label) except (TypeError, ValueError): scale_tool_count = None base = { "scale": scale_label, "task": task_id, "run_dir": str(run_dir) if run_dir else None, "results_match": None, "selected_tools": None, "overhead_planning_ratio": None, "gold_items": len((gold_index.get(task_id, {}) or {}).get("gold_tools", [])) + len((gold_index.get(task_id, {}) or {}).get("gold_servers", [])), "context_tokens": None, "planning_latency_seconds": None, "selection_rate": None, "registered_tool_count": None, "completion_rate": None, "final_result_reached": None, } evaluation = evaluation_index.get(task_id) if evaluation: results = evaluation.get("evaluation_results", {}) base["results_match"] = results.get("results_match") base["completion_rate"] = results.get("completion_rate") base["final_result_reached"] = results.get("final_result_reached") if run_dir is None: return base retrieval_plan = load_json(run_dir / "retrieval_plan.json", {}) run_summary = load_json(run_dir / "run_summary.json", {}) selected_tools = len(selected_tool_names(retrieval_plan)) total_runtime = retrieval_plan.get("total_runtime_seconds") or run_summary.get("total_runtime_seconds") planning_latency = retrieval_plan.get("planning_latency_seconds") or run_summary.get("planning_latency_seconds") context_text = retrieval_plan.get("planning_context_text") context_tokens = rough_token_count(context_text) if isinstance(context_text, str) else 0 registered_tool_count = infer_tool_universe_size(retrieval_plan, mcp_config) base.update( { "selected_tools": selected_tools, "overhead_planning_ratio": (planning_latency / total_runtime) if isinstance(planning_latency, (int, float)) and isinstance(total_runtime, (int, float)) and total_runtime else None, "context_tokens": context_tokens, "planning_latency_seconds": planning_latency, "selection_rate": (selected_tools / scale_tool_count) if scale_tool_count and selected_tools is not None else (selected_tools / registered_tool_count) if registered_tool_count and selected_tools is not None else None, "registered_tool_count": registered_tool_count, } ) return base def csv_cell(value: Any) -> Any: if isinstance(value, bool): return "TRUE" if value else "FALSE" return value def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Summarize Biomni scaling task metrics into JSON/CSV tables.") parser.add_argument("--runs-root", type=Path, required=True) parser.add_argument("--evaluation-json", type=Path, required=True) parser.add_argument("--gold", type=Path, required=True) parser.add_argument("--mcp-config", type=Path, default=None) parser.add_argument("--scale-label", default=None) parser.add_argument("--out-json", type=Path, required=True) parser.add_argument("--out-csv", type=Path, required=True) return parser.parse_args() def main() -> int: args = parse_args() evaluation_payload = load_json(args.evaluation_json, default={}) evaluation_index = { item.get("task_id"): item for item in evaluation_payload.get("results", []) if isinstance(item, dict) } gold_index = load_json(args.gold, default={}) or {} scale_label = args.scale_label or args.runs_root.name.replace("scale_", "") rows = [] for task_id in TASK_ORDER: rows.append( build_task_row( task_id=task_id, run_dir=latest_run_dir(args.runs_root, task_id), evaluation_index=evaluation_index, gold_index=gold_index, mcp_config=args.mcp_config, scale_label=scale_label, ) ) args.out_json.parent.mkdir(parents=True, exist_ok=True) args.out_csv.parent.mkdir(parents=True, exist_ok=True) payload = {"scale": scale_label, "runs_root": str(args.runs_root), "rows": rows} args.out_json.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") fieldnames = [ "scale", "task", "results_match", "selected_tools", "overhead_planning_ratio", "gold_items", "context_tokens", "planning_latency_seconds", "selection_rate", "registered_tool_count", "completion_rate", "final_result_reached", "run_dir", ] with args.out_csv.open("w", encoding="utf-8", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fieldnames) writer.writeheader() for row in rows: writer.writerow({key: csv_cell(row.get(key)) for key in fieldnames}) print(json.dumps(payload, ensure_ascii=False, indent=2)) print(f"Saved task metrics JSON to: {args.out_json}") print(f"Saved task metrics CSV to: {args.out_csv}") return 0 if __name__ == "__main__": raise SystemExit(main())