构建代码助手指南：基于RAG与AST的上下文感知AI编程伴侣

很多人以为开发一个AI编程助手，仅仅意味着调用一下OpenAI等大模型的API。这种做法确实能创建一个通用的聊天机器人，但它远不足以理解你的专属代码库。

一个真正懂行的代码助手，其核心在于一个专为代码设计的上下文感知 RAG（检索增强生成）管道。代码与普通文本不同，它有严格的结构，绝不能简单地按字符数随意分割。

一个合格的代码助手通常包含四个关键模块：代码解析 负责将源文件转换为AST语法树；向量存储 基于语义而非关键词来索引代码片段；仓库地图 为LLM提供全局视角，让其了解项目文件结构和类定义的位置；推理层 则将用户问题、相关代码与仓库结构整合成一个完整的Prompt发送给模型。

代码解析：别用文本分割器

自行构建代码助手时，最常见的错误是直接使用通用的文本分割器。

例如，按1000个字符的固定长度去切割一个Python文件，极有可能把某个函数拦腰截断。AI如果只拿到一个没有函数签名的后半截代码，根本无从知晓参数等关键信息。

正确的做法是基于AST（抽象语法树）进行分块。tree-sitter 是这方面的标准工具，被Atom和Neovim等知名编辑器采用。它能够按照逻辑边界（如类、函数）来切分代码，保持代码结构的完整性。

以下是使用LangChain结合tree-sitter进行代码解析的示例：

from langchain.text_splitter import RecursiveCharacterTextSplitter, Language
from langchain_community.document_loaders.generic import GenericLoader
from langchain_community.document_loaders.parsers import LanguageParser

# 1. 加载代码仓库
# 将加载器指向本地仓库，它会自动处理文件扩展名。
loader = GenericLoader.from_filesystem(
    "./my_legacy_project",
    glob="**/*",
    suffixes=[".py"],
    parser=LanguageParser(language=Language.PYTHON, parser_threshold=500)
)
documents = loader.load()

# 2. 基于AST进行分块
# 这确保了我们不会在类或函数的中间将其切断。
python_splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.PYTHON,
    chunk_size=2000,
    chunk_overlap=200
)

texts = python_splitter.split_documents(documents)

print(f"Processed {len(texts)} semantic code chunks.")
# 示例输出： Processed 452 semantic code chunks.

保持函数和类的完整性至关重要。这样，检索器获取的每一个分块都是一个完整的逻辑单元，而非支离破碎的代码片段。

向量存储方案

代码分块完成后，下一步是将其存储起来以便检索，向量数据库无疑是标准配置。

对于Embedding模型，推荐使用OpenAI的text-embedding-3-large或Voyage AI专为代码优化的模型。这类模型在理解代码语义上表现更佳，能够识别出 def get_users(): 和“获取用户列表”表达的是相同意图。

这里以ChromaDB为例进行演示：

from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings

# 初始化向量数据库
# 理想情况下，应将其持久化到磁盘，避免每次运行时都重新索引。
db = Chroma.from_documents(
    texts,
    OpenAIEmbeddings(model="text-embedding-3-large"),
    persist_directory="./chroma_db"
)

retriever = db.as_retriever(
    search_type="mmr", # 使用最大边际相关性以获得多样性
    search_kwargs={"k": 8} # 获取前8个最相关的代码片段
)

这里有一个细节值得说明：将search_type设置为"mmr"。因为普通的相似度搜索很容易返回五个几乎完全相同的分块，而MMR会强制选取相关但又彼此不同的结果，从而为模型提供更宽广的代码库视野。

上下文构建

仅仅把检索到的代码片段扔给GPT是不够的。模型可能看到了User类的定义，却不知道main.py里是如何实例化它的。它缺少的是项目的全局视角。

因此，解决方法是设计一个系统提示词，引导模型以高级软件工程师的身份来理解代码：

from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

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

# “Stuff”链将所有检索到的文档放入上下文窗口中
prompt = ChatPromptTemplate.from_template("""
You are a Senior Software Engineer assisting with a Python legacy codebase.
Use the following pieces of retrieved context to answer the question.
If the context doesn't contain the answer, say "I don't have enough context."

CONTEXT FROM REPOSITORY:
{context}

USER QUESTION:
{input}

Answer specifically using the class names and variable names found in the context.
""")

combine_docs_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, combine_docs_chain)

# 让我们在一个棘手的遗留函数上测试它
response = rag_chain.invoke({"input": "How do I refactor the PaymentProcessor to use the new AsyncAPI?"})

print(response["answer"])

经过这样的设计，AI便不会凭空捏造不存在的导入，因为它现在可以看到从向量库中检索出的AsyncAPI类定义和PaymentProcessor类。它的回答会变得具体而可靠：“要重构PaymentProcessor，你需要修改_make_request方法。根据上下文，初始化AsyncAPI时需要加上await关键字……”

代码地图：应对大型代码库

上述方案对于中小型项目已经足够有效。但是，当代码库规模膨胀到十万行以上时，仅靠检索几个代码片段还远远不足以覆盖所有上下文。

像Aider、Cursor这类先进工具采用了一种进阶技术，称为“Repo Map”（仓库地图）。其核心思想是将整个代码库压缩成一个轻量级的树状结构，并将其塞入模型的上下文窗口：

src/
  auth/
    login.py:
      - class AuthManager
      - def login(user, pass)
  db/
    models.py:
      - class User

我们的实现思路是：在向模型发送查询之前，先生成一份包含文件名和主要类/函数定义的树状结构大纲，并将其作为系统提示词的一部分附加进去。这样，模型就能说：“地图显示有一个auth_utils.py文件，但当前检索结果中没有它的内容，是否需要查看一下那个文件？” 这极大地增强了AI对大型项目的导航和理解能力。

总结

我们自行构建代码助手的目标，并非要在代码补全速度上与Copilot一较高下，而是在代码理解的深度和个性化上寻求突破。

你可以将内部文档、编码规范、那些只有资深员工才了解的遗留模块逻辑都“喂”给这个助手。它将从一个依赖概率“猜测”的通用AI，转变为一个真正理解你代码库上下文、熟悉你团队习惯的专属智能伙伴。这种深度集成的价值，对于处理遗留系统或复杂项目架构尤其显著。更多关于利用 Python 和 AI 技术解决实际开发问题的讨论，欢迎在云栈社区交流分享。