怎么构建AI生成代码的AlphaCodium引擎,通过 RAG 和自我修正来生成代码

131 阅读17分钟

大家知道 AlphaCodium 吗, 它提出了一种利用控制流进行代码生成的方法。它的核心思想就是通过迭代的方式逐步构建和改进代码答案。它不仅仅依赖于一次性生成代码,而是通过多次测试和修正来优化生成的代码,确保其正确性和鲁棒性。AlphaCodium 在公开测试和 AI 生成的测试上迭代测试并改进针对特定问题的答案。 再问一个问题大家知道 AlphaCodium 是怎么为我们生成代码的吗?我来为大家梳理一下步骤, 首先,AlphaCodium 会理解用户提出的编程问题,并生成一个初始的代码解决方案。这一步通常基于大型语言模型(LLM)的能力,结合问题描述生成初步代码。 其次 AlphaCodium 会使用一组公开测试用例和AI 生成的测试用例来验证初始代码的正确性。如果代码未通过测试,AlphaCodium 会分析失败原因,并生成改进后的代码。 通过多次迭代,AlphaCodium 不断测试和改进代码,直到代码能够通过所有测试用例。每次迭代都会生成新的代码版本,并对其进行测试和修正。最终生成的代码会被格式化为结构化输出,确保其可读性和可维护性。AlphaCodium 采用的的这种迭代式方法为代码生成提供了一种新的思路,结合了测试驱动开发和 AI 的强大生成能力,能够显著提高代码生成的准确性和实用性。然后接下来我们将使用 LangGraph从头实现其中的一些思想:

  • 从用户指定的一组文档开始。
  • 使用长上下文 LLM 来消化这些文档,并通过 RAG(检索增强生成)基于文档回答问题。
  • 调用工具生成结构化输出。
  • 在将解决方案返回给用户之前,执行两项单元测试(检查导入和代码执行)。 在这里插入图片描述 有了架构设计之后,接下来开始我们的代码。

导入文件库

我们先从 LangChain 的网站抓取和处 LCEL(LangChain Expression Language) 文档,使用BeautifulSoup(别名为Soup)解析HTML内容,使用 RecursiveUrlLoader 来递归抓取网页内容。

from bs4 import BeautifulSoup as Soup
from langchain_community.document_loaders.recursive_url_loader import RecursiveUrlLoader

# LCEL docs
url = "https://python.langchain.com/docs/concepts/lcel/"
loader = RecursiveUrlLoader(
    url=url, max_depth=20, extractor=lambda x: Soup(x, "html.parser").text
)
docs = loader.load()

然后按照源URL对文档列表进行排序,反转排序后的列表(可能是为了优先处理更重要或更基础的文档),将所有文档的内容用分隔符 \n\n\n --- \n\n\n 连接成一个大字符串 concatenated_content

d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
d_reversed = list(reversed(d_sorted))
concatenated_content = "\n\n\n --- \n\n\n".join(
    [doc.page_content for doc in d_reversed]
)

LCEL 中构建一个 RAG 链

我们利用 LangChain 和 OpenAI 模型(如 GPT-4o-mini)自动生成针对 LCEL(LangChain Expression Language)问题的可执行代码解决方案。 定义一个提示模版:

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

code_gen_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """你是一个具备 LCEL(LangChain 表达式语言)专业知识的编程助手。以下是完整的 LCEL 文档:\n ------- \n  {context} \n ------- \n
请根据上述提供的文档回答用户的问题。确保你提供的任何代码都可以执行,并且所有必需的导入和变量都已定义。
按照以下结构回答问题:代码解决方案的描述。列出所需的导入。最后列出可执行的代码块。
用户的问题是:""",
        ),
        ("placeholder", "{messages}"),
    ]
)

定义数据模型:

# Data model
class code(BaseModel):
    """LCEL(LangChain 表达式语言)问题的代码解决方案模板。"""

    prefix: str = Field(description="问题描述与方法说明")
    imports: str = Field(description="代码块上面的导入")
    code: str = Field(description="代码块(不包括导入语句)")

构建模型调用链,并执行调用:

expt_llm = "gpt-4o-mini"
llm = ChatOpenAI(temperature=0, model=expt_llm)
code_gen_chain_oai = code_gen_prompt | llm.with_structured_output(code)
question = "如何在 LCEL 中构建一个 RAG 链?"
solution = code_gen_chain_oai.invoke(
    {"context": concatenated_content, "messages": [("user", question)]}
)
print(solution)

得到如下结果:

prefix='在 LCEL 中构建一个 RAG(Retrieval-Augmented Generation)链的过程涉及将检索和生成模型结合起来,以便在生成响应时利用外部知识。以下是构建 RAG 链的步骤和示例代码。' imports='from langchain_core.runnables import RunnableSequence, RunnableParallel\nfrom langchain_core.runnables import RunnableLambda\nfrom langchain_core.runnables import LLM, Retriever' code='# 定义检索器和生成模型\nretriever = Retriever(...)  # 替换为实际的检索器实现\nllm = LLM(...)  # 替换为实际的生成模型实现\n\n# 定义检索和生成的链\nretrieval_chain = RunnableLambda(lambda query: retriever.retrieve(query))\n\n# 定义 RAG 链\nrag_chain = RunnableSequence([\n    retrieval_chain,\n    llm\n])\n\n# 调用 RAG 链\nfinal_output = rag_chain.invoke("你的查询内容")\nprint(final_output)'

上面的代码还存在2个问题,不知道小伙伴们看出来没有,第一个是大模型可能忽略结构化输出要求,返回自由文本,不方便我们结构化解析。第二如果大模型首次返回了不符合格式的响应,比如json语法错误,那这样子并不能达到更精准的反馈,可能偏差还非常大。接下来我们就要通过强制调用工具来正确构建输出(这里的“工具”指的是结构化输出机制,而不是传统的 LangChain 工具 Tool类),然后增加错误重试机制,先设定一个重试值,捕获解析错误或未调用工具的情况,将错误信息附加到消息历史中,重新执行链,引导模型修正输出。最后通过我们前面的 Pydantic 模型 code 定义的字段和类型来进行结构化输出验证。

from langchain_core.prompts import ChatPromptTemplate

# 提示模版
code_gen_prompt_claude = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """<instructions> 您是一位在LCEL(LangChain表达式语言)方面拥有专业知识的编程助手。\n 
这里是LCEL文档:\n ------- \n  {context} \n ------- \n 请根据\n 
上述提供的文档回答用户问题。确保您提供的任何代码都可以执行,包含所有必需的导入和已定义的变量。\n
组织您的回答结构:1)描述代码解决方案的前缀,2)导入语句,3)功能完整的代码块。\n
调用代码工具以正确构建输出。</instructions> \n 这是用户问题:""",
        ),
        ("placeholder", "{messages}"),
    ]
)

expt_llm = "gpt-4o-mini"
llm = ChatOpenAI(temperature=0, model=expt_llm)
structured_llm_claude = llm.with_structured_output(code, include_raw=True)

# 如果工具使用不稳定,请检查错误。模型需要按照指定的格式(如code类)生成响应。因此,“调用工具”实际上是要求模型生成符合该格式的输出,而不是调用外部的函数或API。
def check_claude_output(tool_output):
    """检查解析错误或调用工具失败"""

    # 解析报错
    if tool_output["parsing_error"]:
        # 报告输出和解析错误。
        print("解析错误")
        raw_output = str(tool_output["raw"].content)
        error = tool_output["parsing_error"]
        raise ValueError(
            f"解析输出时出错!请确保调用工具。输出:{raw_output}。\n 解析错误:{error}"
        )

    # tool未被调用
    elif not tool_output["parsed"]:
        print("调用tool失败!")
        raise ValueError(
            "您未使用提供的工具!请确保调用工具以结构化输出。."
        )
    return tool_output


# 带输出检查的链式处理
code_chain_claude_raw = (
    code_gen_prompt_claude | structured_llm_claude | check_claude_output
)
# 错误恢复机制
def insert_errors(inputs):
    """消息中插入工具解析的错误信息"""

    # Get errors
    error = inputs["error"]
    messages = inputs["messages"]
    messages += [
        (
            "assistant",
            f"重试。您需要修复解析错误:{error} \n\n 您必须调用提供的工具.",
        )
    ]
    return {
        "messages": messages,
        "context": inputs["context"],
    }


# 此操作将作为后备链执行
fallback_chain = insert_errors | code_chain_claude_raw
N = 3  # 最大的重试次数
code_gen_chain_re_try = code_chain_claude_raw.with_fallbacks(
    fallbacks=[fallback_chain] * N, exception_key="error"
)


def parse_output(solution):
    """它将返回一个包含 'raw'(原始数据)、'parsed'(解析后数据)和 'parsing_error'(解析错误)的字典。"""

    return solution["parsed"]


# 支持重试机制,用于修复工具调用失败的情况
code_gen_chain = code_gen_chain_re_try | parse_output

# 假如没有重试
#code_gen_chain = code_gen_prompt_claude | structured_llm_claude | parse_output

这下我们来调用验证上面错误检查和重试逻辑:

question = "如何在 LCEL 中构建一个 RAG 链?"
solution = code_gen_chain.invoke(
    {"context": concatenated_content, "messages": [("user", question)]}
)
print(solution)

结果

prefix='以下代码展示了如何在 LCEL 中构建一个检索增强生成(RAG)链。该链结合了检索器和生成模型,以便在生成响应时利用外部知识。' imports='from langchain_core.runnables import RunnableSequence, RunnableParallel\nfrom langchain_core.runnables import RunnableLambda\nfrom langchain_core.runnables import LLM, Retriever' code='# 定义检索器和生成模型\nretriever = Retriever(...)  # 替换为实际的检索器实现\nllm = LLM(...)  # 替换为实际的生成模型实现\n\n# 定义 RAG 链\nrag_chain = RunnableSequence([\n    retriever,  # 首先调用检索器\n    llm  # 然后将检索到的结果传递给生成模型\n])\n\n# 调用 RAG 链\nfinal_output = rag_chain.invoke(query)  # query 是用户输入的查询'

通过我们上面增加结构化输出约束 + 错误重试机制,构建了一个鲁棒的代码生成流水线,特别适合需要严格遵循格式规范的场景(如生成可执行代码)。这种设计思想可推广至其他需要模型遵守特定输出格式的任务。

LangGraph构建 Agent来完成代码的生成

现在我们换LangGraph上场,用它来完成我们上面的任务。 首先定义一个状态结构,将包含与代码生成相关的键(错误、问题、代码生成)。这里我们定义了四个字段, error : 控制流的二进制标志,用于指示是否触发了测试错误,messages : 包含用户问题、错误信息、推理,generation : 代码解决方案, iterations : 尝试次数。

from typing import List
from typing_extensions import TypedDict


class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        error : 控制流的二进制标志,用于指示是否触发了测试错误
        messages : 包含用户问题、错误信息、推理
        generation : 代码解决方案
        iterations : 尝试次数
    """

    error: str
    messages: List
    generation: str
    iterations: int

然后定义 generate 节点,负责生成代码解决方案。这里使用code_gen_chain工具和提供的上下文内容生成代码,也就是我们上面构建的新方式。 如果之前有错误发生,会添加提示消息要求重试,将生成的代码(包括前缀、导入和代码块)添加到对话历史中,递增迭代计数器并返回更新后的状态。

def generate(state: GraphState):
    """
    参数:
    state (dict): 当前的图状态

    返回:
    state (dict): 添加了新键的 state,generation
    """

    print("---生成代码---")

    # State
    messages = state["messages"]
    iterations = state["iterations"]
    error = state["error"]

    # 我们已被路由回生成阶段,并出现了一个错误。
    if error == "yes":
        messages += [
            (
                "user",
                "现在,请重试。调用代码工具来构建输出,包括前缀、导入和代码块:",
            )
        ]

    # Solution
    code_solution = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )
    messages += [
        (
            "assistant",
            f"{code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
        )
    ]

    # Increment
    iterations = iterations + 1
    return {"generation": code_solution, "messages": messages, "iterations": iterations}

再增加一个验证节点,负责验证生成的代码质量。这里我们分两步检查代码:首先测试导入语句是否可执行,然后测试完整代码是否可运行。如果发现错误,会记录错误信息并标记错误状态为"yes"。如果代码通过所有测试,标记错误状态为"no"。最后返回更新后的状态,包括错误信息和标记。

def code_check(state: GraphState):
    """
    参数:
    state (dict): 当前的图状态

    返回:
    state (dict): 添加了新键的 state,error
    """

    print("---检查代码---")

    # State
    messages = state["messages"]
    code_solution = state["generation"]
    iterations = state["iterations"]

    # Get solution components
    imports = code_solution.imports
    code = code_solution.code

    # Check imports
    try:
        exec(imports)
    except Exception as e:
        print("---代码测试失败---")
        error_message = [("user", f"您的解决方案未能通过导入测试。: {e}")]
        messages += error_message
        return {
            "generation": code_solution,
            "messages": messages,
            "iterations": iterations,
            "error": "yes",
        }

    # Check execution
    try:
        exec(imports + "\n" + code)
    except Exception as e:
        print("---代码测试失败---")
        error_message = [("user", f"您的解决方案未能通过导入测试。: {e}")]
        messages += error_message
        return {
            "generation": code_solution,
            "messages": messages,
            "iterations": iterations,
            "error": "yes",
        }

    # No errors
    print("---没有代码测试失败---")
    return {
        "generation": code_solution,
        "messages": messages,
        "iterations": iterations,
        "error": "no",
    }

这里我们要增加一个反思节点,非常重要,当代码生成失败时,需要提供错误分析和反思,还是使用相同的 code_gen_chain 工具来分析错误并提供见解,将反思结果添加到对话历史中,返回更新后的状态,保持代码解决方案和迭代计数不变。

def reflect(state: GraphState):
    """
    参数:
    state (dict): 当前的图状态

    返回:
    state (dict): 添加了新键的 state,generation
    """

    print("---反思生成代码---")

    # State
    messages = state["messages"]
    iterations = state["iterations"]
    code_solution = state["generation"]

    # 增加一个反思
    reflections = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": messages}
    )
    messages += [("assistant", f"这里是关于错误的反思。: {reflections}")]
    return {"generation": code_solution, "messages": messages, "iterations": iterations}

接下来设置 edge 作为流程控制的决策点,根据错误状态和已完成的迭代次数决定下一步行动,如果没有错误或已达到最大迭代次数,结束流程返回"end",当启用反思机制时,引导流程到"reflect"节点,未启用反思时,直接返回"generate"节点继续生成代码。

def decide_to_finish(state: GraphState):
    """
    参数:
    state (dict): 当前的图状态

    返回:
    str: 要调用的下一个节点
    """
    error = state["error"]
    iterations = state["iterations"]

    if error == "no" or iterations == max_iterations:
        print("---完成---")
        return "end"
    else:
        print("---重试进入reflect或者generate---")
        if flag == "reflect":
            return "reflect"
        else:
            return "generate"

最后我们通过上面的节点和边构建自己图:

from langgraph.graph import END, StateGraph, START

workflow = StateGraph(GraphState)

# Define the nodes
workflow.add_node("generate", generate)  # generation solution
workflow.add_node("check_code", code_check)  # check code
workflow.add_node("reflect", reflect)  # reflect

# Build graph
workflow.add_edge(START, "generate")
workflow.add_edge("generate", "check_code")
workflow.add_conditional_edges(
    "check_code",
    decide_to_finish,
    {
        "end": END,
        "reflect": "reflect",
        "generate": "generate",
    },
)
workflow.add_edge("reflect", "generate")
app = workflow.compile()

from IPython.display import display, Image
display(Image(app.get_graph().draw_mermaid_png()))

结果如下图所示: 在这里插入图片描述 我们来调用它试试看:

question = "如何能直接将一个字符串传递给可运行对象(runnable)并用它来构建我的提示所需的输入?"
solution = app.invoke({"messages": [("user", question)], "iterations": 0, "error": ""})

达到下面结果:

---生成代码---
---检查代码---
这是您的提示: Hello, LangChain! 这是附加信息。
---没有代码测试失败---
---完成---

然后打印一下结果:

print(solution["generation"])
prefix='您可以使用 `RunnableLambda` 来将字符串直接传递给可运行对象,并将其用作构建提示的输入。以下是一个示例代码,展示了如何实现这一点。' imports='from langchain_core.runnables import RunnableLambda, RunnableSequence' code='# 定义一个函数,将输入字符串转换为提示\n\ndef create_prompt(input_string):\n    return f"这是您的提示: {input_string}"\n\n# 创建一个 RunnableLambda 对象\nrunnable_prompt = RunnableLambda(create_prompt)\n\n# 创建一个可运行的序列,将字符串传递给 RunnableLambda\ninput_string = "Hello, LangChain!"\nchain = runnable_prompt | RunnableLambda(lambda x: x + \' 这是附加信息。\')\n\n# 执行链并获取最终输出\nfinal_output = chain.invoke(input_string)\nprint(final_output)  # 输出: 这是您的提示: Hello, LangChain! 这是附加信息。'

到这里我们就已经完成我们代码的生成任务了,但是还一个地方需要注意,大家有没有想过一个问题,我们怎么才能保证它生成的code一定是正确可执行的呢,万一有错误呢。大家别担心,下面我们展示一种非常新的手段来验证,langsmith。

LangSmith的Eval(评估)

langsmith的评估主要有下面几个作用: 自动化评估:帮助开发者自动评估AI模型生成内容的质量和性能,不需要手动检查每个输出。 定制评估指标:允许开发者创建自定义评估函数,以针对特定任务或领域进行评估。 对比不同模型或方法:可以比较不同模型或不同实现方法(比如context stuffing和LangGraph方法)的性能差异。 数据集管理:支持创建、克隆和管理评估数据集,便于进行一致性测试和基准测试。 质量控制:帮助识别和解决AI系统中的缺陷或性能问题,确保生成的内容满足质量标准。 性能追踪:随着时间的推移跟踪模型性能的变化,帮助开发团队了解改进是否有效。 LangSmith被用来评估代码生成的质量,检查生成的代码是否能够正确编译和执行,这是一个非常实用的应用场景。 为了使用这种方式评估,我们首先下载一个LCEL问题的公开数据集。然后将其保存为 lcel-teacher-eval。初始化LangSmith客户端:导入langsmith库,创建一个LangSmith客户端实例,克隆公共数据集,尝试从指定URL克隆一个公共数据集到用户的LangSmith账户中,如果克隆失败,输出提示信息(例如,因为LangSmith未正确设置)。

import os
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ["LANGCHAIN_API_KEY"] = "lsv2_pt_9****************"
os.environ["LANGCHAIN_PROJECT"] = "test"


import langsmith

client = langsmith.Client()

# Clone the dataset to your tenant to use it
try:
    public_dataset = (
        "https://smith.langchain.com/public/326674a6-62bd-462d-88ae-eea49d503f9d/d"
    )
    client.clone_public_dataset(public_dataset)
except:
    print("Please setup LangSmith")

这段代码为后续的评估和测试工作做准备,它确保了:所有LangChain操作都会被追踪和记录,LangSmith客户端已正确配置,必要的评估数据集已经可用。 然后我们定义两个评估函数,check_import: 测试生成的导入语句是否能够正确执行, check_execution: 测试导入语句和生成的代码一起运行是否能正确执行。

from langsmith.schemas import Example, Run


def check_import(run: Run, example: Example) -> dict:
    imports = run.outputs.get("imports")
    try:
        exec(imports)
        return {"key": "import_check", "score": 1}
    except Exception:
        return {"key": "import_check", "score": 0}


def check_execution(run: Run, example: Example) -> dict:
    imports = run.outputs.get("imports")
    code = run.outputs.get("code")
    try:
        exec(imports + "\n" + code)
        return {"key": "code_execution_check", "score": 1}
    except Exception:
        return {"key": "code_execution_check", "score": 0}

这两个评估函数会返回分数1(成功)或0(失败),用于在LangSmith中评估代码生成的质量。 我们再构建两个预测函数,用于代码生成。

def predict_base_case(example: dict):
    """Context stuffing"""
    solution = code_gen_chain.invoke(
        {"context": concatenated_content, "messages": [("user", example["question"])]}
    )
    return {"imports": solution.imports, "code": solution.code}


def predict_langgraph(example: dict):
    """LangGraph"""
    graph = app.invoke(
        {"messages": [("user", example["question"])], "iterations": 0, "error": ""}
    )
    solution = graph["generation"]
    return {"imports": solution.imports, "code": solution.code}

第一个predict_base_case 用于上下文填充的方式,是较简单的上下文填充方法,直接使用所有可用上下文,通过调用code_gen_chain,将所有连接的内容作为上下文和示例问题一起传,回一个包含导入语句(imports)和代码(code)的字典。 第二个 predict_langgraph 主要用于LangGraph的方式,调用一个应用图(app)并传入示例问题,将迭代次数设为0,错误设为空字符串,图的"generation"输出中获取解决方案,并返回包含导入语句和代码的字典。

开始评估

导入LangSmith的评估模块evaluate,我们创建一个评估器列表code_evalulator,包含两个之前定义的评估函数,然后定义数据集名称 dataset_name 为 "lcel-teacher-eval" ,先测试第一种上下文填充的方式作为预测函数:

from langsmith.evaluation import evaluate

# Evaluator
code_evalulator = [check_import, check_execution]

# Dataset
dataset_name = "lcel-teacher-eval"


# Run base case
try:
    experiment_results_ = evaluate(
        predict_base_case,
        data=dataset_name,
        evaluators=code_evalulator,
        experiment_prefix=f"test-without-langgraph-{expt_llm}",
        max_concurrency=2,
        metadata={
            "llm": expt_llm,
        },
    )
except:
    print("Please setup LangSmith")

我们在这里设置最大并发数为2, 如果评估过程失败,会捕获异常并打印"Please setup LangSmith"提示用户设置LangSmith,这段代码的目的是评估基础方法(不使用LangGraph的方法)在代码生成任务上的性能,并将结果记录在LangSmith平台上以便后续分析。 在这里插入图片描述 然后再使用predict_langgraph函数(之前定义的基于LangGraph的方法)作为预测函数,使用与之前相同的数据集dataset_name,使用相同的评估器列表code_evalulator。

try:
    experiment_results = evaluate(
        predict_langgraph,
        data=dataset_name,
        evaluators=code_evalulator,
        experiment_prefix=f"test-with-langgraph-{expt_llm}-{flag}",
        max_concurrency=2,
        metadata={
            "llm": expt_llm,
            "feedback": flag,
        },
    )
except:
    print("Please setup LangSmith")

这下我们可以进入 LangSmith平台观察性能,我这里贴个链接给大家看一下,smith.langchain.com/public/3266… 在这里插入图片描述 在这里插入图片描述 在这里插入图片描述 我们在这里可以点击进入到每一步查看生成的code的正确率还有导入正确率是否符合我们要求,通过对比发现代码 LangGraph 的表现优于基础传统上下文的调用方式,主要是增加了重试循环提高了性能。但是我们设置反思并未起到帮助,我这里预估是在重试前进行的反射导致了与直接将错误返回给语言模型(LLM)相比出现了性能倒退。