跳到主要内容

代码专用 RAG:为什么通用检索在代码库中会失败

· 阅读需 11 分钟
Tian Pan
Software Engineer

大多数构建 AI 编程助手的团队都会采用与文档检索相同的现成 RAG 流水线:根据 token 数量对源文件进行分块(chunking),对块进行嵌入(embedding),将其存储在向量数据库中,并通过语义相似性进行查询。这种流水线在处理散文(prose)时表现良好。但在处理代码时,它会悄无声息地失败——而且这些失败很难在聚合指标中显现,因为检索到的代码块看起来似乎合情合理,直到模型生成了错误返回类型的代码、调用了签名错误的函数,或者遗漏了调用图中三层之后才存在的依赖项。

问题不在于嵌入模型或向量数据库,而在于分块策略。代码不是散文。它具有结构属性——依赖图、调用链、类型签名、作用域层级——而基于 token 的分块在检索器看到它们之前就破坏了这些属性。修复这个问题需要重新思考在进入嵌入步骤之前如何分解代码。

基于 Token 的分块究竟破坏了什么

当你每隔 512 个 token 切分一次源文件时,切分点会落在 token 计数用完的任何地方。在自然语言文档中,这只是令人烦恼;但在代码中,这在结构上是灾难性的。

在实现过程中被切分的函数会丢失其返回类型。跨块碎片化的类层次结构会丢失继承链。在任意边界断开的调用链会丢失显示一个函数为何调用另一个函数的上下文。与使用位置分离的类型签名在孤立状态下变得无法解释。在条件块中间中断的控制流无法告诉模型该分支是提前返回还是继续向下执行。

这些并非极端情况——在任何现实的代码库中,它们都是常态。一个包含多个方法的 200 行类几乎肯定会被切分。一个负责验证输入、调用下游 API 并格式化响应的服务函数会跨越多个分块边界。模型接收到语义无效的片段,却仍试图对其进行推理。

下游结果是:一个编程助手检索到了看似相关的上下文,但却基于对返回值、函数签名和数据形状的错误假设生成了代码。检索指标看起来不错,但生成的代码却无法编译。

基于 AST 的分块:尊重语法边界

解决方法是让语言自身的结构来确定分块边界。抽象语法树(AST)将源代码表示为语法单元的层级结构——函数、类、控制结构、变量声明——而不是扁平的 token 流。沿 AST 边界进行分块可以确保每个分块在语义上保持完整。

Tree-sitter 是一个拥有 40 多种语言绑定的解析器库,它使这一方案变得切实可行。它能增量生成具体的语法树,并处理部分或格式错误的源文件,这对于并不总是处于整洁可编译状态的现实代码库至关重要。基于 Tree-sitter 构建的分块器可以在函数边界、类定义和顶级声明处切分代码,而不是在任意的 token 计数处切分。

最近的研究将这一过程形式化为一种递归的“先切分再合并”算法。大型 AST 节点(例如具有许多方法的类)在超过大小限制时会被递归分解。相邻的小兄弟节点(例如一系列短小的工具函数)会被合并,以最大化每个块的信息密度。结果是,分块既能满足大小限制,又不会跨越语法边界。

基准测试结果非常显著。在 RepoEval(一个仓库级代码补全基准测试)上,具有 AST 感知的分块比基于行的分块在检索召回率上提高了 4.3 个百分点,在生成 pass@1 上提高了 2.67 到 5.5 个百分点。在测试自然语言到代码任务解决能力的 SWE-bench 上,同样的方法使 pass@1 提高了 2.3 到 2.7 个百分点。这些并不是边际收益——它们在处理数千次补全的生产系统中会产生叠加效应。

一个实现细节至关重要:具有 AST 感知的分块器应该以字符数而非 token 数来衡量分块大小。不同语言的 token 计数差异巨大(Python 很密集;C++ 头文件很庞大),但字符数保持一致。这避免了嵌入模型在看待名义上“大小”相同的块时产生隐藏的不一致性。

调用图感知检索:顺藤摸瓜

修复分块可以让你获得结构上有效的分块。但它并没有解决更难的问题:在语义上与查询相关的代码,在文本上往往与查询并不相似。

假设一个函数 processPayment 调用了 validateCard,后者又调用了 checkFraudScore,再由其调用 fetchUserHistory。关于支付处理的查询可能会通过语义相似性检索出 processPayment,但实际的 bug 或上下文可能存在于 fetchUserHistory 中。基于 token 的检索止步于第一跳。而调用图感知检索会沿着链条追踪。

代码属性图(CPG)正是为此建模的。CPG 将代码库表示为一个图,其中节点是代码实体(函数、类、变量、文件),边编码了结构关系:CALLS(函数调用)、CONTAINS(文件包含)、INHERITS_FROM(继承)、HAS_ARGUMENT(参数关系)、DEPENDS_ON(导入依赖)。与其通过嵌入距离匹配分块,支持 CPG 的检索器可以进行遍历:找到被检索函数调用的所有函数,找到所有调用该函数的函数,找到所有通过其参数流转的类型。

Token 的节省效果令人震惊。一个开源实现报告称,在对一个包含 52 个文件的仓库进行传统检索时,每个查询消耗了大约 123,000 个 token——每个潜在相关的分块都会传递给模型。而通过调用图引导的检索,同一个仓库每个查询仅需约 1,700 个 token:减少了 71 倍,这是通过仅遍历结构相关的子图,而不是淹没上下文窗口来实现的。

除去成本因素,这同样重要。大上下文窗口仅在一定程度上有所帮助。来自处理大型代码库团队的生产数据一致表明,一旦代码库超过大约 4MB,增加更多的上下文窗口并不能提高检索质量——它只会增加噪音。智能结构化检索是这种规模下唯一的出路。

加载中…
References:Let's stay in touch and Follow me for more thoughts and updates