手写一个轻量级 AI 编程助手：从 0 到 1 实现代码生成与自愈流水线

┌─────────────────────────────────────────────────────────────┐
│                    输入：需求描述                            │
└─────────────────────┬───────────────────────────────────────┘
                      ▼
┌─────────────────────────────────────────────────────────────┐
│  1. CodeGenerator：根据需求 + 上下文生成代码              │
│     - 读取项目规范（.codebuddy/rules.md）                  │
│     - 读取相关已有代码（AST 解析依赖）                     │
│     - 调用 LLM 生成代码                                    │
└─────────────────────┬───────────────────────────────────────┘
                      ▼
┌─────────────────────────────────────────────────────────────┐
│  2. CodeExecutor：在沙箱中执行生成的代码                   │
│     - 写临时文件                                           │
│     - 执行单元测试 / 编译检查                              │
│     - 捕获 stdout / stderr                                 │
└─────────────────────┬───────────────────────────────────────┘
                      ▼
              ┌───────┴───────┐
              │  测试通过？    │
              └───────┬───────┘
                是 ▼       ▼ 否
┌──────────────────────────┐  ┌──────────────────────────────┐
│  3. 输出代码 & 报告      │  │  4. ErrorAnalyzer：分析错误  │
│                          │  │     - 截取关键错误栈          │
│                          │  │     - 归类错误类型            │
└──────────────────────────┘  └──────────────┬───────────────┘
                                              ▼
                                    ┌─────────────────────┐
                                    │ 5. CodeFixer：修复  │
                                    │    结合错误 + 代码  │
                                    │    调用 LLM 修复    │
                                    └──────────┬──────────┘
                                               ▼ (回到 Step 2)

# generator.py
import os
import json
from typing import Optional
from openai import OpenAI

class CodeGenerator:
    def __init__(self, model: str = "gpt-4", rules_path: str = ".codebuddy/rules.md"):
        self.client = OpenAI()
        self.model = model
        self.rules = self._load_rules(rules_path)
    
    def _load_rules(self, path: str) -> str:
        if os.path.exists(path):
            with open(path, "r") as f:
                return f.read()
        return "使用 Python 3.10+，遵循 PEP 8 规范。"
    
    def _collect_context(self, target_file: str) -> dict:
        """收集目标文件所在的模块上下文"""
        context = {"imports": [], "functions": [], "classes": []}
        if not os.path.exists(target_file):
            return context
        
        import ast
        with open(target_file, "r") as f:
            tree = ast.parse(f.read())
        
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    context["imports"].append(alias.name)
            elif isinstance(node, ast.FunctionDef):
                context["functions"].append(node.name)
            elif isinstance(node, ast.ClassDef):
                context["classes"].append(node.name)
        return context
    
    def generate(self, requirement: str, target_file: str) -> str:
        context = self._collect_context(target_file)
        
        prompt = f"""
你是一个资深 Python 工程师。请在保持项目风格一致的前提下，实现以下需求。

## 项目规范
{self.rules}

## 目标文件上下文
- 已有导入: {', '.join(context['imports'])}
- 已有函数: {', '.join(context['functions'])}
- 已有类: {', '.join(context['classes'])}

## 需求
{requirement}

## 输出要求
1. 只输出完整的代码，不要加解释
2. 如果需求是新增函数，输出包含该函数的完整文件内容
3. 如果需求是修改已有函数，输出修改后的完整函数定义

现在请输出代码：
"""
        
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "你是一个代码生成助手，只输出代码，不输出任何解释。"},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
        )
        
        return response.choices[0].message.content

# executor.py
import subprocess
import tempfile
import os
from typing import Tuple

class CodeExecutor:
    def __init__(self, workdir: str = "./workspace"):
        self.workdir = workdir
        os.makedirs(workdir, exist_ok=True)
    
    def execute(self, code: str, filename: str = "generated.py") -> Tuple[bool, str, str]:
        """
        执行代码，返回 (是否成功, stdout, stderr)
        """
        # 写入临时文件
        filepath = os.path.join(self.workdir, filename)
        with open(filepath, "w") as f:
            f.write(code)
        
        # 如果有测试文件，运行 pytest
        test_file = os.path.join(self.workdir, "test_" + filename)
        if os.path.exists(test_file):
            return self._run_pytest(test_file)
        
        # 否则仅做语法检查
        return self._syntax_check(filepath)
    
    def _run_pytest(self, test_file: str) -> Tuple[bool, str, str]:
        try:
            result = subprocess.run(
                ["pytest", test_file, "-v", "--tb=short"],
                cwd=self.workdir,
                capture_output=True,
                text=True,
                timeout=30
            )
            success = result.returncode == 0
            return success, result.stdout, result.stderr
        except subprocess.TimeoutExpired:
            return False, "", "测试执行超时（30秒）"
    
    def _syntax_check(self, filepath: str) -> Tuple[bool, str, str]:
        try:
            result = subprocess.run(
                ["python", "-m", "py_compile", filepath],
                capture_output=True,
                text=True,
                timeout=10
            )
            success = result.returncode == 0
            return success, "语法检查通过", result.stderr
        except subprocess.TimeoutExpired:
            return False, "", "语法检查超时"

# analyzer.py
import re
from typing import Dict, List

class ErrorAnalyzer:
    # 常见错误模式
    PATTERNS = {
        "name_error": re.compile(r"NameError: name '(\w+)' is not defined"),
        "type_error": re.compile(r"TypeError: (.*)"),
        "attribute_error": re.compile(r"AttributeError: '(\w+)' object has no attribute '(\w+)'"),
        "import_error": re.compile(r"ImportError: No module named '(\w+)'"),
        "syntax_error": re.compile(r"SyntaxError: (.*)"),
        "assertion_error": re.compile(r"AssertionError: (.*)"),
    }
    
    def analyze(self, stdout: str, stderr: str) -> Dict:
        """分析错误，返回结构化结果"""
        combined = stdout + "\n" + stderr
        
        # 提取最后 20 行（通常错误栈在末尾）
        lines = combined.strip().split("\n")
        tail = "\n".join(lines[-20:]) if len(lines) > 20 else combined
        
        result = {
            "error_type": "unknown",
            "error_message": "",
            "context_lines": tail,
            "suggested_fix": "",
            "traceback": self._extract_traceback(combined),
        }
        
        # 匹配错误类型
        for err_type, pattern in self.PATTERNS.items():
            match = pattern.search(combined)
            if match:
                result["error_type"] = err_type
                result["error_message"] = match.group(0)
                break
        
        # 根据错误类型给出修复建议方向
        fix_suggestions = {
            "name_error": "检查变量或函数是否已定义，可能需要添加 import 或调整作用域",
            "type_error": "检查参数类型是否匹配，可能需要类型转换或调整函数签名",
            "attribute_error": "检查对象是否有该属性，可能需要修改对象类型或添加属性",
            "import_error": "检查依赖是否已安装，或调整 import 路径",
            "syntax_error": "检查代码语法，可能缺少括号、冒号或缩进有误",
            "assertion_error": "检查测试断言与预期值是否一致，可能业务逻辑有误",
        }
        result["suggested_fix"] = fix_suggestions.get(result["error_type"], "请检查代码逻辑")
        
        return result
    
    def _extract_traceback(self, text: str) -> List[str]:
        """提取完整的错误栈"""
        lines = text.split("\n")
        in_traceback = False
        tb_lines = []
        for line in lines:
            if "Traceback (most recent call last):" in line:
                in_traceback = True
            if in_traceback:
                tb_lines.append(line)
        return tb_lines

# fixer.py
from openai import OpenAI
from typing import Dict

class CodeFixer:
    def __init__(self, model: str = "gpt-4"):
        self.client = OpenAI()
        self.model = model
    
    def fix(self, original_code: str, error_info: Dict) -> str:
        prompt = f"""
你是一个代码修复专家。下面的代码在执行时出现了错误，请修复它。

## 原始代码
```python
{original_code}

{error_info['traceback']}

### 2.5 主流水线：把一切串起来

```python
# pipeline.py
import json
from typing import Dict
from generator import CodeGenerator
from executor import CodeExecutor
from analyzer import ErrorAnalyzer
from fixer import CodeFixer

class AICodingPipeline:
    def __init__(self, max_retries: int = 3):
        self.generator = CodeGenerator()
        self.executor = CodeExecutor()
        self.analyzer = ErrorAnalyzer()
        self.fixer = CodeFixer()
        self.max_retries = max_retries
        self.history = []
    
    def run(self, requirement: str, target_file: str = "generated.py") -> Dict:
        """
        运行完整流水线
        """
        # Step 1: 生成代码
        print(f"[1/5] 正在生成代码...")
        code = self.generator.generate(requirement, target_file)
        self.history.append({"stage": "generate", "code": code})
        
        # Step 2-4: 迭代修复
        for attempt in range(1, self.max_retries + 1):
            print(f"[2/5-{attempt}] 执行代码 (尝试 {attempt}/{self.max_retries})...")
            success, stdout, stderr = self.executor.execute(code, target_file)
            
            if success:
                print(f"[3/5] ✅ 测试通过！")
                return self._result("success", code, stdout, self.history)
            
            # 分析错误
            print(f"[4/5-{attempt}] 分析错误...")
            error_info = self.analyzer.analyze(stdout, stderr)
            self.history.append({"stage": "error", "attempt": attempt, "error": error_info})
            
            if attempt < self.max_retries:
                print(f"[5/5-{attempt}] 尝试修复...")
                code = self.fixer.fix(code, error_info)
                self.history.append({"stage": "fix", "attempt": attempt, "code": code})
            else:
                print(f"[5/5] ❌ 已达到最大重试次数，修复失败")
                return self._result("failed", code, stderr, self.history, error_info)
        
        return self._result("error", code, "未知错误", self.history)
    
    def _result(self, status: str, code: str, logs: str, history: list, error: Dict = None):
        return {
            "status": status,
            "code": code,
            "logs": logs,
            "history": history,
            "error": error,
        }


# 使用示例
if __name__ == "__main__":
    pipeline = AICodingPipeline(max_retries=3)
    
    result = pipeline.run(
        requirement="""
实现一个 http_get 函数，支持：
1. 发送 GET 请求到指定 URL
2. 支持自定义超时时间（默认 5 秒）
3. 支持自定义请求头（字典格式）
4. 返回响应内容（字符串）和状态码
5. 如果超时或请求失败，抛出明确的异常
""",
        target_file="http_utils.py"
    )
    
    print("\n" + "=" * 50)
    print(f"状态: {result['status']}")
    print("=" * 50)
    print(result['code'])
    if result.get('error'):
        print("\n错误信息:", result['error']['error_message'])

# test_pipeline.py
import tempfile
import os
from pipeline import AICodingPipeline

def test_pipeline_on_real_requirement():
    """真实场景测试：实现一个带缓存功能的斐波那契函数"""
    
    pipeline = AICodingPipeline(max_retries=2)
    
    result = pipeline.run(
        requirement="""
实现 fibonacci(n) 函数：
- 输入 n 为非负整数
- 返回第 n 个斐波那契数（从 0 开始，fib(0)=0, fib(1)=1）
- 使用字典缓存来优化重复计算
- 如果 n < 0，抛出 ValueError
""",
        target_file="fib.py"
    )
    
    # 验证生成的代码是否能正确执行
    assert result['status'] == 'success', f"流水线失败: {result.get('error')}"
    
    # 验证生成的代码是否可以正确导入和执行
    temp_dir = tempfile.mkdtemp()
    code_path = os.path.join(temp_dir, "fib.py")
    with open(code_path, "w") as f:
        f.write(result['code'])
    
    import sys
    sys.path.insert(0, temp_dir)
    try:
        import fib
        assert fib.fibonacci(0) == 0
        assert fib.fibonacci(1) == 1
        assert fib.fibonacci(10) == 55
        print("✅ 所有测试通过！")
    finally:
        sys.path.remove(temp_dir)
        import shutil
        shutil.rmtree(temp_dir)


if __name__ == "__main__":
    test_pipeline_on_real_requirement()

def _retrieve_knowledge(self, query: str) -> str:
    """从企业知识库检索相关文档"""
    # 假设知识库是 Markdown 文件目录
    import glob
    docs = []
    for md_file in glob.glob("./llm-knowledge/**/*.md", recursive=True):
        with open(md_file, "r") as f:
            content = f.read()
            if any(kw in content.lower() for kw in query.lower().split()):
                docs.append(f"## {md_file}\n{content[:500]}...")
    return "\n\n".join(docs[:3])

def _lint_check(self, code: str) -> Tuple[bool, str]:
    """使用 flake8 或 pylint 检查代码质量"""
    # 用临时文件执行 lint
    import tempfile
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write(code)
        tmp_path = f.name
    
    try:
        result = subprocess.run(
            ["flake8", tmp_path, "--max-line-length=120"],
            capture_output=True, text=True
        )
        if result.returncode == 0:
            return True, "Lint 检查通过"
        else:
            return False, result.stdout
    finally:
        os.unlink(tmp_path)

重要：如果代码中使用了任何外部库，必须在文件开头添加对应的 import 语句。

if debug:
    with open(f"./debug/generated_{attempt}.py", "w") as f:
        f.write(code)
    with open(f"./debug/error_{attempt}.json", "w") as f:
        json.dump(error_info, f, indent=2)