JetBrains工程师详细记录了为实现精准可溯源的语义代码搜索搭建RAG流水线的全过程,包括解析、分块和向量化各环节的工程取舍。
用 AI 驱动的功能为工具赋能:多个 JetBrains 产品中的 AI 功能
第 1 部分:解析、分块与向量化
一段时间以前,我们着手打造一个尽可能完美的语义代码搜索平台:一个 RAG 管道,让 LLM 智能体能够从真实代码仓库中获取精确的、可引用的证据,而不是仅凭 grep 碰巧搜到的结果。最终的解决方案就是 JetBrains Context。我们让它运行起来了,我们把它投入生产了,一路上也积累了不少踩坑经验。在这一系列文章中,我们将分享那些第一天就该有人告诉我们的东西。
毫无疑问,编码智能体是本 decade 软件开发领域最重大的技术飞跃。智能体和前沿模型正在证明它们在面对看似无法克服的代码复杂性时,能够产出表面可靠代码的能力。
然而,随着越来越多的开发流程转向智能体驱动,智能体的效率以及所生成代码的质量变得愈发重要。问题不在于智能体能否完成任务——只要给予足够的时间和 token 资源,它肯定能——而在于它需要花费多少时间、精力和引导才能生成生产级别的结果。尤其对于大规模代码库,智能体会花费大量时间搜索与其正在开发的功能相关的代码片段,并将它们拉入上下文。
智能体试图定位正确的代码片段时,会求助于传统的代码搜索工具,如关键词搜索和 grep。然而,这些工具的局限性在于,它们要求智能体预先知道要搜索的确切文本。例如,一个智能体寻找会话令牌在哪里被刷新,不能依赖代码中恰好包含了 "refresh" 这个词。智能体要推理抽象领域,需要能够按含义搜索代码(即语义搜索)的能力。这就是检索增强生成(RAG)登场的地方。如果我们能将源代码以捕捉其语义的方式索引,然后允许智能体使用自由文本搜索按需检索相关片段,就创造了一个发挥智能体优势的接口。
与智能体时代的许多伟大想法一样,原生原型实现非常简单。而一个经过良好评估的生产级解决方案绝非如此。在这一系列博文中,我们想分享构建一个有效 RAG 系统所涉及的内容,以及我们在打造自己的 JetBrains Context 过程中走过的弯路。我们将逐一攻克每个阶段,从预处理到存储再到智能体集成,提供一些更技术性的背景和建议。
本系列的第一部分将介绍管道的第一阶段:解析与分块——将原始源文件划分为适当作用域的单元;以及向量化——将这些单元转换为支持语义搜索的表示形式。

解析与分块是优秀 RAG 解决方案中关键的预处理步骤,但往往被忽视。为了让 LLM 能够嵌入或以其他方式索引源代码,我们必须先向其输入原始代码行。这听起来微不足道,对于小型演示项目或许确实如此。然而,生产级系统包含数千个文件,而这些文件又各自跨越数百甚至数千行代码。智能体实际上加剧了这个问题,因为它们往往是多产的写手,进一步膨胀了代码库。每个文件可能包含大量的类、字段和方法,它们之间的关联程度参差不齐。
即便有可能将这些巨大的代码文件整体送入嵌入模型,这种昂贵的操作最终也会适得其反。因为整个文件被嵌入为单一单元,搜索返回的将是整个文件。这与智能体代码探索和导航的目标背道而驰——后者大多关注于找到特定的函数、符号或代码片段。
另一方面,如果我们走向另一个极端,精细地嵌入每一行独立的代码,则会面临另一种问题。这些单独的行在缺乏周围上下文的情况下可能语义无关。一个通用的函数名或注释不值得嵌入,会产生错误的检索结果。从某种意义上说,我们会陷入只见树木不见森林的困境,智能体会被大量通常无关紧要的微结果所淹没。
因此,找到正确的方法将代码分块或分组至关重要。每个分组应包含足够必要的上下文,并代表共同的语义含义。
分块是一种通用技术的名称,指将输入给智能体的内容分割为一组块。对分块采用朴素方法可以是简单地按固定行数将大文件分组。然而,如果我们采用这种方法,会发现结果分组在语义上是错误的。不相关的代码片段会被归为一组,例如 import 语句和函数内容混在一起,导致检索时出错。
为了解决这个问题,我们可以利用每个源文件都有非常清晰定义的结构这一事实。以 Java 为例——import 通常在文件顶部,随后是类定义,其头部前可能有文档注释。该类将包含字段和方法,而这些字段和方法本身也可能有自己的文档注释。了解定义类结构的约定和规则,使我们能够执行更智能的分块,达到周围信息的正确平衡。
在过去的 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 声明附近,等等。
分组之后,执行块归一化,包括:
归一化过程完成后,块会连同元数据一起传递到下一步——嵌入。元数据由相对路径组成,与归一化后的块内容一起被嵌入。
对于嵌入模型的输入应该长什么样,很难给出具体的答案。块大小很重要,但如前所述,更大并不总是更好。此外,与代码一起嵌入的某些元数据可能有用,而某些则可能引入最终降低搜索质量的噪声。
我们选择使用"LLM 即裁判"策略来检查分块,作为评估的一部分。裁判根据一个分块及其源文件,判断边界是否合理。它会寻找意外的人为产物,例如脱节的文档、孤立的闭合语法,或被有意义结构切割成碎片的代码。此外,任何源代码处理流水线的变更也都要经过完整的端到端检索评估。我们将在本系列的下篇文章中回到那个评估流水线。
在完成源代码的预处理之后,我们终于得到了大小合适、语义分组正确的文本分块。下一步是通过一个叫做向量化的过程,将这些碎片转换成本身支持语义搜索的形式。
向量化时,一个嵌入模型读取一段文本,输出一个固定长度的数字列表(一个向量),相当于一个数千维空间中的点。重要的是,这个模型经过训练,使得语义相近的文本在空间中彼此靠近。传统搜索可能会遗漏这种联系,但在这里,一个刷新缓冲区写操作和一个排空待处理队列的函数,尽管没有共同关键词,却可能彼此相邻。因此,向量之间的距离成了关联度的度量。一个查询会被转换为同一空间中的一个位置,结果就是距离该位置最近的所有内容。

字节之争:存储优化
任何对大型代码库进行向量化的尝试都必须考虑成本和性能。单次嵌入很便宜,但一个大型仓库会产生数百万个分块,进而变成数百万个向量,这些向量必须被存储、保存在内存中,并在每个传入查询时进行比较。一个数千维、32 位浮点的向量大约重 16KB,因此数百万个分块的索引会轻松达到数十 GB。在这个规模下,每个向量的字节分配受到成本限制,首要问题迅速从"我们能有多准确?"变成"每个字节能换来什么?"换句话说,我们需要找到一种在尽可能保留搜索质量的同时降低成本的方法。
降低向量成本有两种方法。第一种是减少维度数。现代嵌入模型经过训练,使得向量的前导片段可以独立发挥作用。维度损失同时在多个嵌套前缀长度上应用,将最粗粒度的结构推向最早的那些维度。这意味着你可以将向量截短并重新归一化,它仍然可以检索。或者,你可以保留每个维度,但通过牺牲精度来减少每个向量的字节数。
这两种选项是相互独立的,可以组合使用,这意味着任何存储预算都可以通过维度数量和数值精度的不同组合来满足。真正的问题是,在相同的字节数下,哪种组合检索效果最好。这个权衡远非均衡。假设预算为每个向量 512 字节。你可以将它花在 128 维、保持完整的 32 位精度上,或者花在全部 4,096 维、每维只保留 1 位上。两种方式都恰好满足预算,但测试中你会发现第二种方式检索效果明显更好。
为什么维度比精度更重要
要理解其中的原因,将每个维度看作模型学会对文本提出的一个小问题会很有帮助:这个是关于错误处理的吗?它涉及网络吗?这是测试代码吗?还有几千个类似的主题和问题尚未被命名。(真实的维度比这更模糊,但这是个有用的抽象。)
没有任何单一答案本身有多大意义。当两个分块对许多问题的答案相同时,我们认为它们是相似的。因此,我们应该通过查看问题的覆盖率而不是答案的精确度来评估向量。
在每维 1 位的情况下保留全部 4,096 维,保留了对每个问题的粗略是和否答案。截断到 128 维则精确保留了约 3% 问题的答案,其余的都被丢弃了,而在存活维度上再怎么提高精度,也无法恢复被丢弃维度所携带的信息。从某种意义上说,一份填满勾选的长问卷,胜过一份精确到小数点后六位的短问卷。维度是你想要保留的,精度是你可以牺牲的,而且以后更容易补偿。
因此我们选择保留每个维度,将精度降低到极限,将向量降为每维 1 位,比同一向量用 32 位浮点表示小 32 倍。量化本身出人意料地简单。每个分量大于等于零的变为 1,小于零的变为 0,幅度被丢弃:
改变表示方式也改变了相应的度量方式。余弦相似度需要我们刚刚丢弃的幅度,因此二进制向量改用汉明距离进行比较,即两个比特模式不同的位数。例如,比较 10110100 和 10010110。它们在两位上不同,因此它们之间的距离是 2。在完整长度上,计算同样简单。一个 4,096 位向量存储为 64 个 64 位的字,比较两个向量意味着对每对字进行异或运算,结果中 1 出现的位置就是两个向量不一致的地方,然后再统计 1 的个数。CPU 对每个字都只需一条指令就能完成这些操作,因此一次完整的比较大约需要一百条指令,而原始浮点数的余弦相似度则需要数千次乘法运算。
注意,度量方式从来不是独立的选择。我们选择 1 位精度是为了节省存储空间,而一旦每个分量都变成了符号位,汉明距离就成为唯一有意义的选择。选择精度同时也选择了度量方式。
与未量化向量相比,二进制量化仍然会损失几分召回率。我们在考虑到推理智能体会消费这些结果后接受了这个代价。为 AI 智能体提供输入的代码搜索需要的正确邻域,远比一个完美排序的前 10 名更重要。当智能体询问会话令牌在哪里刷新时,重要的是在第一批十几个结果中能看到相关的少数文件。无论最佳分块排在第二位还是第五位都没有区别,因为智能体会打开候选文件进行阅读。在这个循环中,在为人类设计的三结果 UI 中会明显可见的排名下降,在此处几乎不可见。
二进制量化的局限性
我们做出的权衡有一个更微妙的代价,花了稍长时间才理解。二进制量化不仅牺牲了精度;它还压缩了相似度分数的范围。使用全精度向量,不相关的对可以接近零,而几乎完全相同的对可以接近一,这是一个舒适的宽分布。符号位的表现不同。两个完全无关的向量大约有一半的位纯粹由于偶然而一致,而高度相关的对可能有三分之二的位一致。因此索引中的每个分数,无论相关与否,都落在那个窄带内。
排序在压缩中存活了下来,因为相关结果仍然得分高于不相关结果,但阈值判断就不行了。想象一个功能,在不被询问时主动推荐相关代码,比如在你输入时建议现有实现的某个面板。它最困难的要求是知道何时保持沉默。为了做出这个判断,它需要在"相关"和"不相关"分数之间有一个可用的差距。二进制向量没有留下这个差距。放在那个窄带内的任何截止点,要么对所有内容触发,要么对什么都没有触发。因此,当索引需要绝对相关性判断而不是相对排序时,我们保留 16 位浮点数并支付存储成本。
虽然索引和搜索使用相同的模型,但这两个任务却截然不同。索引受到吞吐量限制,数百万个分块异步处理。GPU 在饱和前大约每批处理 32 个分块。而搜索则需要快速响应。如果不在最多几秒内提供结果,用户就会放弃。因此在部署这些模型时,我们相应地进行优化:一个最大化每分块秒数,另一个最小化首个结果的响应时间。
我们选择了一个指令遵循模型,它在检索两侧之间经过了刻意的不对称训练。值得注意的是,这两侧由截然不同的文本类型表示。查询是自然语言中的简短问题,而文档是一段代码。文档在索引时按原样嵌入,查询则包装了描述检索任务的指令,类似于"给定此搜索查询,找到回答它的代码",这告诉模型文本扮演的是什么角色。我们在整个推理过程中保留这种形式,因为这是模型学习时的形态。
为了让两侧更容易对齐,我们将每个块与其文件路径一起嵌入。路径提供了块本身缺乏的元数据:它位于哪个模块,以及文件是什么。然而在 monorepo 中,路径本身就成为了一个问题。IntelliJ IDEA 的 monorepo 超过一百万个文件。其中间源文件位于距根目录 9 层深的路径中,路径长达 91 个字符,近 10,000 个源文件的路径超过 150 个字符,最长的达到 218 个字符。这还没有算上任何 checkout 根目录的前缀。
这些字符中的大部分用于结构嵌套,对文件没有提供有用的信息。像 src/org/jetbrains/kotlin/idea/k2 这样的段落组合重复了包层级结构,这是编译器需要的,但搜索不需要。同时,那个最长路径末尾的文件只有 24 行代码。如果我们简单地将路径文本原样嵌入到块旁边,会发现路径有时占用的空间比代码本身还多。为了解决这个问题,路径在到达模型之前会被截断,规则是两端都要保留。开头的段落告诉你处于哪个模块,而最后两段——直接的父目录和文件名——告诉你文件是什么。中间部分是可以删除的,而且只有在截断需要时才删除。保留仍能装得下的最长前缀,中间删除的部分用 … 代替,如果即使父目录加文件名仍然太长,就只保留文件名本身。

当用户将搜索范围限定到某个子目录时,同样的规则也适用。显而易见的实现方式是元数据过滤器:正常执行搜索,然后丢弃落在目录之外的结果。我们采取了不同的做法。范围被渲染到查询文本本身中,使用相同的形态、相同的缩写函数和与索引块相同的分隔符。如果一个块以 community/plugins/kotlin 的缩写形式进入索引,那么限定在该目录的查询也以完全相同的格式携带相同的字符串,这样查询向量就会落在它应该匹配的块所在的同一区域。
我们考虑的最后一个设计问题是需要对客户的隐私和安全问题保持关注。公司的源代码往往是其知识产权的核心。将其暴露给第三方云模型,甚至暴露给另一家公司,都会增加无意间暴露敏感数据甚至训练其他模型使用这些数据的风险。
为确保解决这些问题,我们很早就决定遵守若干实践:
避免在我们的系统中存储代码:每个块都包含一个簇引用、一个项目类型、一个文件路径、起始和结束偏移量、一个向量的引用,以及一个可选的元数据字段。没有内容,没有源代码本身的副本。搜索返回的是坐标,而你看到的代码片段是在你的机器上、从你的 checkout 中使用这些坐标组装出来的。服务器只知道在给定路径的字节 4,102–4,890 处存在相关内容,而不知道具体是什么。
不使用数据进行训练:JetBrains Context 构建的每个代码索引都由一个开源权重嵌入模型嵌入,运行在我们运营的 GPU 上。没有嵌入请求离开我们的基础设施——不会发送到 OpenAI,不会发送到 Google,不会发送到任何其他供应商。因此,我们可以保证没有任何数据会被用于训练任何模型。
这些自我施加的设计限制不会带来检索质量方面的代价。我们在自有代码检索基准上对开源权重候选模型与主要提供商的托管嵌入 API 进行了评估,我们的模型表现最佳。开权重嵌入模型现在已经足够好了,有趣的工程已经转移到如何喂养它们、如何服务它们,以及选择保留什么。
在这篇博客文章中,我们涵盖了检索流水线的第一阶段:从原始源文件到可搜索的紧凑向量的旅程。
此时,我们已有数百万个二进制向量以及一种生成更多向量的方法。但我们尚未解决的问题包括:如何高效存储它们、如何创建一个能在毫秒内回答查询的系统、如何持续评估我们的结果以确保我们做出正确的选择,以及如何让 AI 智能体真正使用我们这个崭新的 RAG 工具。
这些主题及更多内容将成为本系列后续部分的主题,我们将在未来几周内发布。与往常一样,欢迎在评论区提出任何问题,或分享您在设计 RAG 解决方案时学到的惨痛教训。我们渴望了解你们发现的不同且有效的方法!同时,欢迎体验 JetBrains Context,目前处于公开预览阶段,已包含在您的 JetBrains 许可证中 😀
感谢您,我们罩着您!
在过去六个月中,JetBrains 的 AI 开发费用增长了大约 10 倍。当成本开始上升时,我们当然注意到了——并意识到我们根本不知道如何系统地控制它们。
这是系列的第 3 部分,我们对编码智能体的公开"token 节省器"插件进行了相同的配对 A/B 基准测试。第 1 部分是 caveman skill(广告称 −65%,实测 −8.5%)。第 2 部分是 rtk(广告称 −60–90%,实测 +7.6%)。我们运行了 80 个配对任务来测试 pony……
今天,我们发布 JetBrains Context,这是一种新的代码库智能层,帮助编码智能体在复杂代码库上更高效地工作并产生更高质量的结果。作为 JetBrains AI for Teams and Organizations 推广的一部分,JetBrains Context 现已开放早期体验。