BioScientist/agent_system/mcp_router.py

from __future__ import annotations

import ast
from dataclasses import dataclass
from pathlib import Path
import shutil


@dataclass(frozen=True)
class MCPToolRoute:
    tool_name: str
    function_name: str
    binary_name: str
    server_script: Path
    server_dir: Path
    image_name: str
    service_name: str


class MCPToolRouter:
    """Dynamic router for MCP tool servers under mcp-servers/."""

    def __init__(self, project_root: str | Path):
        self.project_root = Path(project_root)
        self.servers_root = self.project_root / "mcp-servers"
        self.routes: dict[str, MCPToolRoute] = {}
        self.refresh()

    def refresh(self) -> None:
        self.routes = {}
        if not self.servers_root.exists():
            return

        for script in sorted(self.servers_root.glob("mcp_*/app/*_server.py")):
            server_dir = script.parent.parent
            server_name = server_dir.name.replace("mcp_", "")
            canonical_tool = server_name.replace("-", "_")
            image_name = f"mcp-{server_name}:latest"
            service_name = f"mcp-{server_name}"
            for function_name in self._discover_mcp_functions(script):
                route = MCPToolRoute(
                    tool_name=canonical_tool,
                    function_name=function_name,
                    binary_name=self._infer_binary_name(canonical_tool, function_name),
                    server_script=script,
                    server_dir=server_dir,
                    image_name=image_name,
                    service_name=service_name,
                )
                # 可按函数名或工具名取路由
                self.routes[function_name] = route
                self.routes.setdefault(canonical_tool, route)
                # 为常见对齐步骤补别名
                self._register_aliases(route)

    def resolve(self, tool_name: str) -> MCPToolRoute:
        if tool_name not in self.routes:
            raise KeyError(f"Unsupported tool route: {tool_name}")
        return self.routes[tool_name]

    def has_tool(self, tool_name: str) -> bool:
        return tool_name in self.routes

    def pick_first_available(self, candidates: list[str]) -> MCPToolRoute | None:
        for name in candidates:
            if name in self.routes:
                return self.routes[name]
        return None

    def binary_available(self, tool_name: str) -> bool:
        route = self.resolve(tool_name)
        return shutil.which(route.binary_name) is not None

    @staticmethod
    def _discover_mcp_functions(script_path: Path) -> list[str]:
        try:
            tree = ast.parse(script_path.read_text(encoding="utf-8"))
        except Exception:
            return []

        functions: list[str] = []
        for node in tree.body:
            if not isinstance(node, ast.FunctionDef):
                continue
            if any(MCPToolRouter._is_mcp_tool_decorator(d) for d in node.decorator_list):
                functions.append(node.name)
        return functions

    @staticmethod
    def _is_mcp_tool_decorator(decorator: ast.AST) -> bool:
        # @mcp.tool / @mcp.tool()
        if isinstance(decorator, ast.Call):
            decorator = decorator.func
        return (
            isinstance(decorator, ast.Attribute)
            and isinstance(decorator.value, ast.Name)
            and decorator.value.id == "mcp"
            and decorator.attr == "tool"
        )

    @staticmethod
    def _infer_binary_name(tool_name: str, function_name: str) -> str:
        overrides = {
            "trim_galore": "trim_galore",
            "bowtie2_align": "bowtie2",
            "bowtie2_build": "bowtie2-build",
            "bowtie2_inspect": "bowtie2-inspect",
        }
        if function_name in overrides:
            return overrides[function_name]
        if tool_name in overrides:
            return overrides[tool_name]
        return tool_name.replace("_", "-")

    def _register_aliases(self, route: MCPToolRoute) -> None:
        # 将常见 pipeline 逻辑名映射到具体 MCP 工具函数
        if route.function_name == "bowtie2_align":
            self.routes["bowtie2"] = route
        if route.function_name == "trim_galore":
            self.routes["trim_galore"] = route
            self.routes["trim-galore"] = route
        if route.function_name == "fastqc":
            self.routes["fastqc"] = route
        if route.function_name == "multiqc":
            self.routes["multiqc"] = route