#!/usr/bin/env python3 from __future__ import annotations import argparse import ast import json import os import re import sys from collections import Counter from datetime import datetime from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parent if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from biomni.graph import ToolGraph, ToolSchemaExtractor DEFAULT_MCP_ROOT = PROJECT_ROOT / "biomni_web" / "backend" / "data" / "mcp_generated" DEFAULT_HELP_ROOT = PROJECT_ROOT / "biomni_web" / "backend" / "data" / "merged_prefer_help_txt" DEFAULT_OUTPUT_ROOT = PROJECT_ROOT / "graph_outputs" BIOAGENT_BENCH_SERVERS = { # RNA-seq quantification / DESeq2 tasks "fastqc", "multiqc", "salmon", "kallisto", "hisat2", "star", "subread", "htseq", "bioconductor-tximport", "bioconductor-deseq2", "bioconductor-rsubread", "gffutils", "csvtk", # Variant calling tasks "bwa", "bowtie2", "minimap2", "samtools", "bcftools", "gatk", "gatk4", "freebayes", "bedtools", "pybedtools", # Comparative genomics / annotation tasks "blast", "diamond", "mafft", "orthofinder", "prokka", "busco", "abricate", "local-bio-cache", # Metagenomics tasks "kraken2", "metaphlan", "metaphlan2", } SERVER_CATEGORY_OVERRIDES = { "bioconductor-tximport": "transcriptomics", "htseq": "transcriptomics", "gffutils": "genomics", "fastqc": "transcriptomics", "multiqc": "transcriptomics", "salmon": "transcriptomics", "kallisto": "transcriptomics", "bioconductor-rsubread": "transcriptomics", "subread": "transcriptomics", "local-bio-cache": "genomics", "kraken2": "metagenomics", "metaphlan": "metagenomics", "metaphlan2": "metagenomics", } CUSTOM_SERVER_SPECS = ( { "name": "local-bio-cache", "category": "genomics", "source_path": PROJECT_ROOT / "biomni" / "tool" / "example_mcp_tools" / "local_bio_cache_mcp.py", "summary": ( "Offline local bio cache MCP server for registering FASTA collections, optionally building local BLAST " "databases, running local BLAST-like sequence search with Python fallback, running pairwise sequence " "alignment, and querying cached TF binding tables from GTRD, ENCODE, or ChIP-Atlas." ), "description": ( "Use this server when remote BLAST, UniProt sequence search, or TF binding APIs are unstable. " "It provides local sequence grounding, local UniProt-like search over predownloaded FASTA files, " "and offline TF binding lookup over curated flat files." ), "keywords": [ "offline", "local", "cache", "blast", "blastp", "fasta", "sequence", "alignment", "uniprot", "clinvar", "gtrd", "encode", "chip-atlas", "tf", "binding", "motif", "protein", "database", "variant", ], }, ) SERVER_CATEGORY_HINTS = { "transcriptomics": { "deseq2", "rsem", "star", "hisat2", "subread", "kallisto", "rsubread", "scrnaseq", "rnaseq", }, "pathway_enrichment": { "gseapy", "gsva", "pathway", "gsea", "go.db", }, "single_cell": { "sc", "scvi", "scrna", "seurat", "cellid", "cnmf", "doubletdetection", "cardspa", "pycistopic", "echidna", "spapros", }, "genomics": { "gatk", "bwa", "bcftools", "vcftools", "bamtools", "whatshap", "minimap2", "ngmlr", "tabix", "gff", "sam", "variant", "genome", "bowtie2", }, "proteomics": { "maxquant", "searchgui", "msproteomicstools", "sage-proteomics", }, "utility": { "jq", "csvtk", "coreutils", "urllib3", "ftputil", "dxpy", "pysftp", "json", "csv", }, } def resolve_python_command(python_cmd: str | None = None) -> str: if python_cmd: return python_cmd conda_prefix = Path(os.environ.get("CONDA_PREFIX", "")).expanduser() if conda_prefix: candidate = conda_prefix / "bin" / "python" if candidate.exists(): return str(candidate) return sys.executable def normalize_server_name(server_dir: Path) -> str: name = server_dir.name if name.startswith("mcp_"): name = name[len("mcp_") :] return name def safe_server_name(server_name: str) -> str: return re.sub(r"[^0-9a-zA-Z_]", "_", server_name) def tokenize(text: str) -> list[str]: return re.findall(r"[a-z0-9_\-\.]{2,}", (text or "").lower()) def infer_category(server_name: str, description: str, tools: list[dict[str, Any]], extra_text: str = "") -> str: combined = " ".join( [ server_name, description, extra_text, " ".join(tool.get("name", "") for tool in tools), " ".join(tool.get("description", "") for tool in tools), ] ).lower() scores: dict[str, int] = {} for category, hints in SERVER_CATEGORY_HINTS.items(): scores[category] = sum(1 for hint in hints if hint in combined) best_category = max(scores, key=scores.get) return best_category if scores[best_category] > 0 else "general" def parse_dockerfile_metadata(dockerfile_path: Path) -> dict[str, Any]: if not dockerfile_path.exists(): return {"packages": [], "copied_scripts": [], "cmd": "", "summary": ""} text = dockerfile_path.read_text(encoding="utf-8") packages = re.findall(r"conda install .*? ([a-zA-Z0-9_.+\-]+) -y", text) copied_scripts = re.findall(r"COPY\s+app/([^\s]+)\s+/app/", text) cmd_match = re.search(r'^CMD\s+\[(.*?)\]', text, flags=re.MULTILINE) summary_parts = [] if packages: summary_parts.append(f"Installs packages: {', '.join(sorted(set(packages))[:8])}") if copied_scripts: summary_parts.append(f"Copies scripts: {', '.join(copied_scripts[:4])}") if cmd_match: summary_parts.append(f"Default command: [{cmd_match.group(1)}]") return { "packages": sorted(set(packages)), "copied_scripts": copied_scripts, "cmd": cmd_match.group(1) if cmd_match else "", "summary": ". ".join(summary_parts), } def parse_help_metadata(server_name: str, help_root: Path) -> dict[str, Any]: help_path = help_root / f"{server_name}.txt" if not help_path.exists(): return { "path": "", "summary": "", "description": "", "execution_environment": "", "dependencies": [], "raw_excerpt": "", "text_for_index": "", } text = help_path.read_text(encoding="utf-8") lines = text.splitlines() metadata: dict[str, Any] = { "path": str(help_path.resolve()), "summary": "", "description": "", "execution_environment": "", "dependencies": [], "raw_excerpt": "", "text_for_index": "", } in_header = True header_fields: dict[str, str] = {} for line in lines: stripped = line.strip() if not stripped: if in_header: continue if stripped.startswith("## "): in_header = False continue if in_header and ":" in line: key, value = line.split(":", 1) header_fields[key.strip().lower()] = value.strip() metadata["summary"] = header_fields.get("summary", "") metadata["description"] = header_fields.get("description", "") metadata["execution_environment"] = header_fields.get("execution_environment", "") dependencies = header_fields.get("dependencies", "") metadata["dependencies"] = [item.strip() for item in dependencies.split(",") if item.strip()] if metadata["summary"] or metadata["description"]: excerpt_lines = [] if metadata["summary"]: excerpt_lines.append(metadata["summary"]) if metadata["description"]: excerpt_lines.append(metadata["description"]) metadata["raw_excerpt"] = " ".join(excerpt_lines) else: content_lines = [] for line in lines: stripped = line.strip() if not stripped: continue if stripped.startswith("$ ") or stripped.startswith("[rc="): continue content_lines.append(stripped) if len(content_lines) >= 12: break metadata["raw_excerpt"] = " ".join(content_lines) metadata["text_for_index"] = " ".join( part for part in [ metadata["summary"], metadata["description"], metadata["execution_environment"], " ".join(metadata["dependencies"][:20]), metadata["raw_excerpt"], ] if part ) return metadata def resolve_source_server(server_dir: Path) -> Path | None: app_dir = server_dir / "app" shim_candidates = sorted(app_dir.glob("*_shim_server.py")) raw_candidates = sorted( candidate for candidate in app_dir.glob("*_server.py") if not candidate.name.endswith("_shim_server.py") ) if raw_candidates: return raw_candidates[0] if shim_candidates: shim_text = shim_candidates[0].read_text(encoding="utf-8") match = re.search(r"SOURCE_SERVER\s*=\s*Path\((['\"])(.*?)\1\)", shim_text) if match: source_path = Path(match.group(2)) if source_path.exists(): return source_path return shim_candidates[0] return None def annotation_to_json_type(annotation: ast.AST | None) -> str: if annotation is None: return "string" annotation_text = ast.unparse(annotation).lower() if hasattr(ast, "unparse") else "" if any(name in annotation_text for name in ("int", "float")): return "number" if "bool" in annotation_text: return "boolean" if any(name in annotation_text for name in ("list", "tuple", "set")): return "array" if "dict" in annotation_text: return "object" return "string" def extract_tools_from_source(source_path: Path, server_name: str) -> list[dict[str, Any]]: try: tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) except (OSError, SyntaxError): return [] tools = [] for node in tree.body: if not isinstance(node, ast.FunctionDef) or node.name.startswith("_"): continue description = ast.get_docstring(node) or f"MCP tool {node.name} from {server_name}" properties = {} required = [] args = node.args.args defaults = list(node.args.defaults) first_optional_idx = len(args) - len(defaults) for idx, arg in enumerate(args): arg_name = arg.arg if arg_name in {"self", "cls"}: continue properties[arg_name] = { "type": annotation_to_json_type(arg.annotation), "description": "", } if idx < first_optional_idx: required.append(arg_name) tools.append( { "name": node.name, "description": description.splitlines()[0], "inputSchema": {"properties": properties, "required": required}, } ) return tools def build_fallback_tool(server_name: str, docker_meta: dict[str, Any], source_path: Path | None) -> dict[str, Any]: description_parts = [f"Fallback MCP tool entry for server {server_name}"] if docker_meta.get("packages"): description_parts.append(f"packages: {', '.join(docker_meta['packages'][:5])}") if source_path is not None: description_parts.append(f"source: {source_path.name}") return { "name": f"{safe_server_name(server_name)}_tool", "description": "; ".join(description_parts), "inputSchema": {"properties": {}, "required": []}, } def build_server_entry(server_dir: Path, help_root: Path, *, python_cmd: str | None = None) -> dict[str, Any]: server_name = normalize_server_name(server_dir) dockerfile_path = server_dir / "Dockerfile" docker_meta = parse_dockerfile_metadata(dockerfile_path) help_meta = parse_help_metadata(server_name, help_root) source_path = resolve_source_server(server_dir) tools = extract_tools_from_source(source_path, server_name) if source_path is not None else [] if not tools: tools = [build_fallback_tool(server_name, docker_meta, source_path)] summary_parts = [] if help_meta.get("summary"): summary_parts.append(help_meta["summary"]) if help_meta.get("description"): summary_parts.append(help_meta["description"]) if docker_meta.get("summary"): summary_parts.append(docker_meta["summary"]) if help_meta.get("execution_environment"): summary_parts.append(f"Execution environment: {help_meta['execution_environment']}") if source_path is not None: summary_parts.append(f"Resolved source server: {source_path}") summary = " | ".join(summary_parts) or f"Auto-indexed MCP server for {server_name}" category = infer_category(server_name, summary, tools, extra_text=help_meta.get("text_for_index", "")) category = SERVER_CATEGORY_OVERRIDES.get(server_name, category) keywords = tokenize( " ".join( [ server_name, category, summary, help_meta.get("text_for_index", ""), " ".join(docker_meta.get("packages", [])), " ".join(tool.get("name", "") for tool in tools), " ".join(tool.get("description", "") for tool in tools), ] ) ) python_cmd = resolve_python_command(python_cmd) command = [] shim_candidates = sorted((server_dir / "app").glob("*_shim_server.py")) raw_candidates = sorted( candidate for candidate in (server_dir / "app").glob("*_server.py") if not candidate.name.endswith("_shim_server.py") ) if shim_candidates: command = [python_cmd, str(shim_candidates[0].resolve())] elif raw_candidates: command = [python_cmd, str(raw_candidates[0].resolve())] return { "name": server_name, "safe_name": safe_server_name(server_name), "module": f"mcp_servers.{safe_server_name(server_name)}", "category": category, "summary": summary, "keywords": keywords[:60], "command_available": bool(command), "tools": tools, "tool_names": [tool["name"] for tool in tools], "server_meta": { "enabled": True, "description": summary, "command": command, "command_available": bool(command), "dockerfile": str(dockerfile_path.resolve()) if dockerfile_path.exists() else "", "source_server": str(source_path.resolve()) if source_path is not None and source_path.exists() else "", "packages": docker_meta.get("packages", []), "help_file": help_meta.get("path", ""), "help_summary": help_meta.get("summary", ""), "help_description": help_meta.get("description", ""), "execution_environment": help_meta.get("execution_environment", ""), "dependencies": help_meta.get("dependencies", []), }, } def build_custom_server_entry(spec: dict[str, Any], *, python_cmd: str | None = None) -> dict[str, Any] | None: source_path = Path(spec["source_path"]).resolve() if not source_path.exists(): return None server_name = spec["name"] tools = extract_tools_from_source(source_path, server_name) if not tools: tools = [ { "name": f"{safe_server_name(server_name)}_tool", "description": spec["summary"], "inputSchema": {"properties": {}, "required": []}, } ] python_cmd = resolve_python_command(python_cmd) summary = " | ".join(part for part in [spec.get("summary", ""), spec.get("description", "")] if part) keywords = tokenize( " ".join( [ server_name, spec.get("category", ""), summary, " ".join(spec.get("keywords", [])), " ".join(tool.get("name", "") for tool in tools), " ".join(tool.get("description", "") for tool in tools), ] ) ) command = [python_cmd, str(source_path)] return { "name": server_name, "safe_name": safe_server_name(server_name), "module": f"mcp_servers.{safe_server_name(server_name)}", "category": spec.get("category", "general"), "summary": summary or f"Custom MCP server for {server_name}", "keywords": keywords[:80], "command_available": True, "tools": tools, "tool_names": [tool["name"] for tool in tools], "server_meta": { "enabled": True, "description": summary or f"Custom MCP server for {server_name}", "command": command, "command_available": True, "dockerfile": "", "source_server": str(source_path), "packages": [], "help_file": "", "help_summary": spec.get("summary", ""), "help_description": spec.get("description", ""), "execution_environment": "local python process", "dependencies": [], "custom_server": True, }, } def summarize_graph(tool_graph: ToolGraph) -> dict[str, Any]: node_type_counts = Counter(node["type"] for node in tool_graph.nodes.values()) edge_type_counts = Counter(edge["type"] for edge in tool_graph.edges) category_counts = Counter( tool_graph.server_index[server_name]["entry"].get("category", "general") for server_name in tool_graph.server_index ) capability_counts = Counter() datatype_counts = Counter() stage_counts = Counter() for server_name in tool_graph.server_index: semantics = tool_graph.get_server_semantics(server_name) capability_counts.update(semantics.get("capabilities", [])) datatype_counts.update(semantics.get("datatypes", [])) stage_counts.update(semantics.get("stages", [])) return { "server_count": len(tool_graph.server_entries), "tool_count": len([node for node in tool_graph.nodes.values() if node["type"] == "tool"]), "node_count": len(tool_graph.nodes), "edge_count": len(tool_graph.edges), "node_type_counts": dict(node_type_counts), "edge_type_counts": dict(edge_type_counts), "top_categories": category_counts.most_common(20), "top_capabilities": capability_counts.most_common(20), "top_datatypes": datatype_counts.most_common(20), "top_stages": stage_counts.most_common(20), } def write_json(path: Path, payload: Any) -> None: path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") def load_server_catalog(graph_dir: Path) -> list[dict[str, Any]]: server_catalog_path = graph_dir / "server_catalog.json" if not server_catalog_path.exists(): raise FileNotFoundError(f"server_catalog.json not found under {graph_dir}") payload = json.loads(server_catalog_path.read_text(encoding="utf-8")) if not isinstance(payload, list): raise ValueError(f"Expected a list of server entries in {server_catalog_path}") return payload def filter_server_entries( server_entries: list[dict[str, Any]], *, preset: str = "all", include_servers: set[str] | None = None, ) -> list[dict[str, Any]]: allowed = set(include_servers or []) if preset == "bioagent-bench": allowed.update(BIOAGENT_BENCH_SERVERS) if not allowed: return server_entries return [entry for entry in server_entries if entry["name"] in allowed] def build_graph( mcp_root: Path, help_root: Path, output_root: Path, *, preset: str = "all", include_servers: set[str] | None = None, python_cmd: str | None = None, ) -> Path: if not mcp_root.is_dir(): raise SystemExit(f"MCP root does not exist or is not a directory: {mcp_root}") if not help_root.is_dir(): raise SystemExit(f"Help root does not exist or is not a directory: {help_root}") timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") suffix = "benchmark" if preset == "bioagent-bench" else "all" run_dir = output_root / f"mcp_generated_graph_{suffix}_{timestamp}" run_dir.mkdir(parents=True, exist_ok=True) server_dirs = sorted(path for path in mcp_root.glob("mcp_*") if path.is_dir()) server_entries = [build_server_entry(server_dir, help_root, python_cmd=python_cmd) for server_dir in server_dirs] custom_entries = [ entry for spec in CUSTOM_SERVER_SPECS if (entry := build_custom_server_entry(spec, python_cmd=python_cmd)) is not None ] if custom_entries: existing_names = {entry["name"] for entry in server_entries} for entry in custom_entries: if entry["name"] not in existing_names: server_entries.append(entry) server_entries = filter_server_entries(server_entries, preset=preset, include_servers=include_servers) schema_extractor = ToolSchemaExtractor() tool_graph = ToolGraph(schema_extractor=schema_extractor) tool_graph.build_from_server_entries(server_entries) summary = summarize_graph(tool_graph) summary["preset"] = preset summary["included_servers"] = sorted(entry["name"] for entry in server_entries) server_semantics = { server_name: tool_graph.get_server_semantics(server_name) for server_name in sorted(tool_graph.get_server_names()) } write_json(run_dir / "server_catalog.json", server_entries) write_json(run_dir / "graph_nodes.json", list(tool_graph.nodes.values())) write_json(run_dir / "graph_edges.json", tool_graph.edges) write_json(run_dir / "server_semantics.json", server_semantics) write_json(run_dir / "graph_summary.json", summary) print(f"Built graph for {summary['server_count']} MCP servers.") print(f"Generated {summary['tool_count']} tool nodes, {summary['node_count']} total nodes, {summary['edge_count']} edges.") print(f"Artifacts written to: {run_dir}") return run_dir def main() -> None: parser = argparse.ArgumentParser(description="Build a graph for all MCP servers under mcp_generated.") parser.add_argument("--mcp-root", default=str(DEFAULT_MCP_ROOT), help="Root directory containing mcp_* server dirs.") parser.add_argument( "--help-root", default=str(DEFAULT_HELP_ROOT), help="Directory containing merged help txt files for MCP servers.", ) parser.add_argument("--output-root", default=str(DEFAULT_OUTPUT_ROOT), help="Directory to store graph artifacts.") parser.add_argument( "--preset", choices=["all", "bioagent-bench"], default="all", help="Graph scope preset. bioagent-bench keeps only high-signal servers for local benchmark tasks.", ) parser.add_argument( "--include-server", action="append", default=[], help="Additional server name to include. Can be repeated.", ) parser.add_argument( "--python-cmd", default=None, help="Python interpreter to bake into shim server commands. Defaults to CONDA_PREFIX/bin/python when available.", ) args = parser.parse_args() build_graph( Path(args.mcp_root).resolve(), Path(args.help_root).resolve(), Path(args.output_root).resolve(), preset=args.preset, include_servers=set(args.include_server), python_cmd=args.python_cmd, ) if __name__ == "__main__": main()