构建语义代码搜索RAG流水线:存储、检索与智能体工具实践
第 2 部分:存储、搜索与 Agent 工具
在本系列博客中,我们希望分享过去六个月里构建并推出 Air Context 过程中学到的一些经验。Air Context 是一个面向代码搜索的生产级 RAG 系统。在系列的第一篇文章中,我们介绍了分块、解析和向量化。然而,将代码转化为向量只是整个旅程的一部分。
本文将讨论任何严肃的生产级 RAG 系统都必须解决的下一批问题:如何存储数百万个向量,使索引能够跟上 Git 的变化,以及如何在毫秒级时间内检索到正确的向量。 [LOADING...]
存储
当我们完成代码解析、分块并转化为向量之后,下一个问题就是把这些向量放在哪里。真实世界的代码库可能产生数百万个向量。更复杂的是,代码库并不是静态的。它们不仅经常变化,而且同时存在于数十个分支中;同时,Agent 期望在毫秒级时间内得到答案,并且结果要过滤到正确的仓库、正确的路径,以及仅限用户有权查看的文件。所有这些都需要一个特殊的存储层,能够在给定约束下高效处理搜索操作。
这个存储层的核心是向量索引。向量搜索需要找到与查询向量最接近的向量。逐个将查询与一百万个向量进行比较实在太慢,因此实践中我们会使用近似最近邻(ANN)结构,例如 HNSW 或 IVF。这些结构会跳过大多数比较,返回非常接近最佳匹配的结果,但并不保证一定是最佳匹配。这种速度是有代价的,涉及真正相关结果被返回的比例(也称为召回率)、索引需要多少内存、它应对持续插入和删除的能力,以及它是否能够遵守过滤条件。
除此之外,任何健康的代码库都必须在提交落地时支持多次增量索引更新。为了确保这一点能够大规模实现,设计必须避免为共享同一份代码的每个分支中的每个块都保留单独副本,并且必须保持向量与描述它们的元数据之间的一致性。这些能力都不是嵌入模型提供的。没有强大的存储方案,我们精心提取的向量根本无法使用。
近似搜索的作用
近似搜索算法会牺牲一定的召回率,但能在大型数据集上良好运行。
我们发现,一种混合方法效果非常好:在初始 ANN 搜索之后,再在更有限的向量空间上执行精确搜索。
换句话说,先用低成本的第一轮搜索将整个语料库缩小到一小批可能候选,然后再执行精确比较就变得可负担。近似索引负责规模,精确搜索负责最后一英里的精度。这为我们提供了另一个微调搜索的杠杆,因为我们不仅考虑使用哪些算法,还考虑如何将它们协同使用。 [LOADING...]
过滤式 ANN:细节
用于语义相似度的向量搜索无法感知更严格的搜索约束,例如将结果限制到某个仓库、分支或特定子文件夹。此外,安全和权限约束,例如对可查询数据类型的限制,必须被严格执行。这类元数据过滤需要一种确定性的机制,并且该机制要优先于任何搜索算法。
将过滤器与 ANN 索引结合确实很难,因为这两种机制会相互对抗。HNSW 图的构建假设你可以在邻居之间自由遍历;一旦将遍历限制为满足谓词的节点,图就可能变得不连通,遍历也可能错过有效候选。有三种主要策略可以解决这个问题:
-
预过滤。 这意味着先应用元数据谓词,然后只在幸存集合中搜索。相对于过滤器而言,它提供 100% 的召回率。然而,它会在过滤后的子集上失去 ANN 索引,实际上退化为精确搜索。这种策略最适合高度具体的过滤器,例如单个小型仓库或狭窄的子文件夹。但当过滤器仍然留下大量对象需要搜索时,它就会变得昂贵。
-
元数据与向量索引捆绑。 最强的通用方法是在存储中让过滤成为索引本身的一等公民。换句话说,文件夹权限或仓库等参数是可搜索的,并且会影响向量距离。Qdrant 和 turbopuffer 就是这样构建的。你可以保留 ANN 的大部分速度,并且遵守过滤器,而不必退回到暴力搜索。对于具有真实 ACL 要求的生产级代码库 RAG 来说,这会是一个平衡的选择。
-
后过滤。 先运行 ANN 搜索,然后丢弃不满足谓词的结果。问题当然在于,由于 top-K 候选并不是在考虑该过滤器的情况下选出的,其中大多数或全部最终可能都是无效的,从而留下一个很短、质量很低的列表。pgvector 的迭代扫描模式通过不断获取更多候选,直到有足够多的候选通过过滤器来缓解这个问题,但它会带来延迟成本,而且可能不可预测。后过滤被认为是三种策略中最弱的一种。
跟上增量代码变更
在处理大型代码仓库时,RAG 系统无法承受每次变更都重新索引整个代码库。因此,为了实现可行的索引性能和数据存储限制,索引必须增量执行。
理想情况下,我们希望每个仓库快照都能确定性地聚类成一棵与其目录结构相对应的 Merkle 树。新修订版本会自顶向下进行差异比较,完全未变的子树会原样链接到新快照中,无需重新索引。嵌入还会按块内容哈希进一步去重,因此在被修改的聚类内部,我们也只会索引发生变化的代码。
有了增量机制后,索引更新和删除会以小规模、受锁边界约束的事务逐步执行,并行度有限;而向量索引构建等繁重维护工作则被转移到后台任务,避免数据库 CPU 出现瓶颈。 [LOADING...]
然而,所有这些处理和业务逻辑都会消耗时间和处理资源。构建 HNSW 向量索引是我们数据库流水线中的主要 CPU 瓶颈之一。为了防止系统过载,我们将向量摄入与索引构建解耦。向量先写入,只有当仓库快照完全存储到数据库后,才会构建 HNSW 索引。随后,构建会在严格控制的并行度下异步运行。为了保持 CPU 使用效率,我们使用非阻塞并发索引创建,使每个分片一次最多只能构建一个索引。这样可以让数据库实例保持稳定,因为快照只有在其索引准备就绪后才可搜索。
为了应对单个数据库实例在 CPU、内存和元数据容量方面的限制,我们的架构会在多个分片之间水平扩展,每个分片都有自己的主实例和只读副本。搜索查询由部署在不同地理区域的只读副本提供服务,以保证低延迟。此外,为了维护数据隔离,每个客户都被固定到恰好一个分片,并通过一个简单的应用侧路由层,根据客户标识符将每个事务路由到正确的分片。
利用本地索引处理本地变更
在典型的 SDLC 中,代码首先编写在本地笔记本电脑或台式机上,随后被推送到远程分支,并与代码库的其余部分合并。在这些早期开发阶段,这些代码对单个开发者高度相关,但对其他在自己分支上工作的开发者则不那么相关。在开发者机器上创建本地索引来支持推送前的开发阶段,有助于减轻中央搜索索引的负载和磁盘需求,并确保搜索结果具有相关性。
Air Context 索引可以完全驻留在设备本地,以支持这一用例。对于本地向量存储和搜索,我们采用了 USearch,这是一个紧凑、高性能的向量引擎。向量与元数据是分离的,快照、聚类和块都保存在向量索引旁边的专用元数据存储中。由于这些元数据遵循与服务端完全相同的模型,整个增量流水线可以在本地继续工作。未变化的聚类会跨修订版本重新链接,相同的块则共享同一个存储向量。
在服务端和本地存储之间共享同一套元数据模型,使我们能够支持级联索引,从而允许代码从本地过渡到共享状态。搜索因此变成一个两阶段过程。首先对远程索引运行搜索,然后再对本地索引运行搜索,并使用重排器将两个来源的顶级候选合并成一个结果列表。
通过这种方式,我们既保持了一个始终与本地文件树同步的灵活本地索引,又通过为远程分支共享远程索引来优化存储和推理。此外,当代码被推送到中央源代码库时,索引会透明地从本地过渡到远程,从而为我们节省宝贵时间。
检索后优化
后处理:将候选转化为答案
上文描述的最近邻搜索并不会产生最终结果。它产生的是一个候选列表:向量最接近查询的那些块。虽然这是一个很好的起点,包含大量有用数据,但它并不是最终交付给 Agent 或用户的产品。ANN 结果可能包含近似重复项、同一函数的多个片段,或者因为错误原因而得分很高的块,例如它们与查询共享了几个罕见词。
后处理涵盖将原始候选列表转化为 Agent 所读取文本的全部过程。在我们的流水线中,这包括去除重复项、将属于同一文件的片段分组、决定展示多少结果才值得,以及最重要的重排。与搜索本身相比,这些步骤的每一步成本都很低,因为每一步只处理几十个候选,而不是整个代码库。
重排
重排可以最好地描述为:在把候选交出去之前,再更仔细地看它们一眼。我们的第一轮检索必须很快,因为它通过两个独立计算的向量将每个块与查询进行比较。重排器则可以承受更慢的速度。它的作用是并排查看查询和候选,判断它们实际匹配得有多好。虽然我们无法在庞大索引上使用 LLM 来判断结果相关性,但当结果集较为有限时,这确实是一个可行选项。最终,重排会给我们一个重新排序的列表,其中最佳匹配被移到顶部。
除了这种先做邻居搜索再重排的方法之外,我们还尝试了多种其他检索配置。其中一些在搜索质量上带来了明显提升,但代价是资源使用显著增加、搜索速度变慢。一个例子是 JointRank,这是 ICTIR '25 会议上提出的一种方法,它要求语言模型在一次处理中对许多候选进行排序:将候选拆分成重叠的组,并行对每组排序,再将部分排序结果合并为一个整体排序。
Agent 式搜索
为了进一步提升搜索质量,我们探索了使用 Agent 式搜索工具。与其进行单次检索,我们可以选择使用专门的代码检索 Agent。这个 Agent 可以编排对多个搜索和读取工具的调用,包括基于嵌入的检索。它可以根据查询本身判断应调用哪些工具,并使用不同的阈值。Agent 会自行决定何时搜索、何时打开文件,以及何时已经看到足够的信息。我们限制它可以执行的步数,以避免探索过程失控并耗尽预算。
这个检索 Agent 可以实现为一个子 Agent,这是近年来为保持隔离上下文的辅助 Agent 所建立的一种常见模式。当子 Agent 执行读取时,只有它的结论会返回,因此主 Agent 的上下文可以保持小而专注。例如,Claude Code 允许用户定义自定义子 Agent,主 Agent 可以像调用其他工具一样调用它们。
子 Agent 还可以运行在比编写代码的模型更便宜或更快的模型上,并且可以被限制为只能使用只读工具,以防止它意外更改任何内容。此外,多个子 Agent 可以同时运行,处理相互独立的问题。对于我们的搜索工具来说,这意味着主 Agent 在查找代码上花费的 token 和时间更少,而把更多精力放在它被赋予的任务上。
结果格式化与前置加载片段
后处理阶段最后一个需要考虑的方面,是我们将结果交给 Agent 时所使用的格式。我们的目标是通过防止 Agent 发起额外调用,同时确保它拥有足够数据来评估结果对其原始目标的相关性,从而为 Agent 节省时间。不过,这需要谨慎进行,因为我们预先放入响应的数据越多,Agent 上下文就越膨胀,也可能受到污染。我们采用的策略是让详细程度与每个结果的评估相关性相匹配。
对于相关性分数超过预定义阈值的候选,其结果会包含更多信息,包括代码块本身、行号以及额外元数据。对于估计相关性较低的候选,我们只提供一些基本信息,例如文件名,可能还会附带非常部分的片段。这样可以让 Agent 的注意力集中在可能正确的匹配上,同时仍然提示它还可以去哪里查找。
从理论到实践
在本文中,我们讨论了流水线的第二部分,涵盖存储以及搜索和检索。在存储方面,近似索引承担了主要工作,并在小型候选集上执行精确搜索。仓库、分支和权限过滤器位于索引内部,而不是围绕索引之外。每个快照的 Merkle 树视图通过复用所有未变化的内容,让索引能够跟上 Git;开发者机器上的本地索引则在代码被推送之前承担相应工作。
在检索方面,重排器会更仔细地查看候选,判断它们的相关性。当单次搜索不够时,专门的子 Agent 能够进行更细致的探索。最后,结果格式可以根据我们对每个匹配项的置信度,向 Agent 提供相应详细程度的信息,从而让 Agent 的上下文保持小而专注。
还有一个关键问题我们尚未回答:这一切真的能帮助 Agent 吗?尽管我们做了所有这些优化、算法以及对规模的考量,这套 RAG 机制是否真的对 Agent 性能产生了可衡量的影响?毕竟,Agent,尤其是使用前沿模型时,已经越来越擅长在大型代码库中导航。既然如此,语义搜索工具的附加价值是什么?
这些问题从我们旅程一开始就萦绕在我们心头,非常重要。在下一篇文章中,我们将展示如何将 Air Context 接入 Agent、将其转化为产品需要付出什么,以及我们如何评估它——从检索基准到端到端 Agent 运行。剧透一下:在那些运行中,使用 Air Context 的 Agent 解决任务所需轮次最多减少了 70%。
一如既往,欢迎在评论区提问,或分享你自己来之不易的经验。如果你想在自己的代码上尝试,Air Context 已进入公开预览,并且已包含在你的 JetBrains AI 订阅中。
下次见!