AutoGen/LangGraph多Agent：状态机与路由

from typing import Dict, List, Any, Optionalfrom enum import Enumfrom dataclasses import dataclass, field  class AgentRole(Enum):    """Agent角色枚举"""    RED_TEAM = "red_team"    BLUE_TEAM = "blue_team"    JUDGE = "judge"  @dataclassclass Message:    """Agent间通信消息"""    sender: str    receiver: str    content: str    msg_type: str = "task"    ttl: int = 10  @dataclassclass AgentState:    """Agent状态"""    agent_id: str    role: AgentRole    context: Dict[str, Any] = field(default_factory=dict)    history: List[Message] = field(default_factory=list)  class MessageRouter:    """消息路由器：支持点对点、广播和发布订阅"""     def __init__(self):        self.agents: Dict[str, AgentState] = {}        self.subscriptions: Dict[str, List[str]] = {}     def register(self, agent: AgentState):        """注册Agent"""        self.agents[agent.agent_id] = agent     def subscribe(self, topic: str, agent_id: str):        """订阅主题"""        if topic not in self.subscriptions:            self.subscriptions[topic] = []        self.subscriptions[topic].append(agent_id)     def route_point_to_point(self, msg: Message) -> bool:        """点对点路由"""        if msg.receiver in self.agents:            self.agents[msg.receiver].history.append(msg)            return True        return False     def route_broadcast(self, msg: Message) -> int:        """广播路由：发送给所有Agent"""        count = 0        for agent_id in self.agents:            if agent_id != msg.sender:                clone = Message(                    sender=msg.sender, receiver=agent_id,                    content=msg.content, msg_type=msg.msg_type                )                self.agents[agent_id].history.append(clone)                count += 1        return count     def route_publish(self, topic: str, msg: Message) -> int:        """发布订阅路由"""        subscribers = self.subscriptions.get(topic, [])        count = 0        for sub_id in subscribers:            clone = Message(                sender=msg.sender, receiver=sub_id,                content=msg.content, msg_type=msg.msg_type            )            self.agents[sub_id].history.append(clone)            count += 1        return count  class StateMachine:    """状态机：驱动Agent流程编排"""     def __init__(self):        self.nodes: Dict[str, callable] = {}        self.edges: Dict[str, List[tuple]] = {}        self.checkpoints: List[Dict] = []     def add_node(self, name: str, handler: callable):        """添加节点"""        self.nodes[name] = handler     def add_edge(self, src: str, dst: str,                 condition: callable = None):        """添加条件边"""        if src not in self.edges:            self.edges[src] = []        self.edges[src].append((dst, condition))     def save_checkpoint(self, state: Dict):        """保存检查点"""        self.checkpoints.append(state.copy())     def run(self, initial_state: Dict) -> Dict:        """执行状态机流程"""        state = initial_state.copy()        current = "start"        while current and current in self.nodes:            self.save_checkpoint(state)            state = self.nodes[current](state)            next_node = self._get_next(current, state)            current = next_node        return state     def _get_next(self, node: str,                  state: Dict) -> Optional[str]:        """根据条件边获取下一节点"""        for dst, cond in self.edges.get(node, []):            if cond is None or cond(state):                return dst        return None  # 使用示例print("=== 多Agent消息路由演示 ===")router = MessageRouter()red = AgentState("red_1", AgentRole.RED_TEAM)blue = AgentState("blue_1", AgentRole.BLUE_TEAM)judge = AgentState("judge_1", AgentRole.JUDGE)router.register(red)router.register(blue)router.register(judge)router.subscribe("attack_result", "blue_1")router.subscribe("attack_result", "judge_1") msg = Message("red_1", "blue_1", "对抗样本#42")router.route_point_to_point(msg)print(f"点对点: {blue.history[-1].content}") broadcast = Message("judge_1", "*", "全局通知: 第3轮开始")n = router.route_broadcast(broadcast)print(f"广播: 发送给 {n} 个Agent") pub = Message("red_1", "", "攻击成功")n = router.route_publish("attack_result", pub)print(f"发布订阅: {n} 个订阅者收到")