experiments/bioagent_bench/scripts/summarize_biomni_task_metrics.py

#!/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())