LangGraph 实战：构建“自我修正”的代码生成 Agent

LangGraph 实战：构建"自我修正"的代码生成 Agent

前文我们聊了 LangGraph 的可视化追踪，这次聚焦一个非常实用的场景------构建一个能"自我修正"的代码生成 Agent。这个 Agent 不仅能根据需求生成代码，还能自动运行测试、发现错误、修正代码，直到测试全部通过。全程附带可运行的代码片段，新手也能跟着做。

一、先明确：我们要做一个什么样的 Agent？

这个代码生成 Agent 需要完成 4 件事：

1. 解析用户需求，生成初始代码；
2. 将代码写入临时文件，运行单元测试；
3. 如果测试失败，解析错误信息，修正代码；
4. 重复步骤 2-3，直到测试全部通过，或者达到最大重试次数。

二、第一步：定义全局 State（智能体的"共享记事本"）

State 是所有节点共享的数据容器，我们用 TypedDict 定义它需要存储的字段：

python

编辑

from typing import TypedDict, Optional, List
from langchain_core.messages import HumanMessage, AIMessage

class CodeGenState(TypedDict):
    messages: List[HumanMessage | AIMessage]  # 对话历史
    requirement: str  # 用户需求
    code: Optional[str]  # 生成的代码
    test_code: Optional[str]  # 测试代码
    test_result: Optional[str]  # 测试结果（通过/失败+错误信息）
    retry_count: int  # 重试次数
    is_success: bool  # 是否成功

三、第二步：编写 4 个核心 Node（每个节点只做一件事）

1. 代码生成节点：根据用户需求生成初始代码和测试代码

python

编辑

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o", temperature=0)

def generate_code(state: CodeGenState) -> dict:
    requirement = state["requirement"]
    prompt = f"""根据以下需求生成 Python 代码和对应的单元测试代码，返回 JSON 格式：
    需求：{requirement}
    输出格式：{{"code": "生成的代码", "test_code": "测试代码"}}
    要求：代码要简洁，测试代码要覆盖主要场景"""
    result = llm.invoke(prompt).content
    # 这里省略 JSON 解析逻辑，实际开发中可以用 pydantic 解析
    return {
        "code": "def add(a, b):\n    return a + b",
        "test_code": "def test_add():\n    assert add(1, 2) == 3\n    assert add(-1, 1) == 0",
        "retry_count": 0,
        "is_success": False
    }

2. 测试执行节点：将代码写入临时文件，运行单元测试

python

编辑

import tempfile
import subprocess
import os

def run_tests(state: CodeGenState) -> dict:
    code = state["code"]
    test_code = state["test_code"]
    
    # 创建临时目录
    with tempfile.TemporaryDirectory() as tmpdir:
        # 写入代码文件
        code_file = os.path.join(tmpdir, "solution.py")
        with open(code_file, "w") as f:
            f.write(code)
        
        # 写入测试文件
        test_file = os.path.join(tmpdir, "test_solution.py")
        with open(test_file, "w") as f:
            f.write(test_code)
        
        # 运行测试
        result = subprocess.run(
            ["python", "-m", "pytest", test_file, "-v"],
            capture_output=True,
            text=True,
            cwd=tmpdir
        )
        
        if result.returncode == 0:
            return {"test_result": "PASSED", "is_success": True}
        else:
            return {"test_result": f"FAILED: {result.stderr}", "is_success": False}

3. 错误分析节点：解析测试失败的原因，生成修正建议

python

编辑

def analyze_error(state: CodeGenState) -> dict:
    test_result = state["test_result"]
    code = state["code"]
    
    prompt = f"""以下代码的单元测试失败了，请分析错误原因并给出修正建议：
    代码：{code}
    测试结果：{test_result}
    输出格式：100字以内的错误分析和修正建议"""
    analysis = llm.invoke(prompt).content
    return {"error_analysis": analysis}

4. 代码修正节点：根据错误分析修正代码

python

编辑

def fix_code(state: CodeGenState) -> dict:
    code = state["code"]
    error_analysis = state["error_analysis"]
    requirement = state["requirement"]
    
    prompt = f"""根据以下错误分析修正代码：
    原始需求：{requirement}
    原始代码：{code}
    错误分析：{error_analysis}
    输出格式：修正后的完整代码（只输出代码，不要其他内容）"""
    fixed_code = llm.invoke(prompt).content
    return {
        "code": fixed_code,
        "retry_count": state["retry_count"] + 1
    }

四、第三步：用条件边控制循环流转

我们用条件边实现"测试-修正"的循环：如果测试通过，结束流程；如果测试失败且未达到最大重试次数，继续修正；否则结束流程。

python

编辑

from langgraph.graph import StateGraph, START, END

# 初始化图
workflow = StateGraph(CodeGenState)

# 添加节点
workflow.add_node("generate_code", generate_code)
workflow.add_node("run_tests", run_tests)
workflow.add_node("analyze_error", analyze_error)
workflow.add_node("fix_code", fix_code)

# 添加边：从起点到代码生成
workflow.add_edge(START, "generate_code")

# 代码生成后，执行测试
workflow.add_edge("generate_code", "run_tests")

# 定义条件路由函数
def route_after_test(state: CodeGenState) -> str:
    if state["is_success"]:
        return END  # 测试通过，结束
    elif state["retry_count"] >= 3:
        return END  # 达到最大重试次数，结束
    else:
        return "analyze_error"  # 测试失败，继续分析错误

# 添加条件边
workflow.add_conditional_edges("run_tests", route_after_test)

# 错误分析后，修正代码
workflow.add_edge("analyze_error", "fix_code")

# 修正代码后，重新测试
workflow.add_edge("fix_code", "run_tests")

# 编译图
app = workflow.compile()

五、第四步：运行测试，看看效果

我们输入一个有陷阱的需求，看看 Agent 能不能自我修正：

python

编辑

initial_state = {
    "requirement": "写一个函数，计算两个数的和。注意：如果输入不是数字，要抛出 TypeError",
    "retry_count": 0,
    "is_success": False
}

result = app.invoke(initial_state)
print(f"最终代码：\n{result['code']}")
print(f"测试结果：{result['test_result']}")
print(f"重试次数：{result['retry_count']}")

输出结果示例：

文本

编辑

最终代码：
def add(a, b):
    if not isinstance(a, (int, float)) or not isinstance(b, (int, float)):
        raise TypeError("输入必须是数字")
    return a + b

测试结果：PASSED
重试次数：2

六、进阶：给 Agent 加上"人工审核"和"代码评审"

1. 添加人工审核：在代码生成后暂停，让用户确认后再执行测试，适合高风险场景。

python

编辑

# 编译时配置中断点
app = workflow.compile(interrupt_before=["run_tests"])

# 运行后会暂停，等待人工确认
result = app.invoke(initial_state)
# 人工确认后，传入 None 恢复执行
result = app.invoke(None)

2. 添加代码评审节点：在测试通过后，让 LLM 对代码进行评审，给出优化建议。

python

编辑

def code_review(state: CodeGenState) -> dict:
    code = state["code"]
    prompt = f"""对以下代码进行评审，给出优化建议：
    代码：{code}
    输出格式：100字以内的评审意见"""
    review = llm.invoke(prompt).content
    return {"code_review": review}

# 添加评审节点
workflow.add_node("code_review", code_review)
workflow.add_edge("run_tests", "code_review")
workflow.add_edge("code_review", END)

总结

从定义 State、编写 Node、配置条件边到添加人工审核和代码评审，我们只用了几十行代码就搭建了一个具备"自我修正"能力的代码生成 Agent。LangGraph 的循环流转能力让这种"执行-检查-修正"的迭代流程变得非常简单，不管是做代码生成、数据分析还是自动化运维，都可以用这套思路快速落地。