Beyond_Prompt-based_Retrieval / Biomanus /build_generated_mcp_graph.py
czty's picture
Add files using upload-large-folder tool
d1ce356 verified
Raw
History Blame Contribute Delete
24.5 kB
#!/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()