File size: 4,488 Bytes
a9e46a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | 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
|