引言

在之前的两篇文章 “用 langgraph 生成 FastAPI 代码：测试工程师的新利器 1” 与 “用 langgraph 生成 FastAPI 代码：测试工程师的新利器 2” 中，我们已经初步探讨了借助先进技术辅助测试工程师进行 FastAPI 代码开发的相关内容。而在这篇文章里，我们将进一步深入，聚焦于如何依据已有的代码库来生成全新的 FastAPI 代码。通过结合 RAG 技术、LangChain 和 Qdrant 等工具，我们能让大语言模型（LLM）更加智能地调用历史代码，从而显著提升开发效率。

在这里插入图片描述

1. 使用 RAG 技术生成模型实体类代码

在之前的探索中，我们已经了解到技术辅助开发的潜力。而 RAG（Retrieval Augmented Generation）技术则是进一步提升开发效率和代码风格一致性的关键。在 FastAPI 开发中，模型实体类代码的编写是一项基础且重要的工作。利用 RAG 技术，我们可以从外部知识库中检索相关信息，结合大模型生成准确、符合风格的代码。

实现步骤
  1. 构建知识库：把已有的代码库当作知识库，涵盖 FastAPI 的路由代码、处理函数代码、模型定义等内容。就像我们在之前文章里提到的那些示例代码一样，它们都可以成为这个知识库的一部分。
  2. 检索相关信息：当需要生成新的模型实体类代码时，从知识库中检索与该模型相关的代码片段。
  3. 结合大模型生成代码：将检索到的代码片段作为上下文，输入到大模型中，生成新的模型实体类代码。
# 示例代码，使用RAG技术生成模型实体类代码
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Qdrant
from langchain.chains.question_answering import load_qa_chain
from langchain.llms import OpenAI

# 初始化向量数据库
embeddings = OpenAIEmbeddings()
vectorstore = Qdrant.from_texts(texts, embeddings, url="http://localhost:6333", collection_name="code_collection")

# 加载问答链
chain = load_qa_chain(OpenAI(), chain_type="stuff")

# 提出问题，例如生成一个用户模型类
question = "生成一个FastAPI的用户模型类"
docs = vectorstore.similarity_search(question)
answer = chain.run(input_documents=docs, question=question)
print(answer)

2. LangChain + Qdrant 实现文件切分入库

为了让大模型能够更好地利用历史代码，我们需要将代码文件进行处理并存储到向量数据库中。在之前的基础上，我们使用 LangChain 进行文件处理，Qdrant 作为向量数据库。

实现步骤
  1. 文件切分：使用 LangChain 的TextSplitter将代码文件切分成小块。
  2. 生成向量：使用OpenAIEmbeddings将每个小块代码转换为向量。
  3. 存储到 Qdrant：将向量存储到 Qdrant 数据库中。
# 示例代码，使用LangChain和Qdrant实现文件切分入库
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Qdrant

# 加载代码文件
loader = TextLoader("path/to/your/codefile.py")
documents = loader.load()

# 切分代码文件
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# 生成向量
embeddings = OpenAIEmbeddings()

# 存储到Qdrant
vectorstore = Qdrant.from_documents(docs, embeddings, url="http://localhost:6333", collection_name="code_collection")

3. 设计 Agent tool，让大模型调用历史代码生成新的代码

在之前文章的基础上，为了让大模型能够更灵活地调用历史代码，我们设计 Agent tool。Agent tool 是一种可以被大模型调用的工具，它能执行特定的任务，例如从向量数据库中检索代码片段。

# 示例代码，设计Agent tool
from langchain.agents import Tool
from langchain.vectorstores import Qdrant
from langchain.embeddings import OpenAIEmbeddings

# 初始化向量数据库
embeddings = OpenAIEmbeddings()
vectorstore = Qdrant.from_texts(texts, embeddings, url="http://localhost:6333", collection_name="code_collection")

# 定义Agent tool
def search_code(query):
    docs = vectorstore.similarity_search(query)
    return "\n".join([doc.page_content for doc in docs])

tools = [
    Tool(
        name="Search Code",
        func=search_code,
        description="Use this tool to search for relevant code snippets from the codebase."
    )
]

4. 通过向量相似度匹配找到最相关的历史代码

在代码生成过程中，复用历史代码是提高效率的重要手段。通过向量相似度匹配，我们可以从代码库中找到最相关的历史代码。

# 示例代码，通过向量相似度匹配找到最相关的历史代码
from langchain.vectorstores import Qdrant
from langchain.embeddings import OpenAIEmbeddings

# 初始化向量数据库
embeddings = OpenAIEmbeddings()
vectorstore = Qdrant.from_texts(texts, embeddings, url="http://localhost:6333", collection_name="code_collection")

# 提出查询
query = "FastAPI的路由代码示例"
docs = vectorstore.similarity_search(query)

# 打印最相关的代码片段
print(docs[0].page_content)

5. 优化代码生成的提示词和工具描述

为了简化工具使用流程，提高工具的易用性和效率，我们需要优化代码生成的提示词和工具描述。

# 示例代码，优化后的提示词和工具描述
prompt = "请生成一个FastAPI的商品模型类，包含id、name和price字段"

tools = [
    Tool(
        name="Search Code",
        func=search_code,
        description="该工具用于从代码库中搜索与查询相关的代码片段。使用时，输入具体的查询内容，如'FastAPI的路由代码示例'。"
    )
]

6. main 函数生成代码

在前面我们已经完成了各项准备工作，接下来我们要编写main函数将这些功能整合起来。以下是一个使用 langgraph.graph 中的 StateGraph、START 和 END，结合 DeepSeek 大模型来生成 FastAPI 代码的完整 main.py 代码：

from typing import List
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict
from langchain_core.messages import SystemMessage, HumanMessage
from llm import DeepSeek  # 假设 llm.py 中定义了 DeepSeek 类
from tools import modelsTool, routesTool, handlersTool  # 假设 tools.py 中定义了相应工具

tools = [modelsTool, routesTool, handlersTool]
tools_names = {tool.name: tool for tool in tools}
llm = DeepSeek().bind_tools(tools)

systemMessage = """
你是一个 FastAPI 开发者，擅长编写高效、清晰的 FastAPI 代码。你将根据需求直接输出代码，不要做任何解释和说明，不要将代码放到 ```python ```中。
"""

models_prompt = """
# 模型
生成一个 User 模型，包含 id（整数类型）、name（字符串类型）和 email（字符串类型）。
"""

routes_prompt = """
# 路由
1. GET /users 获取所有用户列表
2. GET /users/{user_id} 根据用户 ID 获取单个用户信息

# 规则
按照 FastAPI 的路由定义方式编写代码，注释中说明路由的功能。
"""

handlers_prompt = """
# 任务
生成 FastAPI 路由所对应的处理函数代码

# 规则
你只需要生成提供的路由代码对应的处理函数，不需要生成额外代码
处理函数的名称在路由代码的注释中已经给出
如果处理函数需要用到模型，则在模型代码中选择

# 路由代码
{routes}

# 模型代码
{models}
"""


class State(TypedDict):
    main: str
    models: list[str]
    routes: list[str]
    handlers: list[str]


def models_node(state):
    message = llm.invoke([SystemMessage(content=systemMessage), HumanMessage(content=models_prompt)])
    for tool_call in message.tool_calls:
        tool_name = tool_call["name"]
        get_tool = tools_names[tool_name]
        result = get_tool.invoke(tool_call["args"])
        state["models"].append(result)
    return state


def routes_node(state):
    message = llm.invoke([SystemMessage(content=systemMessage), HumanMessage(content=routes_prompt)])
    state["routes"] += [message.content]
    return state


def handlers_node(state):
    prompt = handlers_prompt.format(routes=state["routes"], models=state["models"])
    message = llm.invoke([SystemMessage(content=systemMessage), HumanMessage(content=prompt)])
    state["handlers"] += [message.content]
    return state


def main_node(state):
    prompt = f"""
# 任务
生成 FastAPI 主程序代码

# 规则
1. 导入必要的模块和定义的模型、路由、处理函数
2. 创建 FastAPI 应用实例
3. 包含之前生成的路由和处理函数
4. 启动应用，监听端口 8000

# 模型代码
{state["models"]}
# 路由代码
{state["routes"]}
# 处理函数代码
{state["handlers"]}
    """
    message = llm.invoke([SystemMessage(content=systemMessage), HumanMessage(content=prompt)])
    state["main"] += message.content
    return state


if __name__ == "__main__":
    sg = StateGraph(State)

    sg.add_node("models_node", models_node)
    sg.add_node("routes_node", routes_node)
    sg.add_node("handlers_node", handlers_node)
    sg.add_node("main_node", main_node)

    sg.add_edge(START, "models_node")
    sg.add_edge("models_node", "routes_node")
    sg.add_edge("routes_node", "handlers_node")
    sg.add_edge("handlers_node", "main_node")
    sg.add_edge("main_node", END)
    graph = sg.compile()
    code = graph.invoke({"main": "", "routes": [], "handlers": [], "models": []})

    print(code["models"][0])
    print(code["routes"][0])
    print("\n".join(code["handlers"]))
    print(code["main"])

代码说明：

  1. 导入必要的模块和工具：从 langgraph.graph 导入 StateGraph、START 和 END，从 llm 模块导入 DeepSeek 大模型，从 tools 模块导入相应的工具。
  2. 定义提示信息：包括系统提示 systemMessage，以及生成模型、路由和处理函数的提示信息。
  3. 定义状态类型：使用 TypedDict 定义 State 类型，包含 main、models、routes 和 handlers 字段。
  4. 定义节点函数：models_node 用于生成模型代码，routes_node 用于生成路由代码，handlers_node 用于生成处理函数代码，main_node 用于生成主程序代码。
  5. 构建状态图：使用 StateGraph 构建图，添加节点和边，并编译图。
  6. 执行图并输出代码：调用 graph.invoke 方法执行图，输出生成的模型、路由、处理函数和主程序代码。

总结

通过使用 RAG 技术、LangChain 和 Qdrant 等工具，我们可以实现根据已有代码库生成新的 FastAPI 代码。设计 Agent tool 让大模型能够灵活调用历史代码，通过向量相似度匹配实现历史代码的复用与借鉴，同时优化提示词和工具描述可以提高工具的易用性和效率。而 main 函数则将这些功能整合起来，形成一个完整的代码生成系统。这些技术的结合将大大提升开发效率，让开发人员能够更专注于业务逻辑的实现。

# AI
Logo
天启AI社区

GitCode 天启AI是一款由 GitCode 团队打造的智能助手，基于先进的LLM（大语言模型）与多智能体 Agent 技术构建，致力于为用户提供高效、智能、多模态的创作与开发支持。它不仅支持自然语言对话，还具备处理文件、生成 PPT、撰写分析报告、开发 Web 应用等多项能力，真正做到“一句话，让 Al帮你完成复杂任务”。

更多推荐

cover

【感知机】感知机(perceptron)学习算法例题及详解

avatar 天启AI社区
cover

AI云原生：数智化时代AI Agent加速落地的技术基石与实践指南（THS）

avatar 天启AI社区
cover

【AI论文】设计实验室：通过迭代检测与修正来设计幻灯片

avatar 天启AI社区