Biomanus_upload/build_generated_mcp_graph.py

#!/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 / "demo_web" / "backend" / "data" / "mcp_generated"
DEFAULT_HELP_ROOT = PROJECT_ROOT / "demo_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()