MCP 与多 Agent 协作系统

# mcp_agent_bus.py - MCP Agent 协作总线实现

from typing import Dict, List, Any
import time
import uuid
from abc import ABC, abstractmethod

class Agent(ABC):
    """Agent 抽象类"""
    
    def __init__(self, agent_id: str, name: str, agent_type: str):
        self.agent_id = agent_id
        self.name = name
        self.agent_type = agent_type
        self.status = "idle"  # idle, working, error
        self.capabilities = []
        self.registered = False
        self.timestamp = time.time()
    
    @abstractmethod
    def process_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """处理消息"""
        pass
    
    @abstractmethod
    def get_capabilities(self) -> List[str]:
        """获取 Agent 能力"""
        pass
    
    def get_agent_info(self) -> Dict[str, Any]:
        """获取 Agent 信息"""
        return {
            "agent_id": self.agent_id,
            "name": self.name,
            "agent_type": self.agent_type,
            "status": self.status,
            "capabilities": self.get_capabilities(),
            "registered": self.registered,
            "timestamp": self.timestamp
        }

class SimpleAgent(Agent):
    """简单 Agent 实现"""
    
    def __init__(self, agent_id: str, name: str, agent_type: str):
        super().__init__(agent_id, name, agent_type)
        self.capabilities = ["task_processing", "message_handling"]
    
    def process_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """处理消息"""
        print(f"Agent {self.name}: Processing message: {message}")
        return {
            "status": "success",
            "result": f"Message processed by {self.name}",
            "agent_id": self.agent_id
        }
    
    def get_capabilities(self) -> List[str]:
        """获取 Agent 能力"""
        return self.capabilities

class AgentRegistry:
    """Agent 注册中心"""
    
    def __init__(self):
        self.agents = {}
    
    def register_agent(self, agent: Agent) -> bool:
        """注册 Agent"""
        if agent.agent_id not in self.agents:
            self.agents[agent.agent_id] = agent
            agent.registered = True
            agent.timestamp = time.time()
            return True
        return False
    
    def unregister_agent(self, agent_id: str) -> bool:
        """注销 Agent"""
        if agent_id in self.agents:
            agent = self.agents[agent_id]
            agent.registered = False
            del self.agents[agent_id]
            return True
        return False
    
    def get_agent(self, agent_id: str) -> Agent:
        """获取 Agent"""
        return self.agents.get(agent_id)
    
    def list_agents(self) -> List[Dict[str, Any]]:
        """列出所有 Agent"""
        return [agent.get_agent_info() for agent in self.agents.values()]
    
    def find_agents_by_capability(self, capability: str) -> List[Dict[str, Any]]:
        """根据能力查找 Agent"""
        agents = []
        for agent in self.agents.values():
            if capability in agent.get_capabilities():
                agents.append(agent.get_agent_info())
        return agents
    
    def update_agent_status(self, agent_id: str, status: str) -> bool:
        """更新 Agent 状态"""
        agent = self.get_agent(agent_id)
        if agent:
            agent.status = status
            agent.timestamp = time.time()
            return True
        return False

class MessageRouter:
    """消息路由器"""
    
    def __init__(self, agent_registry: AgentRegistry):
        self.agent_registry = agent_registry
    
    def route_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """路由消息"""
        # 解析消息
        recipient_id = message.get("recipient_id")
        message_type = message.get("message_type")
        message_content = message.get("content")
        
        # 查找接收方 Agent
        agent = self.agent_registry.get_agent(recipient_id)
        if not agent:
            return {
                "status": "error",
                "message": f"Agent {recipient_id} not found"
            }
        
        # 处理消息
        try:
            result = agent.process_message(message)
            return {
                "status": "success",
                "result": result,
                "message": f"Message routed to {recipient_id}"
            }
        except Exception as e:
            return {
                "status": "error",
                "message": str(e)
            }
    
    def broadcast_message(self, message: Dict[str, Any]) -> List[Dict[str, Any]]:
        """广播消息"""
        results = []
        for agent in self.agent_registry.agents.values():
            try:
                result = agent.process_message(message)
                results.append({
                    "agent_id": agent.agent_id,
                    "status": "success",
                    "result": result
                })
            except Exception as e:
                results.append({
                    "agent_id": agent.agent_id,
                    "status": "error",
                    "message": str(e)
                })
        return results

class MCP_Agent_Collaboration_Bus:
    """MCP Agent 协作总线"""
    
    def __init__(self):
        self.agent_registry = AgentRegistry()
        self.message_router = MessageRouter(self.agent_registry)
        self.mcp_tools = {
            "agent_register": self._agent_register_tool,
            "agent_unregister": self._agent_unregister_tool,
            "message_send": self._message_send_tool,
            "message_broadcast": self._message_broadcast_tool,
            "agent_list": self._agent_list_tool
        }
    
    def get_mcp_tools(self) -> List[Dict[str, Any]]:
        """获取 MCP 工具定义"""
        return [
            {
                "id": "agent_register",
                "name": "Agent Register Tool",
                "description": "Register an agent with the collaboration bus",
                "parameters": {
                    "agent_id": {
                        "type": "string",
                        "description": "Agent ID",
                        "required": True
                    },
                    "name": {
                        "type": "string",
                        "description": "Agent name",
                        "required": True
                    },
                    "agent_type": {
                        "type": "string",
                        "description": "Agent type",
                        "required": True
                    }
                },
                "returns": {
                    "status": "string",
                    "message": "string"
                }
            },
            {
                "id": "agent_unregister",
                "name": "Agent Unregister Tool",
                "description": "Unregister an agent from the collaboration bus",
                "parameters": {
                    "agent_id": {
                        "type": "string",
                        "description": "Agent ID",
                        "required": True
                    }
                },
                "returns": {
                    "status": "string",
                    "message": "string"
                }
            },
            {
                "id": "message_send",
                "name": "Message Send Tool",
                "description": "Send a message to an agent",
                "parameters": {
                    "recipient_id": {
                        "type": "string",
                        "description": "Recipient agent ID",
                        "required": True
                    },
                    "message_type": {
                        "type": "string",
                        "description": "Message type",
                        "required": True
                    },
                    "content": {
                        "type": "object",
                        "description": "Message content",
                        "required": True
                    }
                },
                "returns": {
                    "status": "string",
                    "result": "object",
                    "message": "string"
                }
            },
            {
                "id": "message_broadcast",
                "name": "Message Broadcast Tool",
                "description": "Broadcast a message to all agents",
                "parameters": {
                    "message_type": {
                        "type": "string",
                        "description": "Message type",
                        "required": True
                    },
                    "content": {
                        "type": "object",
                        "description": "Message content",
                        "required": True
                    }
                },
                "returns": {
                    "status": "string",
                    "results": "array",
                    "message": "string"
                }
            },
            {
                "id": "agent_list",
                "name": "Agent List Tool",
                "description": "List all registered agents",
                "parameters": {},
                "returns": {
                    "status": "string",
                    "agents": "array",
                    "message": "string"
                }
            }
        ]
    
    def call_tool(self, tool_id: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """调用 MCP 工具"""
        if tool_id in self.mcp_tools:
            return self.mcp_tools[tool_id](parameters)
        return {
            "status": "error",
            "message": f"Tool {tool_id} not found"
        }
    
    def _agent_register_tool(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """Agent 注册工具实现"""
        agent_id = parameters.get("agent_id")
        name = parameters.get("name")
        agent_type = parameters.get("agent_type")
        
        # 创建 Agent 实例
        agent = SimpleAgent(agent_id, name, agent_type)
        
        # 注册 Agent
        success = self.agent_registry.register_agent(agent)
        
        if success:
            return {
                "status": "success",
                "message": f"Agent {name} registered successfully"
            }
        else:
            return {
                "status": "error",
                "message": f"Agent {name} already registered"
            }
    
    def _agent_unregister_tool(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """Agent 注销工具实现"""
        agent_id = parameters.get("agent_id")
        
        # 注销 Agent
        success = self.agent_registry.unregister_agent(agent_id)
        
        if success:
            return {
                "status": "success",
                "message": f"Agent {agent_id} unregistered successfully"
            }
        else:
            return {
                "status": "error",
                "message": f"Agent {agent_id} not found"
            }
    
    def _message_send_tool(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """消息发送工具实现"""
        recipient_id = parameters.get("recipient_id")
        message_type = parameters.get("message_type")
        content = parameters.get("content")
        
        # 构造消息
        message = {
            "sender_id": "mcp_bus",
            "recipient_id": recipient_id,
            "message_type": message_type,
            "content": content,
            "timestamp": time.time()
        }
        
        # 路由消息
        result = self.message_router.route_message(message)
        return result
    
    def _message_broadcast_tool(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """消息广播工具实现"""
        message_type = parameters.get("message_type")
        content = parameters.get("content")
        
        # 构造消息
        message = {
            "sender_id": "mcp_bus",
            "message_type": message_type,
            "content": content,
            "timestamp": time.time()
        }
        
        # 广播消息
        results = self.message_router.broadcast_message(message)
        
        return {
            "status": "success",
            "results": results,
            "message": f"Message broadcast to {len(results)} agents"
        }
    
    def _agent_list_tool(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """Agent 列表工具实现"""
        agents = self.agent_registry.list_agents()
        return {
            "status": "success",
            "agents": agents,
            "message": f"Found {len(agents)} registered agents"
        }

# 使用示例
if __name__ == "__main__":
    # 初始化 MCP Agent 协作总线
    bus = MCP_Agent_Collaboration_Bus()
    
    # 获取 MCP 工具定义
    tools = bus.get_mcp_tools()
    print("MCP 工具定义:")
    for tool in tools:
        print(f"- {tool['id']}: {tool['name']}")
    
    # 注册 Agent 1
    result = bus.call_tool("agent_register", {
        "agent_id": "agent-001",
        "name": "Agent 1",
        "agent_type": "general"
    })
    print(f"\n注册 Agent 1:")
    print(f"Status: {result['status']}")
    print(f"Message: {result['message']}")
    
    # 注册 Agent 2
    result = bus.call_tool("agent_register", {
        "agent_id": "agent-002",
        "name": "Agent 2",
        "agent_type": "general"
    })
    print(f"\n注册 Agent 2:")
    print(f"Status: {result['status']}")
    print(f"Message: {result['message']}")
    
    # 列出注册的 Agent
    result = bus.call_tool("agent_list", {})
    print(f"\n列出注册的 Agent:")
    print(f"Status: {result['status']}")
    print(f"Agents: {len(result['agents'])}")
    for agent in result['agents']:
        print(f"  - {agent['agent_id']}: {agent['name']} ({agent['agent_type']})")
    
    # 发送消息给 Agent 1
    result = bus.call_tool("message_send", {
        "recipient_id": "agent-001",
        "message_type": "task",
        "content": {
            "task_id": "task-001",
            "task_type": "data_processing",
            "data": "sample data"
        }
    })
    print(f"\n发送消息给 Agent 1:")
    print(f"Status: {result['status']}")
    print(f"Result: {result['result']}")
    
    # 广播消息
    result = bus.call_tool("message_broadcast", {
        "message_type": "system",
        "content": {
            "message": "System update notification",
            "version": "1.0.0"
        }
    })
    print(f"\n广播消息:")
    print(f"Status: {result['status']}")
    print(f"Results: {len(result['results'])} agents processed the message")

# 安装依赖
pip install requests

# 运行示例代码
python mcp_agent_collaboration_bus.py

MCP 工具定义:
- agent_register: Agent Register Tool
- agent_unregister: Agent Unregister Tool
- message_send: Message Send Tool
- message_broadcast: Message Broadcast Tool
- agent_list: Agent List Tool

注册 Agent 1:
Status: success
Message: Agent Agent 1 registered successfully

注册 Agent 2:
Status: success
Message: Agent Agent 2 registered successfully

列出注册的 Agent:
Status: success
Agents: 2
  - agent-001: Agent 1 (general)
  - agent-002: Agent 2 (general)

发送消息给 Agent 1:
Agent Agent 1: Processing message: {'sender_id': 'mcp_bus', 'recipient_id': 'agent-001', 'message_type': 'task', 'content': {'task_id': 'task-001', 'task_type': 'data_processing', 'data': 'sample data'}, 'timestamp': 1736425200.0}
Status: success
Result: {'status': 'success', 'result': 'Message processed by Agent 1', 'agent_id': 'agent-001'}

广播消息:
Agent Agent 1: Processing message: {'sender_id': 'mcp_bus', 'message_type': 'system', 'content': {'message': 'System update notification', 'version': '1.0.0'}, 'timestamp': 1736425200.0}
Agent Agent 2: Processing message: {'sender_id': 'mcp_bus', 'message_type': 'system', 'content': {'message': 'System update notification', 'version': '1.0.0'}, 'timestamp': 1736425200.0}
Status: success
Results: 2 agents processed the message

# distributed_agent_system.py - 分布式 Agent 系统设计框架实现

from typing import Dict, List, Any
import time
import random

class LoadBalancer:
    """负载均衡器"""
    
    def __init__(self):
        self.agents = []
        self.loads = {}  # agent_id -> load
    
    def register_agent(self, agent_id: str):
        """注册 Agent"""
        self.agents.append(agent_id)
        self.loads[agent_id] = 0
    
    def unregister_agent(self, agent_id: str):
        """注销 Agent"""
        if agent_id in self.agents:
            self.agents.remove(agent_id)
            del self.loads[agent_id]
    
    def update_load(self, agent_id: str, load: int):
        """更新 Agent 负载"""
        if agent_id in self.loads:
            self.loads[agent_id] = load
    
    def select_agent(self) -> str:
        """选择负载最低的 Agent"""
        if not self.agents:
            raise ValueError("No agents available")
        
        # 选择负载最低的 Agent
        min_load = min(self.loads.values())
        candidates = [agent_id for agent_id, load in self.loads.items() if load == min_load]
        
        # 如果有多个负载相同的 Agent，随机选择一个
        return random.choice(candidates)
    
    def get_load_info(self) -> Dict[str, Any]:
        """获取负载信息"""
        return {
            "agents": self.agents,
            "loads": self.loads,
            "total_agents": len(self.agents)
        }

class FaultToleranceManager:
    """容错管理器"""
    
    def __init__(self):
        self.failed_agents = []
        self.recovery_attempts = {}  # agent_id -> attempts
    
    def mark_failed(self, agent_id: str):
        """标记 Agent 失败"""
        if agent_id not in self.failed_agents:
            self.failed_agents.append(agent_id)
            self.recovery_attempts[agent_id] = 0
    
    def mark_recovered(self, agent_id: str):
        """标记 Agent 恢复"""
        if agent_id in self.failed_agents:
            self.failed_agents.remove(agent_id)
            del self.recovery_attempts[agent_id]
    
    def is_failed(self, agent_id: str) -> bool:
        """检查 Agent 是否失败"""
        return agent_id in self.failed_agents
    
    def attempt_recovery(self, agent_id: str) -> bool:
        """尝试恢复 Agent"""
        if agent_id in self.failed_agents:
            self.recovery_attempts[agent_id] += 1
            # 简单的恢复逻辑：如果尝试次数少于 3 次，返回可以恢复
            if self.recovery_attempts[agent_id] < 3:
                return True
        return False
    
    def get_fault_info(self) -> Dict[str, Any]:
        """获取故障信息"""
        return {
            "failed_agents": self.failed_agents,
            "recovery_attempts": self.recovery_attempts,
            "total_failed": len(self.failed_agents)
        }

class DistributedAgentSystem:
    """分布式 Agent 系统"""
    
    def __init__(self):
        self.load_balancer = LoadBalancer()
        self.fault_manager = FaultToleranceManager()
        self.agent_nodes = {}
    
    def add_agent_node(self, node_id: str, agents: List[Dict[str, Any]]):
        """添加 Agent 节点"""
        self.agent_nodes[node_id] = agents
        # 注册 Agent 到负载均衡器
        for agent in agents:
            self.load_balancer.register_agent(agent["agent_id"])
    
    def remove_agent_node(self, node_id: str):
        """移除 Agent 节点"""
        if node_id in self.agent_nodes:
            agents = self.agent_nodes[node_id]
            # 从负载均衡器注销 Agent
            for agent in agents:
                self.load_balancer.unregister_agent(agent["agent_id"])
            del self.agent_nodes[node_id]
    
    def route_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """路由请求到合适的 Agent"""
        # 选择负载最低的 Agent
        agent_id = self.load_balancer.select_agent()
        
        # 检查 Agent 是否失败
        if self.fault_manager.is_failed(agent_id):
            # 尝试恢复 Agent
            if self.fault_manager.attempt_recovery(agent_id):
                print(f"Attempting to recover agent {agent_id}")
            else:
                # 选择另一个 Agent
                agent_id = self.load_balancer.select_agent()
        
        # 模拟处理请求
        print(f"Routing request {request['request_id']} to agent {agent_id}")
        
        # 更新负载
        self.load_balancer.update_load(agent_id, self.load_balancer.loads[agent_id] + 1)
        
        return {
            "status": "success",
            "agent_id": agent_id,
            "request_id": request["request_id"],
            "processed_at": time.time()
        }
    
    def get_system_status(self) -> Dict[str, Any]:
        """获取系统状态"""
        return {
            "agent_nodes": self.agent_nodes,
            "load_balancer": self.load_balancer.get_load_info(),
            "fault_tolerance": self.fault_manager.get_fault_info(),
            "total_nodes": len(self.agent_nodes),
            "total_agents": sum(len(agents) for agents in self.agent_nodes.values())
        }

# 使用示例
if __name__ == "__main__":
    # 初始化分布式 Agent 系统
    system = DistributedAgentSystem()
    
    # 添加 Agent 节点 1
    system.add_agent_node("node-001", [
        {"agent_id": "agent-001", "name": "Agent 1", "type": "general"},
        {"agent_id": "agent-002", "name": "Agent 2", "type": "general"}
    ])
    
    # 添加 Agent 节点 2
    system.add_agent_node("node-002", [
        {"agent_id": "agent-003", "name": "Agent 3", "type": "general"},
        {"agent_id": "agent-004", "name": "Agent 4", "type": "general"}
    ])
    
    # 添加 Agent 节点 3
    system.add_agent_node("node-003", [
        {"agent_id": "agent-005", "name": "Agent 5", "type": "general"},
        {"agent_id": "agent-006", "name": "Agent 6", "type": "general"}
    ])
    
    # 查看系统状态
    status = system.get_system_status()
    print("系统状态:")
    print(f"总节点数: {status['total_nodes']}")
    print(f"总 Agent 数: {status['total_agents']}")
    print(f"负载信息: {status['load_balancer']['loads']}")
    
    # 模拟路由请求
    for i in range(5):
        request = {
            "request_id": f"request-{i+1}",
            "request_type": "data_processing",
            "data": f"sample data {i+1}"
        }
        result = system.route_request(request)
        print(f"\n请求 {request['request_id']} 路由结果:")
        print(f"状态: {result['status']}")
        print(f"处理 Agent: {result['agent_id']}")
    
    # 标记 Agent 失败
    system.fault_manager.mark_failed("agent-001")
    print(f"\n标记 Agent 001 失败")
    
    # 再次路由请求
    request = {
        "request_id": "request-6",
        "request_type": "data_processing",
        "data": "sample data 6"
    }
    result = system.route_request(request)
    print(f"\n请求 {request['request_id']} 路由结果:")
    print(f"状态: {result['status']}")
    print(f"处理 Agent: {result['agent_id']}")
    
    # 查看更新后的系统状态
    status = system.get_system_status()
    print(f"\n更新后的系统状态:")
    print(f"失败的 Agent: {status['fault_tolerance']['failed_agents']}")
    print(f"负载信息: {status['load_balancer']['loads']}")

# 安装依赖
pip install requests

# 运行示例代码
python distributed_agent_system.py

系统状态:
总节点数: 3
总 Agent 数: 6
负载信息: {'agent-001': 0, 'agent-002': 0, 'agent-003': 0, 'agent-004': 0, 'agent-005': 0, 'agent-006': 0}

请求 request-1 路由结果:
Routing request request-1 to agent agent-003
状态: success
处理 Agent: agent-003

请求 request-2 路由结果:
Routing request request-2 to agent agent-001
状态: success
处理 Agent: agent-001

请求 request-3 路由结果:
Routing request request-3 to agent agent-002
状态: success
处理 Agent: agent-002

请求 request-4 路由结果:
Routing request request-4 to agent agent-004
状态: success
处理 Agent: agent-004

请求 request-5 路由结果:
Routing request request-5 to agent agent-005
状态: success
处理 Agent: agent-005

标记 Agent 001 失败

请求 request-6 路由结果:
Attempting to recover agent agent-001
Routing request request-6 to agent agent-006
状态: success
处理 Agent: agent-006

更新后的系统状态:
失败的 Agent: ['agent-001']
负载信息: {'agent-001': 1, 'agent-002': 1, 'agent-003': 1, 'agent-004': 1, 'agent-005': 1, 'agent-006': 1}