forked from xiaohei/taiji-AI-PAD
812 lines
29 KiB
Python
812 lines
29 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
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
|
|
from functools import wraps
|
|
|
|
import redis.asyncio as redis
|
|
import nats
|
|
import httpx
|
|
from schemas import (
|
|
MCPRequest, MCPResponse, MCPError,
|
|
ToolDefinition, ToolResult, ExecutionResult
|
|
)
|
|
from function_registry import get_function_registry
|
|
from sandbox_executor import get_sandbox_executor
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MCPError(Exception):
|
|
"""MCP协议错误"""
|
|
def __init__(self, code: int, message: str, data: Optional[Any] = None):
|
|
self.code = code
|
|
self.message = message
|
|
self.data = data
|
|
super().__init__(message)
|
|
|
|
|
|
class MCPRequestError(MCPError):
|
|
"""MCP请求错误"""
|
|
pass
|
|
|
|
|
|
class MCPResourceNotFound(MCPError):
|
|
"""MCP资源未找到错误"""
|
|
def __init__(self, resource: str):
|
|
super().__init__(-32002, f"Resource not found: {resource}", {"resource": resource})
|
|
|
|
|
|
class MCPToolNotFound(MCPError):
|
|
"""MCP工具未找到错误"""
|
|
def __init__(self, tool: str):
|
|
super().__init__(-32003, f"Tool not found: {tool}", {"tool": tool})
|
|
|
|
|
|
def with_retry(max_retries: int = 3, backoff_factor: float = 1.5):
|
|
"""重试装饰器"""
|
|
def decorator(func):
|
|
@wraps(func)
|
|
async def wrapper(*args, **kwargs):
|
|
last_exception = None
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return await func(*args, **kwargs)
|
|
except (asyncio.TimeoutError, httpx.TimeoutException, redis.ConnectionError) as e:
|
|
last_exception = e
|
|
if attempt < max_retries - 1:
|
|
wait_time = backoff_factor ** attempt
|
|
logger.warning(
|
|
f"Attempt {attempt + 1}/{max_retries} failed for {func.__name__}, "
|
|
f"retrying in {wait_time:.1f}s: {e}"
|
|
)
|
|
await asyncio.sleep(wait_time)
|
|
else:
|
|
logger.error(f"All {max_retries} attempts failed for {func.__name__}")
|
|
except Exception as e:
|
|
# Don't retry on other exceptions
|
|
raise
|
|
raise last_exception
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
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"
|
|
|
|
# 函数注册表和沙箱执行器
|
|
self.function_registry = get_function_registry()
|
|
self.sandbox_executor = get_sandbox_executor(
|
|
timeout=5.0,
|
|
max_memory_mb=100,
|
|
allow_network=False,
|
|
allow_filesystem=False
|
|
)
|
|
|
|
# 支持的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}",
|
|
extra={
|
|
"execution_id": execution_id,
|
|
"agent_id": agent_id,
|
|
"method": request.method,
|
|
"request_id": str(request.id)
|
|
}
|
|
)
|
|
|
|
try:
|
|
# 验证请求
|
|
self._validate_request(request)
|
|
|
|
# 发布执行开始事件
|
|
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 MCPError as e:
|
|
# MCP协议特定错误
|
|
completed_at = datetime.utcnow()
|
|
execution_time = (completed_at - started_at).total_seconds() * 1000
|
|
|
|
logger.error(
|
|
f"MCP协议错误: {e.message}",
|
|
extra={
|
|
"execution_id": execution_id,
|
|
"error_code": e.code,
|
|
"error_data": e.data
|
|
}
|
|
)
|
|
|
|
await self._publish_execution_event(
|
|
"execution.failed",
|
|
{
|
|
"execution_id": execution_id,
|
|
"agent_id": agent_id,
|
|
"method": request.method,
|
|
"execution_time": execution_time,
|
|
"error": e.message,
|
|
"error_code": e.code,
|
|
"timestamp": completed_at.isoformat()
|
|
}
|
|
)
|
|
|
|
return ExecutionResult(
|
|
execution_id=execution_id,
|
|
success=False,
|
|
error=f"[{e.code}] {e.message}",
|
|
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)
|
|
error_trace = traceback.format_exc()
|
|
|
|
logger.error(
|
|
f"MCP请求执行失败: {execution_id}",
|
|
extra={
|
|
"execution_id": execution_id,
|
|
"error": error_msg,
|
|
"traceback": error_trace
|
|
}
|
|
)
|
|
|
|
# 发布执行失败事件
|
|
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
|
|
)
|
|
|
|
def _validate_request(self, request: MCPRequest):
|
|
"""验证MCP请求"""
|
|
# 验证方法是否支持
|
|
if request.method not in self.supported_methods:
|
|
raise MCPRequestError(
|
|
-32601,
|
|
f"Method not found: {request.method}",
|
|
{"method": request.method, "supported_methods": list(self.supported_methods)}
|
|
)
|
|
|
|
# 验证请求ID
|
|
if not request.id:
|
|
raise MCPRequestError(-32600, "Invalid request: missing id")
|
|
|
|
# 验证JSONRPC版本
|
|
if request.jsonrpc != "2.0":
|
|
raise MCPRequestError(
|
|
-32600,
|
|
f"Invalid request: unsupported jsonrpc version {request.jsonrpc}"
|
|
)
|
|
|
|
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"
|
|
}
|
|
}
|
|
|
|
@with_retry(max_retries=3, backoff_factor=1.5)
|
|
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 MCPToolNotFound(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 MCPError:
|
|
# Re-raise MCP errors
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"工具调用失败: {e}", exc_info=True)
|
|
raise MCPError(-32000, f"Tool execution failed: {str(e)}", {"tool": tool_name})
|
|
|
|
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助手')}。\n"
|
|
f"角色定义: {agent_info.get('role', '通用助手')}\n"
|
|
f"目标: {agent_info.get('goal', '帮助用户完成任务')}\n\n"
|
|
f"可用工具: {', '.join(agent_info.get('tools', []))}\n\n"
|
|
"请根据用户的请求，选择合适的工具来完成任务。"
|
|
)
|
|
|
|
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:
|
|
"""
|
|
执行函数工具调用
|
|
|
|
安全机制:
|
|
1. 只允许执行注册表中的函数
|
|
2. 在沙箱环境中执行
|
|
3. 参数验证和资源限制
|
|
4. 超时控制
|
|
"""
|
|
function_name = tool_info.get("name") or tool_info.get("function_name")
|
|
|
|
if not function_name:
|
|
raise ValueError("函数名称不能为空")
|
|
|
|
# 检查函数是否在注册表中
|
|
function_info = self.function_registry.get(function_name)
|
|
if not function_info:
|
|
raise ValueError(
|
|
f"函数 {function_name} 未在注册表中。"
|
|
f"可用函数: {', '.join(self.function_registry.list_all()[:10])}"
|
|
)
|
|
|
|
# 验证参数
|
|
self._validate_function_arguments(function_info, arguments)
|
|
|
|
# 在沙箱中执行函数
|
|
try:
|
|
result = await self.sandbox_executor.execute(
|
|
func=function_info["func"],
|
|
arguments=arguments,
|
|
function_name=function_name
|
|
)
|
|
|
|
logger.info(f"函数 {function_name} 执行成功")
|
|
return result
|
|
|
|
except TimeoutError as e:
|
|
logger.error(f"函数 {function_name} 执行超时: {e}")
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"函数 {function_name} 执行失败: {e}")
|
|
raise RuntimeError(f"函数执行失败: {str(e)}")
|
|
|
|
def _validate_function_arguments(
|
|
self,
|
|
function_info: Dict[str, Any],
|
|
arguments: Dict[str, Any]
|
|
):
|
|
"""验证函数参数"""
|
|
expected_params = {
|
|
param["name"]: param
|
|
for param in function_info.get("parameters", [])
|
|
}
|
|
|
|
# 检查必需参数
|
|
for param_name, param_info in expected_params.items():
|
|
if param_info.get("required", True):
|
|
if param_name not in arguments:
|
|
raise ValueError(f"缺少必需参数: {param_name}")
|
|
|
|
# 检查参数类型（基本验证）
|
|
for param_name, param_value in arguments.items():
|
|
if param_name not in expected_params:
|
|
logger.warning(f"未知参数: {param_name}，将被忽略")
|
|
continue
|
|
|
|
param_info = expected_params[param_name]
|
|
expected_type = param_info.get("type", "any")
|
|
|
|
# 基本类型检查
|
|
if expected_type == "number" and not isinstance(param_value, (int, float)):
|
|
try:
|
|
float(param_value) # 尝试转换
|
|
except (ValueError, TypeError):
|
|
raise ValueError(
|
|
f"参数 {param_name} 应该是数字类型，但得到 {type(param_value).__name__}"
|
|
)
|
|
elif expected_type == "string" and not isinstance(param_value, str):
|
|
raise ValueError(
|
|
f"参数 {param_name} 应该是字符串类型，但得到 {type(param_value).__name__}"
|
|
)
|
|
elif expected_type == "object" and not isinstance(param_value, (dict, list)):
|
|
raise ValueError(
|
|
f"参数 {param_name} 应该是对象类型，但得到 {type(param_value).__name__}"
|
|
)
|
|
|
|
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}")
|
|
|