RAG 系列（二十四）：代码 RAG——让 AI 理解你的代码库

冬奇Lab2026-05-21 21:07

代码和文档的区别

把 Python 文件用 RecursiveCharacterTextSplitter 切 1000 字符的块，再做向量检索------这是最常见的"代码 RAG"实现。问题是它把代码当成了文本：

python 复制代码 
def evaluate_rag(questions, answers, contexts):
    """评估 RAG 系统质量"""
    ...（50行代码）...

按字符切块会：

  • 切断函数，第一半在 chunk A，第二半在 chunk B
  • 丢失函数边界信息（这是 evaluate_rag 函数，不是随机文字）
  • 忽略调用关系（这个函数调用了谁，被谁调用）
  • 失去结构层次（这是 RAGPipeline 类的方法）

代码有三层信息：语义 （这段代码做什么）、结构 （函数/类/模块边界）、调用关系（谁调用谁）。好的代码 RAG 需要把三层都建模。

本文以 llm-in-action 为目标，构建一个能回答"这个函数怎么用"和"给我找所有调用链路"的代码 RAG 系统。

用 AST 切代码，不用字符切

Python 的 ast 模块能把源文件解析成语法树。函数定义是树上的一个节点 ast.FunctionDef，包含完整的开始行、结束行、装饰器列表。用它来分块，切点保证在函数边界上：

python 复制代码 
class _FuncExtractor(ast.NodeVisitor):
    def __init__(self, source: str, rel_path: str):
        self._lines       = source.splitlines()
        self._rel_path    = rel_path
        self._class_stack: list[str] = []
        self.units:        list[CodeUnit] = []

    def visit_ClassDef(self, node: ast.ClassDef):
        # 用栈追踪当前所在类，让方法知道自己的 parent_class
        self._class_stack.append(node.name)
        self.generic_visit(node)
        self._class_stack.pop()

    def _visit_func(self, node):
        # 精确提取函数源码（按行，不按字符）
        src = "\n".join(self._lines[node.lineno - 1 : node.end_lineno])
        unit = CodeUnit(
            name         = node.name,
            kind         = "method" if self._class_stack else "function",
            file         = self._rel_path,
            start_line   = node.lineno,
            end_line     = node.end_lineno,
            source       = src,
            docstring    = ast.get_docstring(node) or "",
            parent_class = self._class_stack[-1] if self._class_stack else "",
            calls        = self._extract_calls(node),   # 提取调用关系
        )
        self.units.append(unit)
        self.generic_visit(node)

    visit_FunctionDef      = _visit_func
    visit_AsyncFunctionDef = _visit_func

调用关系从 ast.Call 节点提取：

python 复制代码 
def _extract_calls(self, node) -> list[str]:
    calls: set[str] = set()
    for child in ast.walk(node):
        if isinstance(child, ast.Call):
            if isinstance(child.func, ast.Name):
                calls.add(child.func.id)           # 直接调用：foo()
            elif isinstance(child.func, ast.Attribute):
                calls.add(child.func.attr)          # 属性调用：self.foo()
    return sorted(calls)

对 llm-in-action 的提取结果

makefile 复制代码 
扫描目录: /mnt/hdd/Database/03_Projects/LLM/llm-in-action
用时: 0.13 秒

Python 文件:  22 个
函数:         188 个（顶层函数）
方法:          37 个（类方法）
代码单元合计: 225 个
涵盖文章目录:  18 个

0.13 秒扫完整个代码库，这是 AST 解析而非运行代码，所以没有任何副作用。

调用图：理解谁调用了谁

提取到函数之间的调用关系后，构建双向邻接表------既能问"X 调用了哪些函数"，也能问"谁调用了 X"：

python 复制代码 
class CallGraph:
    def __init__(self, units: list[CodeUnit]):
        self.callees: dict[str, set[str]] = defaultdict(set)  # caller → called
        self.callers: dict[str, set[str]] = defaultdict(set)  # callee → caller

        known = {u.name for u in units}
        for u in units:
            for callee in u.calls:
                if callee in known:           # 只保留代码库内部的调用
                    self.callees[u.name].add(callee)
                    self.callers[callee].add(u.name)

    def downstream(self, name: str, depth: int = 4) -> list[str]:
        """name 传递调用的所有函数（BFS）"""
        return self._bfs(name, self.callees, depth)

    def upstream(self, name: str, depth: int = 4) -> list[str]:
        """所有传递调用 name 的函数（BFS）"""
        return self._bfs(name, self.callers, depth)

    def shortest_path(self, start: str, end: str) -> Optional[list[str]]:
        """start → end 的最短调用路径"""
        queue: deque[list[str]] = deque([[start]])
        visited: set[str] = {start}
        while queue:
            path = queue.popleft()
            if path[-1] == end:
                return path
            for nxt in self.callees.get(path[-1], set()):
                if nxt not in visited:
                    visited.add(nxt)
                    queue.append(path + [nxt])
        return None

调用图分析结果

makefile 复制代码 
调用图统计:
  有出边的函数:  78 个（调用了其他函数）
  有入边的函数:  92 个（被其他函数调用）
  调用边总数:   168 条

被调用次数最多的函数（代码库的"核心"）：

arduino 复制代码 
get               ← 48 处调用（缓存读取，遍布各 article）
set               ← 10 处调用（缓存写入）
split_documents   ←  5 处调用（文档分块，多个 article 共用）
build_embeddings  ←  4 处调用
query             ←  4 处调用

get 出现 48 次是因为各 article 里的缓存操作（SemanticCache.get、InMemoryCache.get 等）都被识别为对 get 方法的调用------这揭示了 Python 鸭子类型的特点：同名方法合并计数。

调用最广的函数（编排者）：

css 复制代码 
main              → 54 个直接调用
build_self_rag_graph → 6 个直接调用
build_index       → 5 个直接调用
build_ragas_dataset → 5 个直接调用

main 调用 54 个函数------这就是入口函数的特征：它编排整个流程，调用所有子步骤。

调用链路查询

build_self_rag_graph（14-self-rag/self_rag.py）的完整下游调用：

复制代码 
build_self_rag_graph
  ├── make_retrieve_node
  ├── make_filter_node
  ├── make_decide_node
  ├── make_support_node
  ├── make_rag_generate_node
  └── make_direct_generate_node

这正是 Self-RAG 的 StateGraph 节点构建模式：一个工厂函数负责组装所有节点，每个节点是独立的小函数。调用图把这个结构一目了然地展示出来。

build_index（08-ragas-eval/rag_pipeline.py）的下游传递链：

arduino 复制代码 
build_index
  → load_documents
  → build_llm
  → build_embeddings
  → split_documents
  → get（缓存）

这是一个典型的 RAG 初始化序列：加载文档 → 构建 LLM → 构建 Embeddings → 切块 → 缓存。

向量存储：用于代码搜索

代码的向量化有一个工程限制：函数源码可能很长（50--200 行），而嵌入 API 通常有 512 token 限制。

解决方案：分离检索单元和问答上下文。

  • Embedding 内容：函数名 + docstring（短，语义准确，在 token 限制内）
  • 元数据：完整源码（存在 Chroma 的 metadata 字段，用于 LLM 问答时的上下文）
python 复制代码 
sig_line     = u.source.splitlines()[0]
embed_content = f"{full_name}: {u.docstring or sig_line}"[:400]

Document(
    page_content = embed_content,         # 被向量化，用于检索
    metadata = {
        "name":        u.name,
        "file":        u.file,
        "start_line":  u.start_line,
        "source_code": u.source[:2000],   # 不被向量化，用于 LLM 读取
    },
)

问答时，检索找到相关函数，再从 metadata 取完整源码发给 LLM：

python 复制代码 
docs    = vs.similarity_search(question, k=4)
context = "\n\n---\n\n".join(
    d.metadata.get("source_code", d.page_content)[:600] for d in docs
)

语义搜索结果

php 复制代码 
查询: "RAGAS evaluation metrics calculation"
  0.488  RAGPipeline.build_index   (08-ragas-eval/rag_pipeline.py:95)
  0.476  create_ragas_embeddings   (08-ragas-eval/evaluate.py:50)
  0.467  RAGPipeline.query         (08-ragas-eval/rag_pipeline.py:141)

查询: "rate limiting and access control in enterprise RAG"
  0.504  RAGPipeline.__init__      (08-ragas-eval/rag_pipeline.py:78)
  0.497  RateLimiter.__init__      (20-enterprise-rag/enterprise_rag.py:118)

查询: "incremental document indexing with record manager"
  0.296  generate_testset          (08-ragas-eval/generate_qa.py:51)
  0.287  upstream                  (24-code-rag/code_rag.py:157)

查询: "conversational history aware retriever"
  0.400  make_ds                   (18-conversational-rag/conversational_rag.py:428)

第一条（RAGAS）和第二条（企业 RAG 限流）能找到正确文件。第三条（增量更新）没有找到 19-incremental-update------因为那些函数的 docstring 里没有写 "record manager" 这类关键词，只有源码里有。这正是 docstring-only 嵌入策略的局限：函数必须有好的 docstring，搜索才有效。

代码 Embedding 的选型建议

普通文本嵌入模型（BGE、text-embedding-3）对代码的支持是"够用但不好"：能检索 docstring，但理解不了 for i in range(n): acc += arr[i] 是在做累加。

专用代码嵌入模型：

模型 特点
microsoft/codebert-base 代码 + 文档双塔，理解变量名/函数签名
Salesforce/codet5-base 生成式，适合代码补全 + 检索
nomic-ai/nomic-embed-text-v1.5 通用模型但对代码效果好，无 token 限制
voyage-code-2 Voyage AI 的代码专用嵌入，效果最好之一

推荐用法：如果 token 限制不是问题（nomic-embed-text-v1.5 支持 8192 token），直接嵌入完整函数源码，不需要拆分 docstring 和源码。

完整的代码 RAG Pipeline

python 复制代码 
# 构建代码 RAG 系统

# 1. AST 提取所有函数/方法
units = extract_repo(repo_dir)

# 2. 构建调用图
cg = CallGraph(units)

# 3. 向量化（检索用 docstring，问答用 source_code）
vs = build_vectorstore(units, embeddings)

# 三种查询方式
# A: 语义搜索（找相关函数）
hits = vs.similarity_search("embedding caching", k=5)

# B: 调用链路（给定函数名，找所有上下游）
callers  = cg.upstream("build_embeddings")   # → 谁调用了它
callees  = cg.downstream("main")             # → 它调用了谁
path     = cg.shortest_path("main", "get")   # → main 怎么到达 get

# C: LLM 问答（检索 + 生成）
answer = llm_code_qa("如何实现增量更新？", vs, llm)

实验汇总

指标 数值
Python 文件数 22
提取代码单元 225 个（188 函数 + 37 方法）
AST 解析时间 0.13 秒
调用图边数 168 条
向量化时间 5.8 秒
被调用最多的函数 get（48 处）
调用最广的函数 main（54 个直接调用）

完整代码

代码已开源：

github.com/chendongqi/...

核心文件：

  • code_rag.py --- AST 提取、调用图、向量化、搜索、报告

运行方式：

bash 复制代码 
git clone https://github.com/chendongqi/llm-in-action
cd 24-code-rag
cp .env.example .env
pip install -r requirements.txt
python code_rag.py

小结

代码 RAG 和文档 RAG 的核心区别：

文档 RAG 代码 RAG
分块单位 固定大小的文本块 函数/方法（AST 边界）
结构信息 无 类层次、模块层次
调用关系 无 调用图（双向查询）
Embedding 内容 全文 docstring + 签名（或全源码）
查询类型 语义搜索 语义搜索 + 调用链路

三个关键取舍：

  1. AST vs 文本分块：AST 在函数边界切，保留完整单元；文本分块快但破坏结构。生产代码 RAG 用 AST，没有理由不用。
  2. docstring vs 全源码 Embedding：有 token 限制时用 docstring（短且语义集中），但质量依赖 docstring 完备性；有长上下文嵌入模型时直接嵌入全源码。
  3. 调用图 vs 纯向量检索：向量检索找语义相似函数；调用图回答"X 调用了什么"和"谁用了 X"------两者互补，不可替代。

这是 RAG 系列的最后一篇。二十四篇文章走完了从"什么是 RAG"到"如何让 AI 理解代码库"的完整路径。代码全部开源在 llm-in-action，每篇都有可运行的 demo 和真实的评测报告。

参考资料

相关推荐
南屹川1 小时前
【算法】动态规划实战：从入门到精通
人工智能
人工智能培训1 小时前
大模型与传统小模型、传统NLP模型的核心差异解析
人工智能·深度学习·神经网络·机器学习·生成对抗网络
沪漂阿龙1 小时前
面试题详解：智能客服 Agent 系统全栈拆解——Rasa Pro、对话管理、意图识别、GraphRAG、Qwen 与 RAG 优化实战
人工智能·架构
薛定猫AI1 小时前
【深度解析】Gemini Omni 多模态生成与 Agent 化创作工作流：从视频编辑到 UI 生成的技术演进
人工智能·ui·音视频
羊羊小栈1 小时前
AI赋能电力巡检：智能故障预警系统
人工智能·yolo·目标检测·毕业设计·大作业
Python私教2 小时前
视觉 Agent 爬取 vs Playwright 脚本：Browser Use 2026 选型表
人工智能
Python私教2 小时前
Crawlee StagehandCrawler：自然语言点 Load More 的工程化爬虫
人工智能
南屹川2 小时前
【容器化】Docker实战：从入门到生产环境部署
人工智能
海蓝可知天湛2 小时前
Agent&IELTS雅思口语专属语料库
人工智能·github·rag·ielts·skills
热门推荐
01GitHub 镜像站点02Codex 接入 DeepSeek API 完整配置文档03CC-Switch & Claude 基于 Linux 服务器安装使用指南04【踩坑记录 | 第一篇】微软商店无法使用时，如何手动安装 OpenAI Codex？附`.msix`文件系统错误解决方法05几个好用的ip纯净度检测网站06裂开！ChatGPT 居然开始要手机号验证，附详细解决方法07装上就回不去了：CodeGraph 让 AI 编程效率飙升 92%，它到底做了什么？08用了半年 OpenRouter，我换到了 Ofox.ai — 两个 AI API 聚合平台的真实对比09【AI】2026 年具身智能模型和世界模型总结10codex app每次打开重连5次Reconnecting问题解决