构建自己的AI编程助手:基于RAG的上下文感知实现方案

简介: 打造智能代码助手,远不止调用API。需构建专为代码设计的RAG系统:基于AST解析保障分块完整性,向量库实现语义检索,结合仓库地图提供全局结构,再通过推理链整合上下文。如此,AI才能真正理解代码,胜任重构、答疑等复杂任务,成为懂你项目的“资深工程师”。

很多人觉得做个AI助手就是调调OpenAI的接口,其实这样智能做出一个通用聊天机器人。

而代码助手需要专门为代码设计的上下文感知的RAG(Retrieval-Augmented Generation)管道,这是因为代码跟普通文本不一样,结构严格,而且不能随便按字符随便进行分割。

一般的代码助手分四块:代码解析把源文件转成AST语法树;向量存储按语义索引代码片段而非关键词匹配;仓库地图给LLM一个全局视角,知道文件结构和类定义在哪;推理层把用户问题、相关代码、仓库结构拼成一个完整的prompt发给模型。

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

自己做代码助手最常见的坑是直接用文本分割器。

比如按1000字符切Python文件很可能把函数拦腰截断。AI拿到后半截没有函数签名根本不知道参数等具体信息。

而正确做法是基于AST分块。tree-sitter是这方面的标准工具,因为Atom和Neovim都在用。它能按逻辑边界比如类或函数来切分代码。

依赖库是tree_sitter和tree_sitter_languages:

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

# 1. Load the Repository  
# We point the loader to our local repo. It automatically handles extensions.  
loader = GenericLoader.from_filesystem(  
    "./my_legacy_project",  
    glob="**/*",  
    suffixes=[".py"],  
    parser=LanguageParser(language=Language.PYTHON, parser_threshold=500)  
)  
documents = loader.load()  

# 2. Split by AST (Abstract Syntax Tree)  
# This ensures we don't break a class or function in the middle.  
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.")  
 # Example output: 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  

# Initialize the Vector DB  
# Ideally, persist this to disk so you don't re-index every run  
db = Chroma.from_documents(  
    texts,   
    OpenAIEmbeddings(model="text-embedding-3-large"),  
    persist_directory="./chroma_db"  
)  

retriever = db.as_retriever(  
    search_type="mmr", # Maximal Marginal Relevance for diversity  
    search_kwargs={"k": 8} # Fetch top 8 relevant snippets  
 )

这里有个需要说明的细节: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)  

# The "Stuff" chain puts all retrieved docs into the context window  
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)  

# Let's test it on that tricky legacy function  
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,但检索结果里没它的内容,要不要看看那个文件?"

总结

我们做自己做代码助手目标不是在补全速度上跟Copilot较劲,而是在于理解层面的提升。比如说内部文档、编码规范、那些只有老员工才知道的遗留模块都可以喂进去。从一个靠猜的AI,变成一个真正懂你代码库的AI。

目录
相关文章
|
21天前
|
前端开发 算法
深度研究Agent架构解析:4种Agent架构介绍及实用Prompt模板
本文系统梳理了深度搜索Agent的主流架构演进:从基础的Planner-Only,到引入评估反馈的双模块设计,再到支持层次化分解的递归式ROMA方案。重点解析了问题拆解与终止判断两大核心挑战,并提供了实用的Prompt模板与优化策略,为构建高效搜索Agent提供清晰路径。
406 10
深度研究Agent架构解析:4种Agent架构介绍及实用Prompt模板
|
22天前
|
人工智能 监控 机器人
别再往一个智能体里塞功能了:6种多智能体模式技术解析与选型指南
单智能体在功能增多时易陷入“指令迷雾”与“工具过载”,导致失效。本文提出6种多智能体架构模式:顺序流水线、并行扇出、层级监督、路由分发、反思迭代、共识投票,类比团队协作,通过分工提升系统稳定性与扩展性,解决复杂任务下的性能衰减问题。
181 6
别再往一个智能体里塞功能了:6种多智能体模式技术解析与选型指南
|
20天前
|
人工智能 自然语言处理 安全
AI就用阿里云!开通阿里云百炼,免费领取超7000万Tokens、100张生图 + 50秒生视频体验!
阿里云百炼是一站式大模型服务平台,支持文本、图像、视频生成及智能体开发。新用户可免费领取超7000万Tokens、100张图像与50秒视频生成额度,无需绑卡,30秒注册即享,助力个人与企业快速构建AI应用。
|
19天前
|
数据采集 监控 数据可视化
快速上手:LangChain + AgentRun 浏览器沙箱极简集成指南
AgentRun Browser Sandbox 是基于云原生函数计算的浏览器沙箱服务,为 AI Agent 提供安全、免运维的浏览器环境。通过 Serverless 架构与 CDP 协议支持,实现网页抓取、自动化操作等能力,并结合 VNC 实时可视化,助力大模型“上网”交互。
365 43
|
18天前
|
开发者 Python
Python异常处理:如何优雅地捕获与传递错误
Python异常处理:如何优雅地捕获与传递错误
192 136
|
23天前
|
人工智能 测试技术 开发者
AI Coding后端开发实战:解锁AI辅助编程新范式
本文系统阐述了AI时代开发者如何高效协作AI Coding工具,强调破除认知误区、构建个人上下文管理体系,并精准判断AI输出质量。通过实战流程与案例,助力开发者实现从编码到架构思维的跃迁,成为人机协同的“超级开发者”。
1503 106
|
18天前
|
机器学习/深度学习 物联网
什么是大模型微调?与预训练有什么区别?
大模型通过海量数据预训练获得通用能力,成本极高;微调则用少量数据调整模型,适应特定任务。轻量级方法如LoRA、PEFT等仅更新部分参数,降低资源消耗,实现高效定制。
152 5
|
2月前
|
数据可视化 安全 测试技术
Anthropic 开源 Bloom:基于 LLM 的自动化行为评估框架
Anthropic推出开源框架Bloom,可自动化评估大语言模型是否阿谀奉承、有政治倾向或绕过监管等行为。不同于传统基准,Bloom基于配置动态生成测试场景,支持多模型、多样化评估,并提供可视化分析,助力模型安全与对齐研究。(237字)
178 12
Anthropic 开源 Bloom:基于 LLM 的自动化行为评估框架