xiaohei
a0ba54b039
- 修复MCP Server数据库连接和导入问题 - 创建Data Ingestion服务核心模块 - 配置阿里云镜像源提升构建速度 - 添加项目状态报告和文档 - 完善微服务架构基础设施
597 lines
21 KiB
Python
597 lines
21 KiB
Python
"""
|
|
MCP (Model Context Protocol) 协议处理器
|
|
实现MCP协议的核心功能，包括工具管理、资源管理和代理通信
|
|
"""
|
|
|
|
import json
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Any, Dict, List, Optional, Union
|
|
import asyncio
|
|
import logging
|
|
import traceback
|
|
|
|
import redis.asyncio as redis
|
|
import nats
|
|
import httpx
|
|
from schemas import (
|
|
MCPRequest, MCPResponse, MCPError,
|
|
ToolDefinition, ToolResult, ExecutionResult
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MCPProtocolHandler:
|
|
"""MCP协议处理器"""
|
|
|
|
def __init__(
|
|
self,
|
|
redis_client: redis.Redis,
|
|
nats_client: nats.NATS,
|
|
litellm_url: str = "http://litellm-gateway:4000"
|
|
):
|
|
self.redis = redis_client
|
|
self.nats = nats_client
|
|
self.litellm_url = litellm_url
|
|
self.http_client = httpx.AsyncClient(timeout=30.0)
|
|
|
|
# MCP协议版本
|
|
self.protocol_version = "2024-11-05"
|
|
|
|
# 支持的MCP方法
|
|
self.supported_methods = {
|
|
"initialize",
|
|
"tools/list",
|
|
"tools/call",
|
|
"resources/list",
|
|
"resources/read",
|
|
"prompts/list",
|
|
"prompts/get",
|
|
"completion/complete",
|
|
"logging/setLevel"
|
|
}
|
|
|
|
# 工具注册表
|
|
self._tools_registry: Dict[str, ToolDefinition] = {}
|
|
|
|
# 资源注册表
|
|
self._resources_registry: Dict[str, Dict[str, Any]] = {}
|
|
|
|
# 会话管理
|
|
self._sessions: Dict[str, Dict[str, Any]] = {}
|
|
|
|
async def execute_request(self, agent_id: str, request: MCPRequest) -> ExecutionResult:
|
|
"""执行MCP请求"""
|
|
execution_id = str(uuid.uuid4())
|
|
started_at = datetime.utcnow()
|
|
|
|
logger.info(f"开始执行MCP请求: {execution_id}, method: {request.method}")
|
|
|
|
try:
|
|
# 验证方法是否支持
|
|
if request.method not in self.supported_methods:
|
|
raise ValueError(f"不支持的MCP方法: {request.method}")
|
|
|
|
# 发布执行开始事件
|
|
await self._publish_execution_event(
|
|
"execution.started",
|
|
{
|
|
"execution_id": execution_id,
|
|
"agent_id": agent_id,
|
|
"method": request.method,
|
|
"timestamp": started_at.isoformat()
|
|
}
|
|
)
|
|
|
|
# 执行具体方法
|
|
result = await self._dispatch_method(agent_id, request)
|
|
|
|
completed_at = datetime.utcnow()
|
|
execution_time = (completed_at - started_at).total_seconds() * 1000
|
|
|
|
# 发布执行完成事件
|
|
await self._publish_execution_event(
|
|
"execution.completed",
|
|
{
|
|
"execution_id": execution_id,
|
|
"agent_id": agent_id,
|
|
"method": request.method,
|
|
"execution_time": execution_time,
|
|
"success": True,
|
|
"timestamp": completed_at.isoformat()
|
|
}
|
|
)
|
|
|
|
return ExecutionResult(
|
|
execution_id=execution_id,
|
|
success=True,
|
|
result=result,
|
|
execution_time=execution_time,
|
|
started_at=started_at,
|
|
completed_at=completed_at
|
|
)
|
|
|
|
except Exception as e:
|
|
completed_at = datetime.utcnow()
|
|
execution_time = (completed_at - started_at).total_seconds() * 1000
|
|
error_msg = str(e)
|
|
|
|
logger.error(f"MCP请求执行失败: {execution_id}, error: {error_msg}")
|
|
logger.error(traceback.format_exc())
|
|
|
|
# 发布执行失败事件
|
|
await self._publish_execution_event(
|
|
"execution.failed",
|
|
{
|
|
"execution_id": execution_id,
|
|
"agent_id": agent_id,
|
|
"method": request.method,
|
|
"execution_time": execution_time,
|
|
"error": error_msg,
|
|
"timestamp": completed_at.isoformat()
|
|
}
|
|
)
|
|
|
|
return ExecutionResult(
|
|
execution_id=execution_id,
|
|
success=False,
|
|
error=error_msg,
|
|
execution_time=execution_time,
|
|
started_at=started_at,
|
|
completed_at=completed_at
|
|
)
|
|
|
|
async def _dispatch_method(self, agent_id: str, request: MCPRequest) -> Any:
|
|
"""分发MCP方法调用"""
|
|
method = request.method
|
|
params = request.params or {}
|
|
|
|
if method == "initialize":
|
|
return await self._handle_initialize(params)
|
|
elif method == "tools/list":
|
|
return await self._handle_tools_list(agent_id, params)
|
|
elif method == "tools/call":
|
|
return await self._handle_tools_call(agent_id, params)
|
|
elif method == "resources/list":
|
|
return await self._handle_resources_list(agent_id, params)
|
|
elif method == "resources/read":
|
|
return await self._handle_resources_read(agent_id, params)
|
|
elif method == "prompts/list":
|
|
return await self._handle_prompts_list(agent_id, params)
|
|
elif method == "prompts/get":
|
|
return await self._handle_prompts_get(agent_id, params)
|
|
elif method == "completion/complete":
|
|
return await self._handle_completion_complete(agent_id, params)
|
|
elif method == "logging/setLevel":
|
|
return await self._handle_logging_set_level(params)
|
|
else:
|
|
raise ValueError(f"未实现的方法: {method}")
|
|
|
|
async def _handle_initialize(self, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理初始化请求"""
|
|
client_info = params.get("clientInfo", {})
|
|
protocol_version = params.get("protocolVersion")
|
|
|
|
logger.info(f"MCP客户端初始化: {client_info}")
|
|
|
|
return {
|
|
"protocolVersion": self.protocol_version,
|
|
"capabilities": {
|
|
"tools": {
|
|
"listChanged": True
|
|
},
|
|
"resources": {
|
|
"subscribe": True,
|
|
"listChanged": True
|
|
},
|
|
"prompts": {
|
|
"listChanged": True
|
|
},
|
|
"completion": {
|
|
"argument": True
|
|
},
|
|
"logging": {}
|
|
},
|
|
"serverInfo": {
|
|
"name": "taiji-AI-PAD MCP Server",
|
|
"version": "1.0.0"
|
|
}
|
|
}
|
|
|
|
async def _handle_tools_list(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理工具列表请求"""
|
|
try:
|
|
# 从Redis获取Agent的工具列表
|
|
agent_tools_key = f"agent:{agent_id}:tools"
|
|
tool_names = await self.redis.smembers(agent_tools_key)
|
|
|
|
tools = []
|
|
for tool_name in tool_names:
|
|
tool_info = await self._get_tool_info(tool_name)
|
|
if tool_info:
|
|
tools.append({
|
|
"name": tool_info["name"],
|
|
"description": tool_info["description"],
|
|
"inputSchema": tool_info.get("schema", {})
|
|
})
|
|
|
|
return {"tools": tools}
|
|
|
|
except Exception as e:
|
|
logger.error(f"获取工具列表失败: {e}")
|
|
return {"tools": []}
|
|
|
|
async def _handle_tools_call(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理工具调用请求"""
|
|
tool_name = params.get("name")
|
|
arguments = params.get("arguments", {})
|
|
|
|
if not tool_name:
|
|
raise ValueError("工具名称不能为空")
|
|
|
|
logger.info(f"调用工具: {tool_name}, arguments: {arguments}")
|
|
|
|
try:
|
|
# 验证Agent是否有权限使用该工具
|
|
agent_tools_key = f"agent:{agent_id}:tools"
|
|
if not await self.redis.sismember(agent_tools_key, tool_name):
|
|
raise ValueError(f"Agent {agent_id} 无权限使用工具 {tool_name}")
|
|
|
|
# 获取工具信息
|
|
tool_info = await self._get_tool_info(tool_name)
|
|
if not tool_info:
|
|
raise ValueError(f"工具 {tool_name} 不存在")
|
|
|
|
# 执行工具调用
|
|
result = await self._execute_tool(tool_name, tool_info, arguments)
|
|
|
|
return {
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": json.dumps(result.result) if result.success else f"错误: {result.error}"
|
|
}
|
|
],
|
|
"isError": not result.success
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"工具调用失败: {e}")
|
|
return {
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": f"工具调用失败: {str(e)}"
|
|
}
|
|
],
|
|
"isError": True
|
|
}
|
|
|
|
async def _handle_resources_list(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理资源列表请求"""
|
|
try:
|
|
# 从Redis获取Agent的资源列表
|
|
agent_resources_key = f"agent:{agent_id}:resources"
|
|
resource_names = await self.redis.smembers(agent_resources_key)
|
|
|
|
resources = []
|
|
for resource_name in resource_names:
|
|
resource_info = await self._get_resource_info(resource_name)
|
|
if resource_info:
|
|
resources.append({
|
|
"uri": resource_info["uri"],
|
|
"name": resource_info["name"],
|
|
"description": resource_info.get("description"),
|
|
"mimeType": resource_info.get("mimeType", "application/json")
|
|
})
|
|
|
|
return {"resources": resources}
|
|
|
|
except Exception as e:
|
|
logger.error(f"获取资源列表失败: {e}")
|
|
return {"resources": []}
|
|
|
|
async def _handle_resources_read(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理资源读取请求"""
|
|
uri = params.get("uri")
|
|
if not uri:
|
|
raise ValueError("资源URI不能为空")
|
|
|
|
try:
|
|
# 验证权限
|
|
agent_resources_key = f"agent:{agent_id}:resources"
|
|
# 这里应该根据URI找到资源名称
|
|
resource_name = uri.split("/")[-1] # 简化处理
|
|
|
|
if not await self.redis.sismember(agent_resources_key, resource_name):
|
|
raise ValueError(f"Agent {agent_id} 无权限访问资源 {uri}")
|
|
|
|
# 读取资源内容
|
|
content = await self._read_resource_content(uri)
|
|
|
|
return {
|
|
"contents": [
|
|
{
|
|
"uri": uri,
|
|
"mimeType": "application/json",
|
|
"text": json.dumps(content)
|
|
}
|
|
]
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"读取资源失败: {e}")
|
|
raise ValueError(f"无法读取资源 {uri}: {str(e)}")
|
|
|
|
async def _handle_prompts_list(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理提示词列表请求"""
|
|
# 获取Agent相关的提示词
|
|
prompts = [
|
|
{
|
|
"name": "system_prompt",
|
|
"description": "系统提示词",
|
|
"arguments": [
|
|
{
|
|
"name": "context",
|
|
"description": "上下文信息",
|
|
"required": False
|
|
}
|
|
]
|
|
}
|
|
]
|
|
|
|
return {"prompts": prompts}
|
|
|
|
async def _handle_prompts_get(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理获取提示词请求"""
|
|
name = params.get("name")
|
|
arguments = params.get("arguments", {})
|
|
|
|
if name == "system_prompt":
|
|
# 构建系统提示词
|
|
agent_info = await self._get_agent_info(agent_id)
|
|
prompt = f"""你是 {agent_info.get('name', 'AI助手')}。
|
|
角色定义: {agent_info.get('role', '通用助手')}
|
|
目标: {agent_info.get('goal', '帮助用户完成任务')}
|
|
|
|
可用工具: {', '.join(agent_info.get('tools', []))}
|
|
|
|
请根据用户的请求，选择合适的工具来完成任务。"""
|
|
|
|
return {
|
|
"description": "Agent系统提示词",
|
|
"messages": [
|
|
{
|
|
"role": "system",
|
|
"content": {
|
|
"type": "text",
|
|
"text": prompt
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
raise ValueError(f"未知的提示词: {name}")
|
|
|
|
async def _handle_completion_complete(self, agent_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理补全请求"""
|
|
ref = params.get("ref", {})
|
|
argument = params.get("argument", {})
|
|
|
|
# 根据参考信息生成补全建议
|
|
completions = []
|
|
|
|
if ref.get("type") == "tool":
|
|
tool_name = ref.get("name")
|
|
if tool_name:
|
|
tool_info = await self._get_tool_info(tool_name)
|
|
if tool_info and "schema" in tool_info:
|
|
# 基于工具schema生成参数建议
|
|
schema = tool_info["schema"]
|
|
properties = schema.get("properties", {})
|
|
for prop_name, prop_info in properties.items():
|
|
completions.append({
|
|
"type": "text",
|
|
"text": prop_name,
|
|
"insertText": f'"{prop_name}": ""'
|
|
})
|
|
|
|
return {"completion": {"values": completions}}
|
|
|
|
async def _handle_logging_set_level(self, params: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""处理设置日志级别请求"""
|
|
level = params.get("level", "info")
|
|
|
|
# 设置日志级别
|
|
numeric_level = getattr(logging, level.upper(), logging.INFO)
|
|
logging.getLogger().setLevel(numeric_level)
|
|
|
|
logger.info(f"日志级别已设置为: {level}")
|
|
|
|
return {"success": True}
|
|
|
|
async def _execute_tool(
|
|
self,
|
|
tool_name: str,
|
|
tool_info: Dict[str, Any],
|
|
arguments: Dict[str, Any]
|
|
) -> ToolResult:
|
|
"""执行工具调用"""
|
|
start_time = datetime.utcnow()
|
|
|
|
try:
|
|
# 根据工具类型执行不同的逻辑
|
|
category = tool_info.get("category", "api")
|
|
|
|
if category == "api":
|
|
result = await self._execute_api_tool(tool_info, arguments)
|
|
elif category == "function":
|
|
result = await self._execute_function_tool(tool_info, arguments)
|
|
elif category == "llm":
|
|
result = await self._execute_llm_tool(tool_info, arguments)
|
|
else:
|
|
raise ValueError(f"不支持的工具类型: {category}")
|
|
|
|
execution_time = (datetime.utcnow() - start_time).total_seconds() * 1000
|
|
|
|
return ToolResult(
|
|
success=True,
|
|
result=result,
|
|
execution_time=execution_time,
|
|
cost=tool_info.get("cost_per_call", 0.0)
|
|
)
|
|
|
|
except Exception as e:
|
|
execution_time = (datetime.utcnow() - start_time).total_seconds() * 1000
|
|
|
|
return ToolResult(
|
|
success=False,
|
|
error=str(e),
|
|
execution_time=execution_time,
|
|
cost=tool_info.get("cost_per_call", 0.0)
|
|
)
|
|
|
|
async def _execute_api_tool(self, tool_info: Dict[str, Any], arguments: Dict[str, Any]) -> Any:
|
|
"""执行API工具调用"""
|
|
endpoint = tool_info.get("endpoint")
|
|
method = tool_info.get("method", "POST")
|
|
headers = tool_info.get("headers", {})
|
|
timeout = tool_info.get("timeout", 30)
|
|
|
|
if not endpoint:
|
|
raise ValueError("API端点不能为空")
|
|
|
|
# 发送HTTP请求
|
|
response = await self.http_client.request(
|
|
method=method,
|
|
url=endpoint,
|
|
json=arguments,
|
|
headers=headers,
|
|
timeout=timeout
|
|
)
|
|
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def _execute_function_tool(self, tool_info: Dict[str, Any], arguments: Dict[str, Any]) -> Any:
|
|
"""执行函数工具调用"""
|
|
# 这里可以调用本地Python函数
|
|
# 为了安全起见，需要严格的沙箱机制
|
|
raise NotImplementedError("函数工具调用暂未实现")
|
|
|
|
async def _execute_llm_tool(self, tool_info: Dict[str, Any], arguments: Dict[str, Any]) -> Any:
|
|
"""执行LLM工具调用"""
|
|
# 调用LiteLLM网关
|
|
payload = {
|
|
"model": arguments.get("model", "gpt-3.5-turbo"),
|
|
"messages": arguments.get("messages", []),
|
|
"temperature": arguments.get("temperature", 0.7),
|
|
"max_tokens": arguments.get("max_tokens", 150)
|
|
}
|
|
|
|
response = await self.http_client.post(
|
|
f"{self.litellm_url}/chat/completions",
|
|
json=payload,
|
|
headers={"Authorization": "Bearer sk-taiji-master-key"}
|
|
)
|
|
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def _get_tool_info(self, tool_name: str) -> Optional[Dict[str, Any]]:
|
|
"""获取工具信息"""
|
|
tool_key = f"tool:{tool_name}"
|
|
tool_data = await self.redis.get(tool_key)
|
|
|
|
if tool_data:
|
|
return json.loads(tool_data)
|
|
return None
|
|
|
|
async def _get_resource_info(self, resource_name: str) -> Optional[Dict[str, Any]]:
|
|
"""获取资源信息"""
|
|
resource_key = f"resource:{resource_name}"
|
|
resource_data = await self.redis.get(resource_key)
|
|
|
|
if resource_data:
|
|
return json.loads(resource_data)
|
|
return None
|
|
|
|
async def _get_agent_info(self, agent_id: str) -> Dict[str, Any]:
|
|
"""获取Agent信息"""
|
|
agent_key = f"agent:{agent_id}"
|
|
agent_data = await self.redis.get(agent_key)
|
|
|
|
if agent_data:
|
|
return json.loads(agent_data)
|
|
return {}
|
|
|
|
async def _read_resource_content(self, uri: str) -> Any:
|
|
"""读取资源内容"""
|
|
# 这里可以根据URI类型读取不同的资源
|
|
# 例如：文件、数据库、API等
|
|
if uri.startswith("file://"):
|
|
# 读取文件
|
|
file_path = uri[7:] # 移除file://前缀
|
|
with open(file_path, 'r') as f:
|
|
return f.read()
|
|
elif uri.startswith("http://") or uri.startswith("https://"):
|
|
# 读取HTTP资源
|
|
response = await self.http_client.get(uri)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
else:
|
|
raise ValueError(f"不支持的资源类型: {uri}")
|
|
|
|
async def _publish_execution_event(self, event_type: str, data: Dict[str, Any]):
|
|
"""发布执行事件"""
|
|
try:
|
|
if self.nats:
|
|
await self.nats.publish(
|
|
f"mcp.{event_type}",
|
|
json.dumps(data).encode()
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"发布事件失败: {e}")
|
|
|
|
async def register_tool(self, tool_definition: ToolDefinition) -> bool:
|
|
"""注册工具"""
|
|
try:
|
|
tool_key = f"tool:{tool_definition.name}"
|
|
tool_data = tool_definition.dict()
|
|
|
|
await self.redis.setex(
|
|
tool_key,
|
|
3600, # 1小时过期
|
|
json.dumps(tool_data)
|
|
)
|
|
|
|
logger.info(f"工具注册成功: {tool_definition.name}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"工具注册失败: {e}")
|
|
return False
|
|
|
|
async def register_agent_tool(self, agent_id: str, tool_name: str) -> bool:
|
|
"""为Agent注册工具"""
|
|
try:
|
|
agent_tools_key = f"agent:{agent_id}:tools"
|
|
await self.redis.sadd(agent_tools_key, tool_name)
|
|
await self.redis.expire(agent_tools_key, 3600)
|
|
|
|
logger.info(f"Agent {agent_id} 工具注册成功: {tool_name}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Agent工具注册失败: {e}")
|
|
return False
|
|
|
|
async def close(self):
|
|
"""清理资源"""
|
|
try:
|
|
await self.http_client.aclose()
|
|
except Exception as e:
|
|
logger.error(f"资源清理失败: {e}")
|
|
|