Agent协作与Swarm

from dataclasses import dataclass
from typing import List, Dict
import time

@dataclass
class Message:
    """Agent消息"""
    from_agent: str
    to_agent: str
    content: str
    message_type: str = "request"

class Agent:
    def __init__(self, name: str, role: str):
        self.name = name
        self.role = role
        self.message_queue: List[Message] = []

    def send_message(self, to_agent: 'Agent', content: str):
        """发送消息"""
        msg = Message(
            from_agent=self.name,
            to_agent=to_agent.name,
            content=content
        )
        to_agent.message_queue.append(msg)
        print(f"{self.name} → {to_agent.name}: {content[:30]}...")

    def process_messages(self) -> List[Message]:
        """处理消息"""
        responses = []
        while self.message_queue:
            msg = self.message_queue.pop(0)
            response = self.handle_message(msg)
            if response:
                responses.append(response)
        return responses

    def handle_message(self, msg: Message) -> Message | None:
        """处理消息（子类实现）"""
        print(f"{self.name} 收到: {msg.content[:30]}...")
        return None

class ResearchAgent(Agent):
    def handle_message(self, msg: Message) -> Message:
        """研究员Agent：收集信息"""
        if "research" in msg.content.lower():
            result = f"已完成关于{msg.content}的研究"
            return Message(
                from_agent=self.name,
                to_agent=msg.from_agent,
                content=result
            )
        return super().handle_message(msg)

class WriterAgent(Agent):
    def handle_message(self, msg: Message) -> Message:
        """写作者Agent：生成文档"""
        if "write" in msg.content.lower():
            result = f"已根据{msg.content}生成文档"
            return Message(
                from_agent=self.name,
                to_agent=msg.from_agent,
                content=result
            )
        return super().handle_message(msg)

class ManagerAgent(Agent):
    """管理者Agent：协调其他Agent"""
    def __init__(self, name: str, agents: List[Agent]):
        super().__init__(name, "manager")
        self.agents = agents

    def coordinate_task(self, task: str):
        """协调任务"""
        print(f"
【{self.name}】分配任务: {task}")

        # 步骤1：指派研究Agent
        research_agent = next(a for a in self.agents if a.role == "researcher")
        self.send_message(research_agent, f"请研究：{task}")

        # 模拟执行
        research_agent.process_messages()

        # 步骤2：指派写作者Agent
        writer_agent = next(a for a in self.agents if a.role == "writer")
        self.send_message(writer_agent, f"请基于研究结果写文档")

        # 模拟执行
        writer_agent.process_messages()

# 创建Swarm
research = ResearchAgent("ResearchBot", "researcher")
writer = WriterAgent("WriterBot", "writer")
manager = ManagerAgent("ManagerBot", [research, writer])

# 运行任务
manager.coordinate_task("AI技术发展趋势")