Building a RAG Pipeline for Semantic Code Search: A Developer Diary and Field Notes

为语义代码搜索构建 RAG 管道：开发日志与实战笔记 - JetBrains 博客

JetBrains AI

通过 JetBrains 产品中的 AI 功能增强您的工具

关注

• 订阅：

• RSS RSS

查看更多

智能体 AI

AI

为语义代码搜索构建 RAG 管道：开发日志与实战笔记

Adam Malek

Ashot Kazaryan

第一部分：解析、分块与向量化

不久前，我们开始打造一个理想的语义代码搜索平台：一个 RAG 管道，它能为 LLM 智能体提供来自真实代码库的精确可引用证据，而不是依赖 grep 随机返回的结果。最终的解决方案是 JetBrains Context。我们成功实现了它，将其投入生产环境，并在这个过程中积累了大量经验教训。在这一系列文章中，我们将分享那些希望有人在第一天就告诉我们的关键要点。

毫无疑问，智能体技术是过去十年软件开发领域最大的技术突破。面对看似无法解决的代码复杂性，智能体和前沿模型正在证明其生成看似可靠代码的能力。

然而，随着越来越多的开发流程转向智能体驱动，智能体的效率和生成代码的质量变得越来越重要。问题的关键不在于智能体是否能在足够的时间和 token 资源下完成任务（它当然可以），而在于需要投入多少时间、精力和人工干预才能生成生产级结果。特别是对于大规模代码库，智能体需要花费大量时间搜索与当前功能相关的代码片段并将其纳入上下文。

为什么语义搜索至关重要

在定位正确的代码片段时，智能体会依赖传统的代码搜索工具，如关键字搜索和 grep。然而这些工具存在局限性，它们要求智能体事先知道要搜索的具体文本。例如，当智能体需要查找会话令牌刷新的位置时，无法依赖代码中恰好包含"refresh"这个词。要处理抽象领域的问题，智能体需要通过语义搜索（即根据代码含义进行搜索）的能力。这就是检索增强生成（RAG）技术发挥作用的地方。如果我们能以捕捉代码语义的方式对源代码进行索引，并允许智能体通过自由文本搜索按需检索相关代码片段，就能构建一个充分发挥智能体优势的接口。

从原型到生产环境

在智能体时代，许多伟大的想法在原型阶段都极其简单。但要打造一个经过充分验证的生产级解决方案却远非易事。在这一系列博客文章中，我们希望分享构建有效 RAG 系统所需的关键要素，以及我们在开发 JetBrains Context 过程中走过的弯路。我们将逐步探讨每个阶段，从预处理到存储和智能体集成，提供更深入的技术背景和建议。

本系列文章的第一部分将介绍管道的初始阶段：解析与分块（将原始源文件划分为适当范围的单元），以及向量化（将这些单元转换为支持语义搜索的表示形式）。

解析与分块的精细 AST

解析和分块是优质RAG解决方案中的关键预处理步骤，但常常被忽视。为了让LLM能够对源代码进行嵌入或其他形式的索引，我们必须首先将原始代码行输入模型。这听起来可能微不足道，在小型演示项目中确实可能如此。然而，生产级系统包含成千上万的文件，每个文件可能包含数百甚至上千行代码。更糟糕的是，智能体往往会产生大量代码，进一步扩大代码库规模。每个文件可能包含多个类、字段和方法，它们之间的相关性各不相同。

确定合适的分块大小

即使能够将这些庞大的代码文件完整地嵌入到模型中，这种昂贵的做法最终也会适得其反。由于整个文件被作为一个单元嵌入，搜索结果将返回整个文件内容。这与智能体代码探索和导航的目标背道而驰，因为后者主要关注查找特定函数、符号或代码片段。

另一方面，如果我们采取极端做法，将每行代码单独进行细粒度嵌入，又会面临不同问题。这些孤立的代码行在缺乏上下文时可能语义无关紧要。通用的函数名或注释不值得嵌入，会导致检索结果错误。从某种意义上说，我们会陷入"只见树木不见森林"的困境，智能体将被大量琐碎且无意义的微结果淹没。

因此，找到合适的分块方法至关重要，需要将代码划分为范围适当的组。每个组应包含足够的上下文信息，并能表达常见的语义含义。

为什么固定大小分块存在局限性

"分块"是将要输入智能体的内容划分为多个块的通用技术名称。一种简单的分块方法可能是将大文件按固定行数分割。但采用这种方法会导致语义错误的分组结果，例如将不相关的代码片段组合在一起（如导入语句与函数内容混杂），从而导致检索错误。

为解决这个问题，我们可以利用源文件结构通常较为明确这一特性。以Java为例——导入语句通常位于文件顶部，随后是带有可选文档注释的类定义。类中包含字段和方法，这些元素本身也可能包含文档注释。了解定义类结构的惯例和规则，使我们能够实现更智能的分块，达到上下文信息的合理平衡。

解析与结构感知分块

在过去的26年中，JetBrains团队开发了能够智能适应各种语言特有怪癖、不规则性、惯例和细微差别的解析器。这些解析器与其他工具一起构成了我们内部的JetBrains Code Engine平台，JetBrains Context正是基于该平台开发的。在撰写本文时，JetBrains Context已支持对九种主要语言（Kotlin、Java、Python、JavaScript、TypeScript、C#、PHP、Go 和 Rust）进行解析和结构感知的分块处理。对于其他语言，我们的实现会退回到基于行的简单分割方式，以确保任何语言或文档都能被索引和搜索。

解析器使我们能够将源文件拆分为携带其代表信息的语法节点流——例如注释、空白符、修饰符列表等。随后，分块算法会处理该流并应用逻辑来决定每个分块的范围。基于节点类型和大小及其子节点，算法会做出判断。如果某个节点超过大小阈值但没有子节点，它将退回到更基础的分割策略。

某些语言特有的结构即使超出首选大小也会被保留为单个切片。例如，文档注释、注解、可见性修饰符和关键字会与声明保持在一起；后缀（通常是闭合语法）则会与它们所闭合的结构保持关联。此外，还有一些语言特定的清理操作，例如移除常见且语义无意义的Java注解（如@NotNull或@Override）。

该算法与Zhang等人于2025年提出的cAST有一些相似之处。我们的实现和cAST都保留能容纳的最大语法单元，仅对过大单元进行细分，并将较小的相邻单元组合在一起，以避免通常语义无意义的小块。最大的区别在于，我们为实现加入了更多语言语义，例如将Python装饰器与定义保持在一起，将KDocs与Kotlin声明放在一起等。

分组完成后，会进行分块规范化处理，包括以下步骤：

• 剪除首尾空白符
• 删除空行
• 保留相对缩进的同时移除公共缩进

规范化处理完成后，分块会与包含相对路径等元数据信息一起传递到下一步——嵌入处理。

在完成对源代码的预处理后，我们最终得到了一些文本块，这些文本块在大小和语义检索的分组上应该已经达到了理想状态。接下来的任务是通过一种称为向量化（vectorization）的过程，将这些片段转换为能够支持语义搜索的形式。

通过向量化，嵌入模型（embedding model）会读取一段文本并输出一个固定长度的数字列表（即向量），这相当于一个几千维空间中的点。重要的是，该模型经过训练后，相似语义的文本会聚集在一起。传统搜索可能无法发现这种联系，但在这里，一个刷新缓冲写入操作的函数和一个清空待处理队列的函数即使没有共同关键词，也可能在向量空间中彼此接近。因此，向量之间的距离就成为了相关性的度量标准。查询会被转换为同一空间中的一个位置，而结果则是距离该位置最近的点。

为字节优化：存储优化

对大型代码库进行向量化时，必须同时考虑成本和性能。单个嵌入向量的成本较低，但大型仓库会产生数百万个文本块，进而生成数百万个向量，这些向量需要被存储、保留在内存中，并与每个查询进行比较。一个几千维的32位浮点向量大约占用16KB空间，因此几百万个文本块在任何管理之前就会累积到几十GB的索引。在如此规模下，每个向量的字节分配变得受成本限制，首要问题也迅速从“我们能有多准确？”转变为“每字节能获得多少信息？”换句话说，我们需要找到一种方法，在尽可能保留搜索质量的同时降低成本。

降低向量成本有两种方式。第一种是减少维度。现代嵌入模型经过训练后，向量的前几个维度已经能够独立工作。维度损失会同时应用于多个嵌套的前缀长度，将最粗粒度的结构推入最早维度。这意味着你可以截断向量并重新归一化，它仍然可以检索。另一种方式是保留所有维度，但通过牺牲精度来减少每个维度的字节数，从而减少每个向量的总字节数。

这两种方式相互独立，可以结合使用，这意味着任何存储预算都可以通过不同维度数量和数值精度的组合来满足。真正的问题是，对于相同字节数，哪种组合的检索效果最好。这种权衡并不均衡。假设每个向量的预算是512字节，你可以将其用于保持完整32位精度的128个维度，或者用于每个维度仅用1位的全部4096个维度。两者都正好符合预算，但在测试中你会发现，第二种选项的检索效果明显更好。

为什么维度比精度更重要

要理解这一点，可以将每个维度视为模型针对文本学习提出的一个小问题：是否涉及错误处理？是否涉及网络？是否是测试代码？还有几千个类似但未命名的主题和问题。（实际维度比这更模糊，但这是一个有用的抽象。）

没有单一答案本身意义不大。我们判断两个片段是否相似，是看它们对许多问题的回答是否一致。因此，我们应通过考察问题覆盖范围来评估向量，而非关注答案的精确性。

将所有4,096个维度保留为1位比特，可以大致保留对每个问题的“是”或“否”答案。压缩至128个维度时，仅能对3%的问题给出非常精确的答案，其余信息被丢弃。无论在保留维度上投入多少精度，都无法恢复被丢弃维度所携带的信息。某种意义上，用勾选标记填满的长问卷，胜过用六位小数填满的短问卷。维度是你需要保留的；精度是你能够承受损失且后续更容易补偿的。

因此我们选择保留所有维度，将精度降低至极限，将每个向量压缩为1位比特，这比32位浮点数表示的相同向量缩小了32倍。量化过程本身出人意料地简单：所有大于等于零的分量变为1，所有负值分量变为0，幅值被丢弃：

表示方式的改变随之改变了度量方式。余弦相似度需要我们刚刚丢弃的幅值，因此二进制向量改用汉明距离进行比较，这仅仅是两个比特序列不一致的位置数量。例如比较10110100和10010110，它们在两个位置存在差异，因此距离为2。在完整长度下，计算依然保持简单。一个4,096比特向量以64个64比特字存储，比较两个向量意味着对每个字对执行异或操作，这会在两个向量不一致的位置留下1，然后统计1的数量。CPU对每个字执行这些操作仅需一条指令，因此完整比较仅需约百条指令，而原始浮点数的余弦相似度计算需要数千次乘法。

请注意，度量方式从未是独立决策。我们选择1位精度是为了节省存储空间，一旦每个分量都成为符号位，汉明距离是唯一有意义的比较方式。精度选择决定了度量方式。

二进制量化仍会导致相对于未量化向量召回率下降几个点。在考虑推理代理将使用这些结果后，我们接受了这一代价。为代理提供代码搜索时，找到正确邻域比获得完美排序的前10结果重要得多。当代理询问会话令牌刷新位置时，关键在于相关文件能在前12个结果中出现。最佳片段是排在第二还是第五并无区别，因为代理无论如何都会打开候选结果并阅读它们。在这一流程中，会导致人类界面中明显可见的三结果排序退化，在代理视角下却几乎不可见。

二进制量化的局限性

我们所做的权衡带来了一个更微妙的代价，我们花了更长时间才理解。二进制量化不仅会牺牲准确性，还会压缩相似性分数的*范围*。在全精度向量中，无关的向量对可以得到接近零的分数，而近似重复的向量对可以得到接近一的分数，这种分布范围很宽。符号位的表现则不同。两个完全无关的向量中，大约一半的符号位会偶然一致，而强相关的向量对可能有三分之二的符号位一致。因此，索引中的每个分数，无论是否相关，都会落在那个狭窄的区间内。

排序功能可以承受这种压缩，因为相关结果仍然会高于不相关结果，但阈值判断则无法承受。想象一个无需用户主动请求就能提供相关代码的特性，比如在输入时自动推荐现有实现的面板。它最困难的要求是知道何时保持沉默。为了做出这个判断，它需要“相关”和“不相关”分数之间存在可利用的差距。二进制向量无法提供这样的差距。任何放置在狭窄区间内的截断点，要么会触发所有结果，要么会触发零个结果。因此，当索引需要绝对相关性判断而非相对排序时，我们会保留16位浮点数并承担存储成本。

嵌入范围

尽管索引和搜索使用相同的模型，但这两项任务的差异极大。索引受吞吐量限制，需要异步处理数百万个片段。GPU在每个批次处理约32个片段后就会饱和。另一方面，搜索需要快速响应。如果用户在几秒内得不到结果，他们就会放弃。因此在部署这些模型时，我们根据各自需求进行优化：一个模型用于最大化每秒处理的片段数，另一个模型用于最小化首次结果返回时间。

我们选择了遵循指令的模型，该模型在检索的两个方面进行了刻意的不对称训练。重要的是，这两个方面由截然不同的文本类型表示。查询是自然语言中的简短问题，而文档是代码片段。文档在索引时直接嵌入。查询则被包装在描述检索任务的指令中，例如“根据此搜索查询，找到回答它的代码”，这告诉模型文本所扮演的角色。我们在推理时保留这种结构，因为这是模型学习到的形态。

为了使两方面更容易对齐，我们会在每个代码片段中同时嵌入其文件路径。路径提供了片段本身缺乏的元数据：它所属的模块以及文件的名称。然而在单体仓库中，路径本身却成为一个问题。IntelliJ IDEA 的单体仓库包含超过一百万个文件。其中中位数源文件位于91个字符长的路径下，且有近1万个源文件的路径长度超过150个字符，最长的路径达到218个字符。这还只是在任何检出根目录前缀之前的情况。

保护源代码

我们还考虑了最后一个设计因素。我们非常重视客户隐私和安全问题。公司的源代码通常是其知识产权的核心。将源代码暴露给第三方云模型，甚至暴露给其他公司，都会增加无意中泄露敏感数据的风险，甚至可能训练出其他模型来使用这些数据。

为了确保我们妥善处理这些担忧，我们决定从一开始就遵循以下实践：

• 避免在我们的系统中存储代码：一个代码块包含一个集群引用、项目类型、文件路径、起始和结束偏移量、一个向量的引用以及一个可选的元数据字段。不保存任何内容，也不保存源代码本身的副本。搜索返回的是坐标信息，你看到的代码片段是通过你的本地代码库，利用这些坐标信息在你的机器上组装生成的。服务器只知道某个路径的4,102–4,890字节位置存在相关内容，而不知道具体内容是什么。

• 不使用数据进行训练：JetBrains Context构建的每个代码索引都通过我们在运营的GPU上运行的开源权重嵌入模型进行处理。所有嵌入请求都保留在我们的基础设施内，不会发送给OpenAI、Google或其他任何供应商。因此，我们可以保证没有任何数据会被用于训练其他模型。

这些自我施加的设计限制在检索质量方面没有任何成本。我们用自己的代码检索基准测试了开源权重候选模型与主要供应商的托管嵌入API，结果显示我们的模型表现最佳。目前开源权重嵌入模型已经足够成熟，有趣的技术挑战已经转向了如何向模型提供数据、如何服务模型以及选择保留哪些信息。

作为插曲的总结

在本文中，我们介绍了检索流程的初始阶段：从原始源文件到准备就绪的紧凑向量的旅程。

到目前为止，我们已经拥有数以百万计的二进制向量，并且掌握了生成更多向量的方法。尚未解决的问题包括如何高效存储这些向量、如何构建一个能在毫秒级响应查询的系统、如何持续评估结果以确保做出正确决策，以及如何让代理真正使用我们精心打造的RAG装置。

这些主题以及更多内容将是本系列后续文章的重点，我们将在未来几周内陆续发布。一如既往，欢迎在评论区提问，或分享你在设计RAG解决方案过程中积累的宝贵经验。我们非常期待了解你发现的其他有效且富有创意的方法！在此期间，欢迎访问目前处于公开预览阶段的JetBrains Context，它已包含在你的JetBrains许可证中 😀

下次再见！

• 分享

• Facebook

• Twitter

• LinkedIn

上一篇

我们迈出的第一步：控制AI支出