AI 编程助手 Agent：RAG 增强下的代码理解和自动补全方案

AI 编程助手 Agent：RAG 增强下的代码理解和自动补全方案

一、深度引言与场景痛点

大家好，我是赵咕咕。

Copilot、Cursor 这些 AI 编程助手大家都用过。体验很分裂对吧？写通用逻辑时补全又快又准，但一碰到公司内部框架、私有库、老的代码规范，补全就变成"随机抽卡"——有时候准得出奇，有时候建议的东西根本没法用。

这不是模型能力的问题。是上下文不足的问题。

通用代码补全的上下文只有你当前打开的文件、光标前几行代码和导入语句。但实际开发中，真正决定"这行代码该怎么写"的信息分散在项目各处：README 里的架构说明、内部 SDK 的源码、同类模块的实现模式、团队的编码规范文档……这些信息都没被喂给模型，它当然只能靠"猜"。

这篇文章，我来聊聊用 RAG（检索增强生成）技术给代码补全装上"项目记忆"，让它真正理解你在写的项目。我会从场景痛点出发，拆解核心原理，给出生产级实现，最后聊聊边界和取舍。

二、底层机制与原理深度剖析

2.1 传统补全 vs RAG 增强补全

传统代码补全的流程极其简单：截取光标前的 N 行代码做 prompt，丢给模型，模型返回补全。这个过程的"视野"只有几百行代码。

RAG 增强补全的思路是：在送 prompt 给模型之前，先用当前代码上下文去项目知识库中检索最相关的信息，拼进 prompt 一起送进去。这样模型的"视野"就拓展到了整个项目。

2.2 核心架构

这里面有三个关键设计：

代码语义分割：不能像处理自然语言一样按长度切分代码。代码的语义单元是函数、类、模块。用 AST 解析按语法边界切分，每个分片是一个完整可编译的函数体，附带它的 docstring、参数签名和类型注解。这样检索出来的结果是一个"可理解的代码片段"，而不是半截函数。

混合检索：纯向量检索有时不够——你正在写一个调用 RedisClient 的代码，向量检索可能返回一堆 Redis 配置代码，但实际你更需要的是项目中其他文件如何调用 RedisClient 的模式。所以需要结构信息辅助——通过 AST 分析出的调用关系图，沿着调用链来召回相关代码。

反馈闭环：用户接受补全 = 正反馈，拒绝 = 负反馈。长期积累下来，检索排名会越来越准。这是一个"越用越聪明"的自增强系统。

2.3 检索策略的关键选择

代码检索和文档检索有三个本质差异：

• 精度优先于召回：代码补全的场景下，返回 3 个高相关片段远好于 10 个半相关片段。因为 prompt 窗口有限，被低质内容占满反而降低补全质量。
• 结构优先于文本：你在函数 A 里调用函数 B，那 B 的签名和实现就是最高相关的内容——比任何语义相似度都重要。
• 时间衰减：3 个月前改过的那段代码，大概率比 1 年前的那段更相关——因为代码库是在持续演化的。

三、生产级代码实现

下面给出一个基于 async/await 的 RAG 增强代码补全引擎实现：

import asyncio
import hashlib
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams

logger = logging.getLogger(__name__)

@dataclass
class CodeChunk:
"""代码语义切片。"""
file_path: str
function_name: str | None = None
class_name: str | None = None
start_line: int = 0
end_line: int = 0
source_code: str = ""
docstring: str = ""
dependencies: list[str] = field(default_factory=list)
chunk_id: str = ""

def __post_init__(self):
if not self.chunk_id:
raw = f"{self.file_path}:{self.function_name or self.class_name}:{self.start_line}"
self.chunk_id = hashlib.sha256(raw.encode()).hexdigest()[:16]

class CodeIndexer:
"""离线阶段：代码索引构建。"""

def __init__(self, embedding_model: str = "text-embedding-3-small"):
self._embeddings = OpenAIEmbeddings(model=embedding_model)
self._client = QdrantClient(path="./qdrant_code_db")

async def build_index(self, project_root: Path) -> None:
"""解析项目代码并构建向量索引。"""
chunks = await self._parse_project(project_root)

if not self._client.collection_exists("code_chunks"):
self._client.create_collection(
collection_name="code_chunks",
vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
)

vector_store = QdrantVectorStore(
client=self._client,
collection_name="code_chunks",
embedding=self._embeddings,
)

# 构建文本表示：函数签名 + docstring + 关键代码片段
texts = []
metadatas = []
for chunk in chunks:
text_repr = (
f"[{chunk.class_name or 'module'}] {chunk.function_name or ''}: "
f"{chunk.docstring}\\n{chunk.source_code[:200]}"
)
texts.append(text_repr)
metadatas.append({
"chunk_id": chunk.chunk_id,
"file_path": chunk.file_path,
"function_name": chunk.function_name or "",
"class_name": chunk.class_name or "",
"start_line": chunk.start_line,
"dependencies": ",".join(chunk.dependencies),
})

# 批量写入
batch_size = 50
for i in range(0, len(texts), batch_size):
batch_texts = texts[i:i + batch_size]
batch_meta = metadatas[i:i + batch_size]
await asyncio.to_thread(
vector_store.add_texts, batch_texts, batch_meta
)
logger.info("已索引 %d/%d 个代码块", min(i + batch_size, len(texts)), len(texts))

async def _parse_project(self, project_root: Path) -> list[CodeChunk]:
"""用 AST 解析项目，按函数/类边界切分代码。"""
import ast
chunks = []

for py_file in project_root.rglob("*.py"):
if "test" in py_file.name or "__pycache__" in str(py_file):
continue
try:
source = py_file.read_text(encoding="utf-8")
tree = ast.parse(source)
relative_path = str(py_file.relative_to(project_root))

for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
docstring = ast.get_docstring(node) or ""
deps = self._extract_calls(node)
chunks.append(CodeChunk(
file_path=relative_path,
function_name=node.name,
class_name=None,
start_line=node.lineno,
end_line=node.end_lineno or node.lineno,
source_code=ast.get_source_segment(source, node) or "",
docstring=docstring,
dependencies=deps,
))
except SyntaxError:
logger.warning("跳过语法错误文件: %s", py_file)
return chunks

@staticmethod
def _extract_calls(node: ast.AST) -> list[str]:
"""提取函数内的所有函数调用名称。"""
calls = set()
for child in ast.walk(node):
if isinstance(child, ast.Call):
if isinstance(child.func, ast.Name):
calls.add(child.func.id)
elif isinstance(child.func, ast.Attribute):
calls.add(child.func.attr)
return sorted(calls)

class CodeCompletionEngine:
"""在线阶段：RAG 增强的代码补全。"""

PROMPT = ChatPromptTemplate.from_messages([
("system", """你是一个代码补全助手。请根据以下来自项目中的相关代码片段，补全给定上下文中的代码。

规则：
1. 优先模仿检索到的代码片段的风格和模式
2. 如果检索结果中有相关函数签名，请直接使用
3. 保持与项目一致的命名规范和错误处理方式
4. 只输出需要补全的代码，不要重复已有的上下文"""),
("human", """【项目中的相关代码】
{retrieved_code}

【当前文件上下文】
{current_context}

请补全以下位置（光标在 |CURSOR| 处）的代码:
{code_before_cursor}|CURSOR|{code_after_cursor}"""),
])

def __init__(self, llm_model: str = "gpt-4o"):
self._llm = ChatOpenAI(model=llm_model, temperature=0.1)
self._client = QdrantClient(path="./qdrant_code_db")
self._embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

async def complete(
self,
code_before: str,
code_after: str = "",
current_file: str = "",
top_k: int = 5,
) -> str:
"""给定光标前后的代码，返回补全建议。"""
try:
# 1. 检索相关代码
retrieved = await self._retrieve(code_before, current_file, top_k)

# 2. 组装 prompt
retrieved_text = "\\n\\n—\\n\\n".join(
f"// {r['file_path']}:{r.get('function_name', '')}\\n{r['source']}"
for r in retrieved
)
context = (
f"// 当前文件: {current_file}\\n"
f"{code_before[-2000:]}" # 截取最近 2000 字符
)

# 3. LLM 推理
chain = self.PROMPT | self._llm | StrOutputParser()
result = await chain.ainvoke({
"retrieved_code": retrieved_text,
"current_context": context,
"code_before_cursor": code_before[-500:],
"code_after_cursor": code_after[:200],
})

return result.strip()

except Exception as e:
logger.error("代码补全失败: %s", e)
# 降级：返回空补全（IDE 侧可以展示错误提示）
return ""

async def _retrieve(
self, query_code: str, current_file: str, top_k: int
) -> list[dict[str, Any]]:
"""混合检索：语义相似 + 文件内优先。"""
if not self._client.collection_exists("code_chunks"):
return []

vector_store = QdrantVectorStore(
client=self._client,
collection_name="code_chunks",
embedding=self._embeddings,
)

try:
# 语义检索
results = await vector_store.asimilarity_search_with_score(
query_code[-1000:],
k=top_k * 2, # 多取一些再过滤
)

scored = []
for doc, score in results:
metadata = doc.metadata or {}

# 同文件加分
file_bonus = 0.15 if metadata.get("file_path") == current_file else 0
final_score = (1 – score) + file_bonus # cosine distance 转相似度

scored.append({
"source": doc.page_content,
"file_path": metadata.get("file_path", ""),
"function_name": metadata.get("function_name", ""),
"score": final_score,
})

# 按最终分数排序
scored.sort(key=lambda x: x["score"], reverse=True)
return scored[:top_k]

except Exception as e:
logger.error("检索失败: %s", e)
return []

async def main():
project_root = Path("./my_project")
indexer = CodeIndexer()
engine = CodeCompletionEngine()

# 离线索引（首次或代码变更后执行）
await indexer.build_index(project_root)

# 在线补全
completion = await engine.complete(
code_before="""
import asyncio
from our_sdk import DatabaseClient

async def fetch_user_orders(user_id: str) -> list[dict]:
client = DatabaseClient()
""",
code_after="""
return orders
""",
current_file="services/order_service.py",
)
print("补全结果:\\n", completion)

if __name__ == "__main__":
asyncio.run(main())

几个值得关注的设计决策：

• asimilarity_search_with_score 拿原始分数，不做简单的 Top-K 截断。这让我们可以在应用层做二次排序——比如同文件加分、最近修改时间加权。纯向量距离只反映语义相似度，反映不了"上下文相关度"。
• AST 级别解析，不用正则。ast.parse 能正确处理装饰器、async 函数、类型注解，不会像正则那样被字符串里的 def 误导。
• 降级策略：检索失败、LLM 调用失败都优雅返回空结果，不会阻塞编辑器。
• 离线索引与在线推理分离：CodeIndexer 是构建时跑的，CodeCompletionEngine 是运行时跑的。两个阶段的依赖完全隔离。

四、边界分析与架构权衡

4.1 RAG 代码补全的适用场景

4.2 延迟 vs 质量

最大的工程权衡是：检索增加了延迟。从向量库检索 + rerank 大约 50-200ms，加上 LLM 推理 500-3000ms。用户对代码补全的延迟容忍度通常在 500ms 以内。

缓解方案：

• 流式输出：不等 LLM 全量推理完，逐 token 推送。用户能感知到第一个字符出现的时间。
• 分级触发：代码块级的补全（函数实现）用 RAG 增强，行级的补全（补完一行）直接走基线模型，不走 RAG。
• 本地小模型：对延迟敏感的 IDE 内场景，用本地部署的 7B 模型替代云端大模型。但这对硬件有要求。

4.3 索引更新的策略

代码库随时在变动。什么时候重建索引？

• IDE 内场景：每次文件保存时增量更新该文件的索引。Qdrant 支持单点 upsert，不需要全量重建。
• CI/CD 场景：每次合并到主分支后触发全量重建，保证索引和最新代码同步。
• 历史版本索引：如果要支持多分支，每个分支独立索引，检索时根据当前分支选择对应索引。

4.4 安全与隐私

代码是企业最敏感的数据之一。把所有代码发给云端 Embedding API（如 OpenAI Embeddings）是一个需要评估的选择。

替代方案：

• 用本地 Embedding 模型（如 BAAI/bge-large-zh-v1.5、thenlper/gte-base），全流程不出本地。
• 用代码脱敏预处理：替换字符串字面量、令牌化变量名，再送云端。但这会损失部分语义信息。

五、总结

RAG 增强的代码补全，本质上是在做一件事：让模型看见它需要看见的东西。

通用代码模型的强大在于它见过全网的代码。但你的项目是它没见过的。RAG 弥补了这个 gap——在推理前，把项目中最相关的代码片段塞进 prompt，模型就能写出行之有效的代码，而不是"看起来像那么回事"的幻觉。

工程落地上，三个关键点要记住：

代码补全的下一个范式一定是"上下文感知"的。RAG 是目前最务实的实现路径。

下一篇预告：LangChain 与 FastAPI 集成，用流式 SSE 把你的 Agent 变成好用的 REST API。