Biomanus/experiments/bioinfomcp_benchmark/scripts/build_benchmark_subset.py

#!/usr/bin/env python3
"""Build a deterministic 500-server benchmark subset for BioinfoMCP conversion."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path("/225040511/project/Hypo_Bio_OS")
DEFAULT_GRAPH_DIR = PROJECT_ROOT / "graph_outputs" / "mcp_generated_graph_all_20260514_100123"
DEFAULT_HELP_ROOT = PROJECT_ROOT / "biomni_web" / "backend" / "data" / "merged_prefer_help_txt"
DEFAULT_MCP_ROOT = PROJECT_ROOT / "biomni_web" / "backend" / "data" / "mcp_generated"
DEFAULT_OUT_JSON = PROJECT_ROOT / "experiments" / "bioinfomcp_benchmark" / "configs" / "benchmark_subset_500.json"
DEFAULT_OUT_CSV = PROJECT_ROOT / "experiments" / "bioinfomcp_benchmark" / "configs" / "benchmark_subset_500.csv"


def load_server_catalog(graph_dir: Path) -> list[dict[str, Any]]:
    return json.loads((graph_dir / "server_catalog.json").read_text(encoding="utf-8"))


def effective_code_lines(path: Path) -> int:
    count = 0
    for line in path.read_text(encoding="utf-8", errors="ignore").splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        count += 1
    return count


def stable_key(seed: int, text: str) -> str:
    return hashlib.sha1(f"{seed}:{text}".encode("utf-8")).hexdigest()


def quantile_bucket(value: float, boundaries: tuple[float, float]) -> str:
    low, high = boundaries
    if value <= low:
        return "low"
    if value <= high:
        return "mid"
    return "high"


def allocate_counts(total: int, groups: dict[str, int]) -> dict[str, int]:
    total_available = sum(groups.values())
    if total >= total_available:
        return dict(groups)

    raw = {name: total * size / total_available for name, size in groups.items()}
    counts = {name: min(groups[name], math.floor(value)) for name, value in raw.items()}
    remainder = total - sum(counts.values())
    order = sorted(groups, key=lambda name: (raw[name] - counts[name], groups[name]), reverse=True)
    while remainder > 0:
        progressed = False
        for name in order:
            if counts[name] < groups[name]:
                counts[name] += 1
                remainder -= 1
                progressed = True
                if remainder == 0:
                    break
        if not progressed:
            break
    return counts


def allocate_bucket_counts(total: int, bucket_sizes: dict[str, int]) -> dict[str, int]:
    nonzero = {name: size for name, size in bucket_sizes.items() if size > 0}
    if not nonzero:
        return {name: 0 for name in bucket_sizes}
    if total >= sum(nonzero.values()):
        return {name: bucket_sizes[name] for name in bucket_sizes}

    counts = {name: 0 for name in bucket_sizes}
    if total >= len(nonzero):
        for name in nonzero:
            counts[name] = 1
        remaining = total - len(nonzero)
    else:
        remaining = total

    extras = {name: nonzero[name] - counts[name] for name in nonzero}
    extra_counts = allocate_counts(remaining, extras) if remaining > 0 else {name: 0 for name in nonzero}
    for name, value in extra_counts.items():
        counts[name] += value
    return counts


def build_candidates(graph_dir: Path, help_root: Path, mcp_root: Path) -> list[dict[str, Any]]:
    help_names = {path.stem for path in help_root.glob("*.txt")}
    entries = []
    for entry in load_server_catalog(graph_dir):
        name = entry["name"]
        if name not in help_names:
            continue
        if " copy" in name:
            continue
        source_value = entry.get("server_meta", {}).get("source_server", "")
        if not source_value:
            continue
        source_server = Path(source_value)
        if not source_server.exists() or not source_server.is_file():
            continue
        server_dir = mcp_root / f"mcp_{name}"
        if not server_dir.exists():
            continue

        help_path = help_root / f"{name}.txt"
        help_text = help_path.read_text(encoding="utf-8", errors="ignore")
        help_lines = len(help_text.splitlines())
        help_chars = len(help_text)
        gold_loc = effective_code_lines(source_server)
        tool_count = len(entry.get("tools", []))

        entries.append(
            {
                "server_name": name,
                "category": entry.get("category", "general"),
                "server_dir": str(server_dir),
                "help_path": str(help_path),
                "gold_source_path": str(source_server),
                "tool_count": tool_count,
                "gold_code_lines": gold_loc,
                "help_lines": help_lines,
                "help_chars": help_chars,
                "keywords": entry.get("keywords", []),
                "summary": entry.get("summary", ""),
            }
        )
    return entries


def attach_complexity(candidates: list[dict[str, Any]]) -> None:
    help_values = sorted(item["help_lines"] for item in candidates)
    loc_values = sorted(item["gold_code_lines"] for item in candidates)
    tool_values = sorted(item["tool_count"] for item in candidates)

    def percentile_bounds(values: list[int]) -> tuple[float, float]:
        n = len(values)
        return values[n // 3], values[(2 * n) // 3]

    help_bounds = percentile_bounds(help_values)
    loc_bounds = percentile_bounds(loc_values)
    tool_bounds = percentile_bounds(tool_values)

    for item in candidates:
        help_bucket = quantile_bucket(item["help_lines"], help_bounds)
        loc_bucket = quantile_bucket(item["gold_code_lines"], loc_bounds)
        tool_bucket = quantile_bucket(item["tool_count"], tool_bounds)
        score = {"low": 0, "mid": 1, "high": 2}[help_bucket]
        score += {"low": 0, "mid": 1, "high": 2}[loc_bucket]
        score += {"low": 0, "mid": 1, "high": 2}[tool_bucket]
        if score <= 1:
            complexity = "low"
        elif score <= 3:
            complexity = "mid"
        else:
            complexity = "high"
        item["complexity"] = complexity


def select_subset(candidates: list[dict[str, Any]], subset_size: int, seed: int) -> list[dict[str, Any]]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for item in candidates:
        grouped[item["category"]].append(item)

    category_targets = allocate_counts(subset_size, {name: len(items) for name, items in grouped.items()})
    selected: list[dict[str, Any]] = []

    for category, items in sorted(grouped.items()):
        buckets: dict[str, list[dict[str, Any]]] = defaultdict(list)
        for item in items:
            buckets[item["complexity"]].append(item)
        for bucket_items in buckets.values():
            bucket_items.sort(key=lambda x: stable_key(seed, x["server_name"]))

        bucket_targets = allocate_bucket_counts(
            category_targets[category],
            {name: len(buckets.get(name, [])) for name in ("low", "mid", "high")},
        )
        for complexity in ("low", "mid", "high"):
            selected.extend(buckets.get(complexity, [])[: bucket_targets[complexity]])

    selected.sort(key=lambda x: (x["category"], x["server_name"]))
    return selected


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--graph-dir", type=Path, default=DEFAULT_GRAPH_DIR)
    parser.add_argument("--help-root", type=Path, default=DEFAULT_HELP_ROOT)
    parser.add_argument("--mcp-root", type=Path, default=DEFAULT_MCP_ROOT)
    parser.add_argument("--subset-size", type=int, default=500)
    parser.add_argument("--seed", type=int, default=20260514)
    parser.add_argument("--out-json", type=Path, default=DEFAULT_OUT_JSON)
    parser.add_argument("--out-csv", type=Path, default=DEFAULT_OUT_CSV)
    args = parser.parse_args()

    candidates = build_candidates(args.graph_dir, args.help_root, args.mcp_root)
    attach_complexity(candidates)
    selected = select_subset(candidates, args.subset_size, args.seed)

    category_counts = Counter(item["category"] for item in selected)
    complexity_counts = Counter(item["complexity"] for item in selected)
    payload = {
        "metadata": {
            "subset_size": len(selected),
            "seed": args.seed,
            "graph_dir": str(args.graph_dir),
            "help_root": str(args.help_root),
            "mcp_root": str(args.mcp_root),
            "candidate_count": len(candidates),
            "category_counts": dict(category_counts),
            "complexity_counts": dict(complexity_counts),
        },
        "items": selected,
    }

    args.out_json.parent.mkdir(parents=True, exist_ok=True)
    args.out_json.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")

    with args.out_csv.open("w", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(
            fh,
            fieldnames=[
                "server_name",
                "category",
                "complexity",
                "tool_count",
                "gold_code_lines",
                "help_lines",
                "help_chars",
                "server_dir",
                "help_path",
                "gold_source_path",
            ],
            extrasaction="ignore",
        )
        writer.writeheader()
        writer.writerows(selected)

    print(json.dumps(payload["metadata"], indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())