experiments/bioagent_bench/scripts/generate_mcp_scale_configs.py

#!/usr/bin/env python3
"""Generate Biomni MCP configs at approximately 100, 500, 1000, 1500, and 2000 tools.

The generated YAML files use explicit tool metadata so Biomni's ``A1.add_mcp``
does not need to start every MCP server during registration.
"""

from __future__ import annotations

import argparse
import ast
import json
import re
import sys
from pathlib import Path
from typing import Any

try:
    import yaml
except ImportError as exc:  # pragma: no cover
    raise SystemExit("PyYAML is required: pip install pyyaml") from exc


DEFAULT_TARGETS = (100, 500, 1000, 1500, 2000)
REPO_ROOT = Path(__file__).resolve().parents[3]
DEFAULT_MCP_ROOT = REPO_ROOT / "mcp_generated"
DEFAULT_OUT_DIR = Path(__file__).resolve().parents[1] / "configs"
DEFAULT_GOLD = Path(__file__).resolve().parents[1] / "gold_tools.json"
DEFAULT_EXECUTION_PYTHON = Path("/225040511/miniconda3/envs/biomni_e1/bin/python")

PRIORITY_SERVER_STEMS = [
    "mcp_kallisto",
    "mcp_salmon",
    "mcp_kraken2",
    "mcp_bracken",
    "mcp_csvtk",
    "mcp_megahit",
    "mcp_kaiju",
    "mcp_fastp",
    "mcp_spades",
    "mcp_scanpy",
    "mcp_scanpy-cli",
    "mcp_scanpy-scripts",
    "mcp_seurat-scripts",
    "mcp_bioconductor-deseq2",
    "mcp_bioconductor-edger",
    "mcp_bioconductor-limma",
    "mcp_bioconductor-clusterprofiler",
    "mcp_star",
    "mcp_hisat2",
    "mcp_subread",
    "mcp_samtools",
    "mcp_bcftools",
    "mcp_bwa",
    "mcp_gatk4",
    "mcp_gatk",
    "mcp_snpeff",
    "mcp_freebayes",
    "mcp_orthofinder",
    "mcp_blast",
    "mcp_blast-legacy",
    "mcp_diamond",
    "mcp_gffread",
    "mcp_seqkit",
    "mcp_multiqc",
    "mcp_fastqc",
    "mcp_trimmomatic",
]


def safe_server_name(dirname: str) -> str:
    name = dirname[4:] if dirname.startswith("mcp_") else dirname
    return re.sub(r"[^A-Za-z0-9_]", "_", name)


def choose_server_file(server_dir: Path) -> Path | None:
    app_dir = server_dir / "app"
    if not app_dir.exists():
        return None
    shims = sorted(app_dir.glob("*_shim_server.py"))
    servers = sorted(p for p in app_dir.glob("*_server.py") if not p.name.endswith("_shim_server.py"))
    if shims:
        return shims[0]
    if servers:
        return servers[0]
    return None


def source_file_for_metadata(server_file: Path) -> Path:
    if not server_file.name.endswith("_shim_server.py"):
        return server_file
    local_name = server_file.name.replace("_shim_server.py", "_server.py")
    local = server_file.with_name(local_name)
    return local if local.exists() else server_file


def annotation_to_json_type(node: ast.AST | None) -> str:
    if node is None:
        return "string"
    text = ast.unparse(node).lower()
    if any(x in text for x in ("int", "float", "double")):
        return "number"
    if "bool" in text:
        return "boolean"
    if any(x in text for x in ("list", "tuple", "set", "sequence")):
        return "array"
    return "string"


def is_mcp_tool_decorator(dec: ast.AST) -> bool:
    target = dec.func if isinstance(dec, ast.Call) else dec
    if isinstance(target, ast.Attribute) and target.attr == "tool":
        return True
    if isinstance(target, ast.Name) and target.id in {"tool", "mcp_tool"}:
        return True
    return False


def parse_tools(source_file: Path, server_stem: str) -> list[dict[str, Any]]:
    try:
        tree = ast.parse(source_file.read_text(encoding="utf-8"), filename=str(source_file))
    except Exception:
        return []

    tools: list[dict[str, Any]] = []
    for node in tree.body:
        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            continue
        if node.name.startswith("_"):
            continue
        if not any(is_mcp_tool_decorator(dec) for dec in node.decorator_list):
            continue
        doc = ast.get_docstring(node) or f"MCP tool {node.name} from {server_stem}"
        first_line = " ".join(doc.strip().splitlines()[0].split()) if doc.strip() else ""
        properties: dict[str, dict[str, Any]] = {}
        required: list[str] = []
        defaults = [None] * (len(node.args.args) - len(node.args.defaults)) + list(node.args.defaults)
        for arg, default in zip(node.args.args, defaults, strict=False):
            if arg.arg in {"self", "cls"}:
                continue
            spec: dict[str, Any] = {
                "type": annotation_to_json_type(arg.annotation),
                "description": "",
            }
            if default is None:
                required.append(arg.arg)
            else:
                with contextlib_suppress(Exception):
                    spec["default"] = ast.literal_eval(default)
            properties[arg.arg] = spec
        tools.append(
            {
                "name": node.name,
                "description": first_line or f"MCP tool {node.name} from {server_stem}",
                "inputSchema": {"properties": properties, "required": required},
            }
        )
    return tools


class contextlib_suppress:
    def __init__(self, *exceptions: type[BaseException]):
        self.exceptions = exceptions or (Exception,)

    def __enter__(self):
        return None

    def __exit__(self, exc_type, exc, traceback):
        return exc_type is not None and issubclass(exc_type, self.exceptions)


def discover_servers(mcp_root: Path, execution_python: Path) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for server_dir in sorted(p for p in mcp_root.iterdir() if p.is_dir() and p.name.startswith("mcp_")):
        server_file = choose_server_file(server_dir)
        if server_file is None:
            continue
        metadata_file = source_file_for_metadata(server_file)
        tools = parse_tools(metadata_file, server_dir.name)
        if not tools:
            continue
        records.append(
            {
                "server_dir": server_dir.name,
                "server_name": safe_server_name(server_dir.name),
                "path": str(server_dir),
                "command": [str(execution_python), str(server_file)],
                "metadata_source": str(metadata_file),
                "tool_count": len(tools),
                "tools": tools,
            }
        )
    return records


def priority_order(records: list[dict[str, Any]], gold_path: Path) -> list[dict[str, Any]]:
    gold_servers: list[str] = []
    if gold_path.exists():
        gold = json.loads(gold_path.read_text(encoding="utf-8"))
        for entry in gold.values():
            gold_servers.extend(entry.get("gold_servers", []))
    priority = []
    seen = set()
    for name in [*PRIORITY_SERVER_STEMS, *gold_servers]:
        if name not in seen:
            priority.append(name)
            seen.add(name)
    rank = {name: i for i, name in enumerate(priority)}
    return sorted(records, key=lambda r: (rank.get(r["server_dir"], 10_000), r["server_dir"]))


def select_for_target(records: list[dict[str, Any]], target: int) -> list[dict[str, Any]]:
    selected: list[dict[str, Any]] = []
    total = 0
    for record in records:
        if total >= target:
            break
        selected.append(record)
        total += record["tool_count"]
    return selected


def write_yaml(selected: list[dict[str, Any]], out_path: Path) -> None:
    config = {"mcp_servers": {}}
    for record in selected:
        config["mcp_servers"][record["server_name"]] = {
            "enabled": True,
            "description": f"Generated MCP server config from {record['server_dir']} ({record['tool_count']} tools)",
            "command": record["command"],
            "source_dir": record["path"],
            "metadata_source": record["metadata_source"],
            "tools": record["tools"],
        }
    out_path.write_text(yaml.safe_dump(config, sort_keys=False, width=120), encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--mcp-root", type=Path, default=DEFAULT_MCP_ROOT)
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT_DIR)
    parser.add_argument("--gold", type=Path, default=DEFAULT_GOLD)
    parser.add_argument("--execution-python", type=Path, default=DEFAULT_EXECUTION_PYTHON)
    parser.add_argument("--targets", type=int, nargs="+", default=list(DEFAULT_TARGETS))
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    execution_python = args.execution_python if args.execution_python.exists() else Path(sys.executable)
    records = priority_order(discover_servers(args.mcp_root, execution_python), args.gold)
    if not records:
        raise SystemExit(f"No MCP servers with parseable tools found under {args.mcp_root}")

    manifest: dict[str, Any] = {
        "mcp_root": str(args.mcp_root),
        "execution_python": str(execution_python),
        "available_server_count": len(records),
        "available_tool_count": sum(r["tool_count"] for r in records),
        "targets": {},
    }
    (args.out_dir / "mcp_server_inventory.json").write_text(
        json.dumps(records, indent=2, ensure_ascii=False), encoding="utf-8"
    )

    for target in args.targets:
        selected = select_for_target(records, target)
        total = sum(r["tool_count"] for r in selected)
        out_path = args.out_dir / f"mcp_scale_{target}.yaml"
        write_yaml(selected, out_path)
        manifest["targets"][str(target)] = {
            "config": str(out_path),
            "server_count": len(selected),
            "tool_count": total,
            "servers": [r["server_dir"] for r in selected],
        }

    (args.out_dir / "mcp_scale_manifest.json").write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8"
    )
    print(json.dumps(manifest, indent=2, ensure_ascii=False))
    return 0


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