Biomanus_upload/mcp_register.py

#!/usr/bin/env python3
from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path

from biomni.agent import A1


PROJECT_ROOT = Path(__file__).resolve().parent
MCP_ROOT = PROJECT_ROOT / "biomni_web" / "backend" / "data" / "mcp_generated"
CONFIG_PATH = PROJECT_ROOT / "strata_mcp_smoke_config_shim.yaml"

# Start with a small set of MCP servers that have already shown successful tool discovery.
SELECTED_SERVERS = [
    "bioconductor-cardspa",
    "bioconductor-catscradle",
    "jq",
]


def find_server_script(server_name: str) -> Path:
    server_dir = MCP_ROOT / f"mcp_{server_name}" / "app"
    shim_candidates = sorted(server_dir.glob("*_shim_server.py"))
    if shim_candidates:
        return shim_candidates[0].resolve()

    raw_candidates = sorted(
        candidate
        for candidate in server_dir.glob("*_server.py")
        if not candidate.name.endswith("_shim_server.py")
    )
    if raw_candidates:
        return raw_candidates[0].resolve()

    raise FileNotFoundError(f"No MCP server script found for {server_name}: {server_dir}")


def write_smoke_config(selected_servers: list[str]) -> Path:
    lines = [
        "# Auto-generated smoke-test MCP config",
        "",
        "mcp_servers:",
    ]
    python_cmd = os.getenv("BIOMNI_MCP_PYTHON", os.sys.executable)

    for server_name in selected_servers:
        server_script = find_server_script(server_name)
        lines.extend(
            [
                f"  {server_name}:",
                "    enabled: true",
                f'    command: ["{python_cmd}", "{server_script}"]',
                f'    description: "Smoke-test MCP server for {server_name}"',
            ]
        )

    CONFIG_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return CONFIG_PATH


def build_agent() -> A1:
    # The smoke tests below directly call MCP tool wrappers, so a real LLM key is
    # only needed if you later switch this script back to agent.go(...).
    api_key = os.getenv("DEEPSEEK_API_KEY") or os.getenv("OPENAI_API_KEY") or "EMPTY"
    base_url = os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")
    model = os.getenv("DEEPSEEK_MODEL_NAME", "deepseek-chat")

    return A1(
        path="./data",
        llm=model,
        source="Custom",
        base_url=base_url,
        api_key=api_key,
        expected_data_lake_files=[],
    )


def run_smoke_tests(agent: A1) -> list[dict]:
    temp_dir = Path(tempfile.gettempdir()) / "biomni_mcp_smoke"
    temp_dir.mkdir(parents=True, exist_ok=True)
    missing_input = temp_dir / "missing_input.rds"

    test_cases = [
        {
            "tool_name": "cardspa",
            "task": "Run cardspa on a placeholder spatial object to verify MCP invocation.",
            "kwargs": {
                "sce_object": str(missing_input),
                "output_path": str(temp_dir / "cardspa_output.rds"),
                "phenotype_col": "cell_type",
            },
        },
        {
            "tool_name": "catscradle_build_neighborhoods",
            "task": "Build neighborhoods from a placeholder RDS file to verify CatsCradle MCP invocation.",
            "kwargs": {
                "input_rds": str(missing_input),
                "output_rds": str(temp_dir / "catscradle_neighborhoods.rds"),
            },
        },
        {
            "tool_name": "catscradle_gene_centric_analysis",
            "task": "Run gene-centric analysis on a placeholder RDS file to verify CatsCradle MCP invocation.",
            "kwargs": {
                "input_rds": str(missing_input),
                "output_rds": str(temp_dir / "catscradle_gene_centric.rds"),
            },
        },
        {
            "tool_name": "jq_process_json",
            "task": "Run jq on a small real JSON file to verify end-to-end MCP tool execution.",
            "kwargs": {
                "jq_filter": "[.[] | .score] | add / length",
                "input_files": [str(_write_demo_json(temp_dir))],
            },
        },
    ]

    results: list[dict] = []
    for case in test_cases:
        tool_name = case["tool_name"]
        wrapper = agent.get_custom_tool(tool_name)
        if wrapper is None:
            results.append(
                {
                    "tool_name": tool_name,
                    "task": case["task"],
                    "status": "not_registered",
                    "detail": "Tool wrapper not found after MCP registration.",
                }
            )
            continue

        try:
            tool_result = wrapper(**case["kwargs"])
            results.append(
                {
                    "tool_name": tool_name,
                    "task": case["task"],
                    "status": "called",
                    "detail": tool_result,
                }
            )
        except Exception as exc:  # pragma: no cover - smoke test reporting
            results.append(
                {
                    "tool_name": tool_name,
                    "task": case["task"],
                    "status": "call_failed",
                    "detail": str(exc),
                }
            )

    return results


def _write_demo_json(output_dir: Path) -> Path:
    demo_file = output_dir / "demo_scores.json"
    payload = [
        {"name": "sample_a", "score": 10},
        {"name": "sample_b", "score": 25},
        {"name": "sample_c", "score": 40},
    ]
    demo_file.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    return demo_file


def run_biomni_real_task(agent: A1) -> str:
    temp_dir = Path(tempfile.gettempdir()) / "biomni_mcp_smoke"
    temp_dir.mkdir(parents=True, exist_ok=True)
    demo_file = _write_demo_json(temp_dir)

    query = (
        "Use MCP tool jq_process_json to process the JSON file "
        f"'{demo_file}'. "
        "Compute three things: "
        "(1) average score, "
        "(2) max score item name, "
        "(3) number of records. "
        "Please call the MCP tool directly and then report final numeric results."
    )
    _, final_answer = agent.go(query)
    return final_answer


def main() -> None:
    config_path = write_smoke_config(SELECTED_SERVERS)
    print(f"Smoke-test config written to: {config_path}")

    agent = build_agent()
    agent.add_mcp(config_path=str(config_path))

    registered_tools = sorted(agent.list_custom_tools())
    print("\n===== REGISTERED MCP TOOLS =====")
    for tool_name in registered_tools:
        print(tool_name)

    print("\n===== SMOKE TEST RESULTS =====")
    for result in run_smoke_tests(agent):
        print(json.dumps(result, ensure_ascii=False, indent=2))

    print("\n===== BIOMNI REAL TASK (MCP-DRIVEN) =====")
    try:
        final_answer = run_biomni_real_task(agent)
        print(final_answer)
    except Exception as exc:  # pragma: no cover - runtime integration reporting
        print(f"Biomni real-task run failed: {exc}")


if __name__ == "__main__":
    main()