04db423416
- 70 skills with code and documentation - Add .gitignore (ignore __pycache__, output/, temp/, venv/) - Clean up test intermediates and caches
160 lines
4.9 KiB
Python
160 lines
4.9 KiB
Python
"""
|
|
MCP (Model Context Protocol) 适配器
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from .base import ProtocolAdapter, Connection, AgentInfo
|
|
|
|
try:
|
|
from mcp import ClientSession, StdioServerParameters
|
|
from mcp.client.stdio import stdio_client
|
|
HAS_MCP = True
|
|
except ImportError:
|
|
HAS_MCP = False
|
|
|
|
|
|
class MCPAdapter(ProtocolAdapter):
|
|
"""MCP 协议适配器"""
|
|
|
|
protocol_name = "mcp"
|
|
|
|
def __init__(self, config_dir: Optional[Path] = None):
|
|
self.config_dir = config_dir or Path(__file__).parent.parent / "config"
|
|
self._sessions: Dict[str, Any] = {}
|
|
|
|
async def connect(self, agent_config: dict) -> Connection:
|
|
"""建立连接"""
|
|
if not HAS_MCP:
|
|
raise ImportError("请安装 mcp 库: pip install mcp")
|
|
|
|
command = agent_config.get("command")
|
|
args = agent_config.get("args", [])
|
|
env = agent_config.get("env", {})
|
|
|
|
if not command:
|
|
raise ValueError("MCP Agent 配置必须包含 command")
|
|
|
|
full_env = os.environ.copy()
|
|
for k, v in env.items():
|
|
if v.startswith("${") and v.endswith("}"):
|
|
env_var = v[2:-1]
|
|
full_env[k] = os.environ.get(env_var, "")
|
|
else:
|
|
full_env[k] = v
|
|
|
|
server_params = StdioServerParameters(
|
|
command=command,
|
|
args=args,
|
|
env=full_env
|
|
)
|
|
|
|
read, write = await stdio_client(server_params).__aenter__()
|
|
session = ClientSession(read, write)
|
|
await session.__aenter__()
|
|
await session.initialize()
|
|
|
|
agent_id = agent_config.get("id", "")
|
|
self._sessions[agent_id] = {
|
|
"session": session,
|
|
"read": read,
|
|
"write": write,
|
|
}
|
|
|
|
return Connection(
|
|
agent_id=agent_id,
|
|
protocol=self.protocol_name,
|
|
endpoint=f"{command} {' '.join(args)}",
|
|
session=session,
|
|
metadata=agent_config
|
|
)
|
|
|
|
async def call(
|
|
self,
|
|
connection: Connection,
|
|
method: str,
|
|
params: dict,
|
|
timeout: float = 30.0
|
|
) -> dict:
|
|
"""调用 MCP 工具"""
|
|
session: ClientSession = connection.session
|
|
|
|
result = await asyncio.wait_for(
|
|
session.call_tool(method, params),
|
|
timeout=timeout
|
|
)
|
|
|
|
if hasattr(result, "content"):
|
|
content = result.content
|
|
if isinstance(content, list) and len(content) > 0:
|
|
first = content[0]
|
|
if hasattr(first, "text"):
|
|
return {"success": True, "result": first.text}
|
|
return {"success": True, "result": str(first)}
|
|
return {"success": True, "result": content}
|
|
|
|
return {"success": True, "result": result}
|
|
|
|
async def discover(self, capability: str = "") -> List[AgentInfo]:
|
|
"""发现 Agent(从本地配置)"""
|
|
agents_file = self.config_dir / "agents.yaml"
|
|
if not agents_file.exists():
|
|
return []
|
|
|
|
import yaml
|
|
with open(agents_file) as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
agents = []
|
|
for agent in config.get("agents", []):
|
|
if agent.get("protocol") != "mcp":
|
|
continue
|
|
|
|
if capability and capability.lower() not in agent.get("id", "").lower():
|
|
continue
|
|
|
|
agents.append(AgentInfo(
|
|
id=f"{agent['id']}@mcp",
|
|
protocol="mcp",
|
|
name=agent.get("name", agent["id"]),
|
|
endpoint=f"{agent.get('command', '')} {' '.join(agent.get('args', []))}",
|
|
metadata=agent
|
|
))
|
|
|
|
return agents
|
|
|
|
async def close(self, connection: Connection):
|
|
"""关闭连接"""
|
|
agent_id = connection.agent_id
|
|
if agent_id in self._sessions:
|
|
session_info = self._sessions.pop(agent_id)
|
|
session = session_info.get("session")
|
|
if session:
|
|
await session.__aexit__(None, None, None)
|
|
|
|
async def get_methods(self, connection: Connection) -> List[dict]:
|
|
"""获取 MCP Server 支持的工具"""
|
|
session: ClientSession = connection.session
|
|
|
|
result = await session.list_tools()
|
|
|
|
tools = []
|
|
if hasattr(result, "tools"):
|
|
for tool in result.tools:
|
|
tools.append({
|
|
"name": tool.name,
|
|
"description": getattr(tool, "description", ""),
|
|
"inputSchema": getattr(tool, "inputSchema", {}),
|
|
})
|
|
|
|
return tools
|
|
|
|
def validate_config(self, agent_config: dict) -> bool:
|
|
"""验证配置"""
|
|
return "command" in agent_config
|