Agent 工程化实践：构建 AI 编程助手驱动的自动化代码生成流水线

# tools.py
import subprocess
import tempfile
import os
from pathlib import Path
from typing import Dict, Any, Optional
from git import Repo
import pytest
from flake8.api import legacy as flake8

class ToolResult:
    def __init__(self, success: bool, output: str, error: Optional[str] = None):
        self.success = success
        self.output = output
        self.error = error

def generate_code_tool(query: str) -> ToolResult:
    """调用 LLM 生成代码（实际被 Agent 内部调用，此处为工具封装）"""
    # 此工具实际由 Agent 的 LLM 直接生成，但为了统一接口，我们仍定义为一个“伪工具”
    # 真正的代码生成是在 Agent 的思考步骤中由 LLM 输出，但这里我们保留一个工具供显式调用
    pass

def write_test_tool(code: str, test_code: str) -> ToolResult:
    """将测试代码写入 tests/ 目录"""
    try:
        os.makedirs("tests", exist_ok=True)
        with open("tests/test_generated.py", "w") as f:
            f.write(test_code)
        return ToolResult(True, "Test file written successfully.")
    except Exception as e:
        return ToolResult(False, "", str(e))

def run_lint_tool(file_path: str) -> ToolResult:
    """使用 flake8 检查代码风格"""
    style_guide = flake8.get_style_guide()
    report = style_guide.check_files([file_path])
    if report.total_errors == 0:
        return ToolResult(True, "Lint passed.")
    else:
        errors = [f"{err}" for err in report._deferred_print]
        return ToolResult(False, "", "\n".join(errors))

def run_test_tool() -> ToolResult:
    """运行 pytest，返回结果"""
    result = subprocess.run(["pytest", "-v", "--tb=short"], capture_output=True, text=True)
    if result.returncode == 0:
        return ToolResult(True, result.stdout)
    else:
        return ToolResult(False, result.stdout, result.stderr)

def commit_and_pr_tool(branch_name: str, commit_msg: str, repo_path: str = ".") -> ToolResult:
    """提交代码并创建 PR（需配置 GitLab/ GitHub token）"""
    try:
        repo = Repo(repo_path)
        # 切换或创建新分支
        if branch_name not in repo.branches:
            repo.create_head(branch_name)
        repo.head.reference = repo.branches[branch_name]
        repo.head.reset(index=True, working_tree=True)
        # 添加所有变更
        repo.git.add(A=True)
        repo.index.commit(commit_msg)
        # push 到远程
        origin = repo.remote(name="origin")
        origin.push(refspec=f"{branch_name}:{branch_name}", set_upstream=True)
        # 创建 PR（调用 GitHub API）
        # 此处省略具体 API 调用，使用 PyGithub 或 requests
        return ToolResult(True, f"PR created from branch {branch_name}")
    except Exception as e:
        return ToolResult(False, "", str(e))

# agent.py
from langchain.tools import StructuredTool
from langchain.pydantic_v1 import BaseModel, Field
from typing import Type

class CodeGenInput(BaseModel):
    requirement: str = Field(description="用户需求描述")

class TestGenInput(BaseModel):
    code: str = Field(description="待测代码")
    test_code: str = Field(description="测试代码")

# 定义工具实例
tools = [
    StructuredTool.from_function(
        func=lambda req: generate_code_tool(req),  # 实际我们会用 LLM 直接生成，这里作为占位
        name="generate_code",
        description="根据需求生成 Python 代码",
        args_schema=CodeGenInput,
    ),
    StructuredTool.from_function(
        func=write_test_tool,
        name="write_test",
        description="将测试代码写入 tests/test_generated.py",
        args_schema=TestGenInput,
    ),
    StructuredTool.from_function(
        func=run_lint_tool,
        name="run_lint",
        description="对指定文件运行 flake8 检查",
        args_schema=FileInput,  # 需定义
    ),
    StructuredTool.from_function(
        func=run_test_tool,
        name="run_test",
        description="运行 pytest 测试",
        args_schema=EmptyInput,
    ),
    StructuredTool.from_function(
        func=commit_and_pr_tool,
        name="commit_and_pr",
        description="提交代码并创建 PR",
        args_schema=CommitInput,  # 需定义
    ),
]

from langchain.agents import create_openai_tools_agent, AgentExecutor
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder

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

prompt = ChatPromptTemplate.from_messages([
    ("system", """你是一位资深的 Python 工程师，负责根据用户需求生成高质量的代码。
你需要按以下步骤操作：
1. 先编写实现代码，保存为 main.py
2. 再编写对应的 pytest 测试代码，保存为 tests/test_main.py
3. 然后运行 flake8 检查 main.py，若有错误则修复
4. 最后运行 pytest，确保全部通过
5. 所有检查通过后，提交代码并创建 PR
你的所有输出必须包含代码块（```python ... ```）或者调用工具。
"""),
    MessagesPlaceholder(variable_name="chat_history", optional=True),
    ("human", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad"),
])

# 我们只需要提供三个工具：run_lint, run_test, commit_and_pr，以及一个自定义的“写文件”工具
# 但为了简化，我们将 write_test 和 write_code 合并为文件写入工具
def write_file_tool(filepath: str, content: str) -> str:
    try:
        Path(filepath).parent.mkdir(parents=True, exist_ok=True)
        with open(filepath, "w") as f:
            f.write(content)
        return f"File {filepath} written successfully."
    except Exception as e:
        return f"Error: {str(e)}"

write_file = StructuredTool.from_function(
    func=write_file_tool,
    name="write_file",
    description="将内容写入指定文件",
)

# 重新定义工具列表
tools = [write_file, run_lint_tool_wrapped, run_test_tool_wrapped, commit_and_pr_tool_wrapped]

agent = create_openai_tools_agent(llm, tools, prompt)
executor = AgentExecutor(agent=agent, tools=tools, verbose=True, max_iterations=10)

# main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

app = FastAPI()

class TaskRequest(BaseModel):
    requirement: str
    repo_url: str  # 可选，指定仓库

class TaskResponse(BaseModel):
    status: str
    pr_url: str | None = None
    logs: str

@app.post("/generate")
async def generate_code(req: TaskRequest):
    # 将需求注入，并设置工作目录（需克隆 repo）
    try:
        # 实际需先 git clone 到临时目录
        result = executor.invoke({"input": req.requirement})
        # 解析结果，提取 PR 链接（从工具输出）
        logs = result["output"]
        # 简化返回
        return TaskResponse(status="success", logs=logs)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

import docker
client = docker.from_env()

def run_in_sandbox(code: str, test_code: str) -> dict:
    container = client.containers.run(
        "python:3.11-slim",
        command=["bash", "-c", "echo '{}' > main.py && echo '{}' > test_main.py && pytest".format(code, test_code)],
        detach=True,
        remove=True,
        mem_limit="512m",
        network_disabled=True,  # 禁止网络访问
    )
    logs = container.logs(stdout=True, stderr=True)
    return {"logs": logs.decode()}

from langchain.callbacks import FileCallbackHandler
handler = FileCallbackHandler("agent.log")
executor = AgentExecutor(..., callbacks=[handler])