使用LangGraph构建简单的AI Coding Agent，你学会了吗？

【摘要】如何将底层的语言模型技术转化为实用、高效的编程工具，仍然是研究人员和开发者共同面临的挑战。作为LangChain生态系统的重要延伸和创新，LangGraph通过将语言模型与图结构相结合，构建了一个灵活且强大的框架，使得开发智能、可定制的AI编程助手成为可能。本文介绍了LangGraph的工作原理和核心特性，分享了基于Python编写Pytest单元测试的具体动手方法。

【作者】李杰，专注于Java虚拟机技术、云原生技术领域的探索与研究。

在人工智能飞速发展的时代，代码生成已经不再是遥不可及的科幻概念。近年来，随着大型语言模型(Large Language Models，LLMs)的不断演进，我们见证了人工智能在辅助编程领域展现出的巨大潜力。这些模型通过对海量代码和自然语言数据进行学习，逐渐获得了理解和生成编程语言的能力，为程序员提供了智能辅助编码的全新可能。

然而，如何将这些底层的语言模型技术转化为实用、高效的编程工具，仍然是研究人员和开发者共同面临的挑战。毕竟，简单地将模型输出直接作为代码是远远不够的，我们需要更加智能和人性化的人机交互方式，使得AI助手能够真正理解软件开发人员的意图，并提供定制化、符合上下文的编码建议。

在这一背景下，LangGraph凭借其创新的设计理念，为我们带来了全新的视角和解决方案。作为LangChain生态系统中的一颗新星，LangGraph通过将语言模型与图结构相结合，构建了一个灵活且强大的框架，使得开发智能、可定制的AI编程助手成为可能。

LangGraph的核心理念在于将程序代码表示为一种图结构，并将语言模型的输出与这些图结构相关联。通过这种方式，LangGraph可以更好地理解和推理代码的语义和结构，从而生成更加准确、符合上下文的代码建议。同时，LangGraph还提供了丰富的定制选项，允许开发者根据自身需求调整语言模型的行为和输出，大大提高了助手的可用性和实用性。

1. 什么是LangGraph，如何工作？

LangGraph可被视为LangChain生态系统的一次重要延伸和创新，赋予了人工智能编码代理全新的协作和任务管理能力。

虽然LangChain为我们提供了构建可利用多种工具执行编程任务的 AI 代理的框架，但在连贯地跨步骤协调多个链或参与者上，仍然还存在一些局限性。然而，在复杂的实际编程场景中，不同的子任务往往存在着相互依赖和交互的需求，能够灵活地在多个参与者之间协作是创建完成端到端复杂任务的AI代理的关键行为。

正是基于这样的认知，LangGraph框架便应运而生。其将AI代理工作流程建模为一个循环图形结构，其中每个节点代表一个函数或LangChain中可运行的对象，而边则表示节点之间的连接和依赖关系。通过这种图形化的建模方式，LangGraph赋予了AI代理跨步骤协调和任务管理的能力，使其能够更好地应对现实世界中复杂的编程挑战。

LangGraph的核心思想是将编程任务建模为一个图结构，其中节点代表任务的子目标，边则表示子目标之间的依赖关系。通过与语言模型的交互，LangGraph可以动态地生成和修改这个图结构，逐步完成复杂的编程任务。

图：LangGraph核心工作流

与传统的线性编程范式不同，LangGraph的图结构方法赋予了AI助手更强的理解和推理能力。它能够捕捉编程任务中的上下文关联和约束条件，并根据实时反馈动态调整代码生成策略。这种灵活性使得AI助手可以更好地适应不同程序员的编码习惯和偏好，提供高度个性化的辅助体验。

同时，LangGraph还提供了一套丰富的工具和API，支持开发者根据自身需求，定制和扩展AI助手的功能。无论是集成特定领域的知识库，还是添加特定的编程约束和规则，开发者都可以充分发挥创意，打造出符合自身场景的智能编程助手。

2. LangGraph核心特性

LangGraph作为一个灵活且强大的框架，其核心设计理念体现在以下几个主要特点上，具体可参考如下：

1）节点化的任务抽象

在LangGraph中，任何函数或 LangChain中可运行的对象，如工具、代理等，都可以被抽象为一个节点。这种节点化的设计使得复杂的编程任务可以被分解为一系列可组合的子任务单元，每个节点负责执行特定的功能或逻辑，从而提高了整体任务的模块化和可维护性。

2）边缘驱动的流程控制

节点之间通过有向边进行连接，这些边不仅定义了节点之间的执行依赖关系，更重要的是，它们赋予了LangGraph以流程控制和数据流转的能力。通过合理设计边的方向和属性，开发者可以灵活地编排节点的执行顺序、数据传递方式，甚至实现条件分支和循环等高级控制流。

3）有状态图形的任务管理

LangGraph的核心是一种称为“有状态图形”的数据结构。与传统的静态图形不同，有状态图形在执行过程中会动态管理和更新一个状态对象，该对象不仅存储了任务的上下文信息和中间结果，更重要的是，它为图形的动态演化提供了支持。在执行过程中，节点可以根据状态对象的变化，动态地修改图形的拓扑结构，添加或删除节点和边，从而实现自适应的任务调度和管理。

4）知识库和约束条件的集成

LangGraph不仅支持将LangChain中的各种工具和代理集成为节点，还可以将特定领域的知识库、规则库、约束条件等外部资源无缝地融入图形结构中。通过这种集成特性，LangGraph可以赋予AI助手更强的领域知识和上下文理解能力，使其行为更加符合人机交互的期望，提供更加智能化和个性化的编程体验。

5）开放的扩展接口

3. 基于Python编写Pytest单元测试

在软件开发的实践中，单元测试是一个确保代码质量和可维护性的关键环节。然而，手动编写单元测试往往是一项耗时且容易出错的工作，这在一定程度上阻碍了开发效率的提升。正是基于这一认识，我们将在本文中构建一个智能代理，旨在为Python 中带有方法的类自动生成Pytest单元测试，从而为开发者节省宝贵的时间和精力。

在这一创新性的尝试中，我们将借助LangChain生态系统中的LangGraph框架，将整个单元测试生成过程建模为一个有状态的图形结构。每个节点代表了测试生成流程中的一个关键步骤，如解析源代码、生成测试用例、执行测试等，而边则定义了这些步骤之间的依赖关系和数据流转。

通过将复杂的任务分解为可组合的子任务单元，并由专门设计的LangChain代理来执行每个节点的逻辑，我们的智能代理将能够高效地完成端到端的单元测试生成工作。更重要的是，借助LangGraph灵活的图形建模能力，我们可以根据实际需求，动态地调整和优化生成流程，实现自适应的任务调度和管理。

1）配置环境

此步骤中，我们需要构建自己的测试或实验环境，然后将相关依赖库进行安装，具体可参考如下：

# pip install langgraph langchain langchain_openai coloramafrom typing import TypedDict, Listimport coloramaimport os
from langchain_openai import ChatOpenAIfrom langchain_core.messages import SystemMessagefrom langchain_core.messages import HumanMessagefrom langchain_core.runnables import RunnableConfig
from langgraph.graph import StateGraph, ENDfrom langgraph.pregel import GraphRecursionError...

2）设置LLM

完成基础环境的构建后，现在，我们需要指定将在这个项目中使用的语言模型（LLM）。具体选择哪种模型取决于任务的需求和资源的可用性。我们可以选择一些专有且功能强大的模型，如GPT-4、Gemini Ultra或GPT-3.5。这些模型具有强大的自然语言处理能力，适用于各种复杂任务。

此外，还有一些开放访问的模型可供使用，如Mixtral和Llama-2。这些模型开源并且灵活，可以根据具体需求进行调整和优化。在本项目中，由于涉及到编写代码，我们可以选择一些经过微调的专用编码模型，例如DeepSeekCoder-33B或Llama-2编码器，这些模型在处理编程语言和代码生成方面表现出色。

为了进行LLM推理，我们有多个平台可供选择，如Anyscale、Abacus和Together AI。每个平台都有其独特的优势和特性。

在本次项目中，我们将使用Together AI平台来进行DeepSeekCoder模型的推断。Together AI提供了稳定和高效的推理服务，能够满足我们在代码生成和处理方面的需求。

from langchain_openai import ChatOpenAIllm = ChatOpenAI(base_url="https://api.together.xyz/v1",    api_key="your-key",    model="deepseek-ai/deepseek-coder-33b-instruct")

3）定义代理及节点

在LangGraph框架中，有状态图形(Stateful Graph)是一个核心概念，允许我们在执行过程中动态地管理和更新任务的状态对象。这个状态对象不仅存储了任务的上下文信息和中间结果，更重要的是，它为图形的动态演化提供了支持，使得我们的智能代理能够根据状态变化自适应地调整执行策略。

在构建单元测试生成代理的过程中，定义一个合理的AgentState类就显得尤为重要。这个类将作为整个执行流程的状态容器，负责维护和传递代理在各个步骤中所需的关键数据。通过这种方式，不同的节点之间可以高效地共享上下文信息，确保整个任务的连贯性和一致性。

在这里，我们将定义一个代理状态，负责在整个执行过程中跟踪代理的状态。这主要是一个TypedDict类，具有维护代理状态的实体。

class AgentState(TypedDict):    class_source: str    class_methods: List[str]    tests_source: str

在上述的代码中，AgentState类将采用Python中TypedDict的形式，这是一种新的数据类型注释方式，可以为字典类型提供更清晰的类型约束。在AgentState中，我们定义了以下三个关键字段：

class_source：该字段用于存储待测试的Python类的源代码字符串，作为整个流程的输入。

class_methods：这是一个列表，用于存储Python类中定义的所有方法名称。该信息将在后续的解析和测试用例生成步骤中发挥重要作用。

tests_source：最终生成的Pytest单元测试代码将存储在这个字段中，作为整个执行流程的输出。

通过将这些关键数据封装在AgentState中，我们的代理将能够在整个执行过程中高效地访问和操作它们，确保任务的连续性和一致性。

定义好AgentState，我们需要添加节点。那么，节点到底是什么？在LangGraph 中，节点是执行单个操作的函数或任何可运行对象，如Langchain工具。在我们的案例中，我们可以定义几个节点，例如查找类方法的函数，将单元测试推断和更新到状态对象的函数，以及将其写入测试文件的函数。

import_prompt_template = """Here is a path of a file with code: {code_file}.Here is the path of a file with tests: {test_file}.Write a proper import statement for the class in the file."""# Discover the class and its methods.def discover_function(state: AgentState):    assert os.path.exists(code_file)    with open(code_file, "r") as f:        source = f.read()    state["class_source"] = source
    # Get the methods.    methods = []    for line in source.split("\n"):        if "def " in line:            methods.append(line.split("def ")[1].split("(")[0])    state["class_methods"] = methods
    # Generate the import statement and start the code.    import_prompt = import_prompt_template.format(        code_file=code_file,        test_file=test_file    )    message = llm.invoke([HumanMessage(content=import_prompt)]).content    code = extract_code_from_message(message)    state["tests_source"] = code + "\n\n"
    return state

# Add a node to for discovery.workflow.add_node(    "discover",    discover_function)...

在上述代码片段中，我们定义了一个用于发现代码的函数。它从AgentState class_source元素中提取代码，将类解剖成单个方法，并带有提示将其传递给LLM。输出存储在AgentState的tests_source元素中。当然，依据实际的场景需要，我们还可以定义更多的节点信息。

4）执行工作流

构建完成了LangGraph工作流程的各个组件之后，我们进入这个智能单元测试生成代理的最后一步 - 编译和运行整个工作流程。当工作流程成功执行完毕后，invoke方法将返回一个结果对象，其中包含了生成的Pytest 单元测试代码。

# Create the app and run itapp = workflow.compile()inputs = {}config = RunnableConfig(recursion_limit=100)try:    result = app.invoke(inputs, config)    print(result)except GraphRecursionError:    print("Graph recursion limit reached.")

至此 ，我们不仅完成了智能单元测试生成代理的构建，更重要的是，我们充分展示了LangGraph及其生态系统在辅助软件开发中的强大能力和广阔前景。凭借LangGraph灵活的图形建模、LangChain生态系统中丰富的AI组件，我们能够将人工智能技术融入到了软件开发的关键环节中，为开发者带来了全新的智能化编程体验。

作为一款高效的框架，LangGraph主要应用于构建循环状态多行为体代理系统。它弥补了原始LangChain框架中的一些关键不足。作为LangChain的扩展，LangGraph使我们能够利用整个LangChain生态系统中的各种强大功能和工具。

通过使用LangGraph，开发者不仅可以从LangChain丰富的生态系统中受益，还可以更高效地实现复杂的自动化工作流程。这意味着，不论是处理循环状态还是协调多个代理行为，LangGraph都能够提供强大的支持，从而提升系统的整体性能和灵活性。