纯Python代码从零构建一个AI多智能体系统

# agent.py
import json
import ollama
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class Message:
    role: str       # "user" | "agent" | "system"
    sender: str     # 发送者名称
    content: str


@dataclass
class Agent:
    name: str
    role: str
    system_prompt: str
    model: str = "qwen3:8b"
    history: list = field(default_factory=list)

    def think(self, messages: list[Message]) -> str:
        """思考并回复"""
        # 构建完整的消息列表
        full_messages = [
            {"role": "system", "content": self.system_prompt}
        ]
        for msg in messages:
            full_messages.append({
                "role": "user" if msg.sender != self.name else "assistant",
                "content": f"[来自 {msg.sender}] {msg.content}"
            })

        # 调用模型
        response = ollama.chat(
            model=self.model,
            messages=full_messages
        )
        return response["message"]["content"]

# agents.py
from agent import Agent

# 研究员 Agent：负责技术调研
researcher = Agent(
    name="Researcher",
    role="技术研究员",
    system_prompt="""你是一个技术研究员。你的任务是：
1. 分析用户需求中涉及的技术点
2. 给出最佳实践和技术方案建议
3. 列出关键注意事项和潜在风险

输出格式：
## 技术分析
（分析内容）

## 推荐方案
（方案详情）

## 注意事项
（风险和注意点）"""
)

# 程序员 Agent：负责写代码
coder = Agent(
    name="Coder",
    role="高级程序员",
    system_prompt="""你是一个高级 Python 程序员。你的任务是：
1. 根据需求和调研结果编写高质量的 Python 代码
2. 代码必须包含完整的类型注解和 docstring
3. 包含错误处理和日志记录
4. 输出可直接运行的完整代码

输出格式：
## 实现思路
（简述思路）

## 完整代码
```python
（代码）
```"""
)

# 审查员 Agent：负责代码审查
reviewer = Agent(
    name="Reviewer",
    role="代码审查专家",
    system_prompt="""你是一个严格的代码审查专家。你的任务是：
1. 审查代码的正确性、安全性和性能
2. 检查是否遵循最佳实践
3. 给出修改建议

评分标准（满分 10 分）：
- 正确性：代码逻辑是否正确
- 安全性：是否存在安全漏洞
- 可读性：命名、注释、代码结构
- 性能：是否有明显的性能问题

输出格式：
## 评分：X/10
## 优点
（列出优点）
## 问题
（列出问题）
## 修改建议
（具体建议）"""
)

# coordinator.py
from agent import Agent, Message


class Coordinator:
    """协调者：负责调度多个 Agent 协作"""

    def __init__(self, researcher, coder, reviewer):
        self.researcher = researcher
        self.coder = coder
        self.reviewer = reviewer
        self.log = []

    def _log(self, phase: str, content: str):
        """记录执行日志"""
        self.log.append({"phase": phase, "content": content})
        print(f"\n{'='*60}")
        print(f"📋 [{phase}]")
        print(f"{'='*60}")
        print(content[:500] + "..." if len(content) > 500 else content)

    def run(self, task: str, max_revisions: int = 2) -> dict:
        """
        执行完整的多智能体协作流程

        Args:
            task: 用户需求描述
            max_revisions: 最大修改轮数

        Returns:
            包含所有阶段输出的字典
        """
        results = {"task": task, "phases": []}

        # ========== 第一阶段：需求分析 + 技术调研 ==========
        print("\n🔍 第一阶段：技术调研")
        research_msg = Message(
            role="user", sender="Coordinator",
            content=f"请分析以下需求的技术方案：\n{task}"
        )
        research_result = self.researcher.think([research_msg])
        self._log("技术调研", research_result)
        results["phases"].append({
            "phase": "research", "output": research_result
        })

        # ========== 第二阶段：代码生成 ==========
        print("\n💻 第二阶段：代码生成")
        code_msg = Message(
            role="user", sender="Coordinator",
            content=f"""请根据以下需求和调研结果编写代码：

## 需求
{task}

## 技术调研结果
{research_result}

请直接输出完整可运行的 Python 代码。"""
        )
        code_result = self.coder.think([code_msg])
        self._log("代码生成", code_result)
        results["phases"].append({
            "phase": "code", "output": code_result
        })

        # ========== 第三阶段：代码审查（可循环） ==========
        current_code = code_result
        revision = 0

        while revision <= max_revisions:
            print(f"\n🧪 第三阶段：代码审查（第 {revision + 1} 次）")
            review_msg = Message(
                role="user", sender="Coordinator",
                content=f"请审查以下代码：\n{current_code}"
            )
            review_result = self.reviewer.think([review_msg])
            self._log(f"代码审查（第{revision+1}次）", review_result)

            # 检查评分（简单解析）
            score = self._extract_score(review_result)
            print(f"\n📊 当前评分：{score}/10")

            if score >= 8 or revision == max_revisions:
                results["phases"].append({
                    "phase": "review",
                    "output": review_result,
                    "score": score
                })
                break

            # 分数不够，让程序员修改
            print(f"\n🔄 评分不足，进入修改...")
            fix_msg = Message(
                role="user", sender="Coordinator",
                content=f"""请根据审查意见修改代码：

## 审查意见
{review_result}

## 原始代码
{current_code}

请输出修改后的完整代码。"""
            )
            current_code = self.coder.think([fix_msg])
            self._log(f"代码修改（第{revision+1}次）", current_code)
            revision += 1

        results["final_code"] = current_code
        results["revisions"] = revision + 1

        return results

    def _extract_score(self, review_text: str) -> int:
        """从审查结果中提取评分"""
        for line in review_text.split("\n"):
            if "评分" in line and "/" in line:
                try:
                    score = int(line.split(":")[-1].strip().split("/")[0])
                    return score
                except (ValueError, IndexError):
                    pass
        return 5  # 默认评分

# main.py
from agents import researcher, coder, reviewer
from coordinator import Coordinator
import json


def main():
    print("🚀 AI 研发小队已就绪！")
    print("=" * 60)

    # 初始化协调者
    coordinator = Coordinator(researcher, coder, reviewer)

    # 示例任务
    task = "用 Python 写一个文件批量重命名工具，支持按日期、序号、正则匹配重命名，需要命令行参数解析和日志记录"

    # 执行任务
    results = coordinator.run(task, max_revisions=2)

    # 输出最终结果
    print("\n" + "=" * 60)
    print("✅ 任务完成！")
    print("=" * 60)
    print(f"\n📊 总修改轮数：{results['revisions']}")
    print(f"\n💻 最终代码：\n{results['final_code']}")

    # 保存结果
    with open("output.json", "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print("\n💾 完整结果已保存到 output.json")


if __name__ == "__main__":
    main()

# tools.py
import os
import subprocess


def write_file(path: str, content: str) -> str:
    """写入文件"""
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    return f"文件已写入: {path}"


def run_python(filepath: str) -> str:
    """运行 Python 文件并返回结果"""
    try:
        result = subprocess.run(
            ["python", filepath],
            capture_output=True, text=True, timeout=30
        )
        output = result.stdout
        if result.returncode != 0:
            output += f"\n❌ 错误:\n{result.stderr}"
        return output
    except subprocess.TimeoutExpired:
        return "执行超时（30秒）"


# 将工具注册到 Agent
AVAILABLE_TOOLS = {
    "write_file": {
        "func": write_file,
        "description": "将内容写入指定文件",
        "params": ["path", "content"]
    },
    "run_python": {
        "func": run_python,
        "description": "运行 Python 文件",
        "params": ["filepath"]
    }
}