构建面向语义代码搜索的RAG流水线:开发者日志与实战笔记
第 1 部分:解析、分块与向量化
不久前,我们着手构建自己能打造的最好的语义代码搜索平台:一个 RAG 流水线,为 LLM 智能体提供来自真实代码仓库的精确、可引用证据,而不是让 grep 碰巧搜到什么就用什么。最终的解决方案是 JetBrains Context。我们让它跑了起来,把它带进了生产环境,也一路积累了不少伤疤。在本系列文章中,我们将分享那些希望一开始就有人告诉我们的经验。
编码智能体无疑是我们这十年软件开发领域最大的技术飞跃。面对看似难以逾越的代码复杂性,智能体和前沿模型正在证明它们有能力产出表面上可靠的代码。
然而,随着越来越多的开发流程由智能体驱动,智能体的效率以及所产出代码的质量变得越来越重要。问题与其说在于智能体能否完成任务——只要给予足够的时间和 token 资源,它当然能完成——不如说在于它需要多少时间、精力和引导才能生成生产级成果。尤其是对于大规模代码库,智能体会花费大量时间搜索与其正在开发的功能相关的代码片段,并将它们拉入上下文。
为什么语义搜索很重要
为了定位正确的代码片段,智能体会求助于关键词搜索和 grep 等传统代码搜索工具。然而,这些工具有其局限:它们要求智能体预先知道要搜索的确切文本。例如,智能体要查找会话令牌在何处刷新,不能指望代码中好心地包含“refresh”这个词。为了对抽象领域进行推理,智能体需要按含义搜索代码的能力,也就是语义搜索。这正是检索增强生成(RAG)发挥作用的地方。如果我们能够以捕捉语义的方式索引源代码,然后允许智能体通过自由文本搜索按需检索相关片段,我们就创建了一个能发挥智能体优势的接口。
从原型到生产
与智能体时代的许多伟大想法一样,原生的原型实现极其简单。但经过充分评估的生产级解决方案绝非如此。在本系列博客文章中,我们希望分享构建有效 RAG 系统所涉及的内容,以及我们在打造自己的 JetBrains Context 旅程中走过的弯路。我们将逐一讨论每个阶段,从预处理到存储和智能体集成,并提供更多技术背景和建议。
本系列的第一部分将涵盖流水线的初始阶段:解析与分块——将原始源文件划分为范围恰当的单元;以及向量化——将这些单元转换为支持语义搜索的表示。 [LOADING...]
解析与分块的精妙 AST
解析与分块是优秀 RAG 解决方案中关键的预处理步骤,但常常被忽视。为了让 LLM 能够嵌入或以其他方式索引源代码,我们必须先将原始代码行提供给它。这听起来可能微不足道,对于小规模演示项目来说大概确实如此。然而,生产级系统包含数千个文件,而每个文件又可能包含数百甚至数千行。如果说有什么影响,那就是智能体让问题雪上加霜,因为它们往往是多产的写手,会进一步使代码库膨胀。每个文件都可能包含大量类、字段和方法,而它们之间的关联程度各不相同。
找到合适的分块大小
即使能够将这些庞大的代码文件完整放入嵌入模型,这种代价高昂的做法最终也会适得其反。因为整个文件被嵌入为单个单元,搜索会返回整个文件。这与智能体代码探索和导航的目标背道而驰,因为后者主要关注查找特定函数、符号或代码片段。
另一方面,如果我们走向另一个极端,对每一行代码分别进行细粒度嵌入,就会面临另一种问题。如果没有周围上下文,这些单独的行在语义上可能无足轻重。泛泛的函数名或注释不值得嵌入,并且会产生错误的检索结果。从某种意义上说,我们会只见树木不见森林,而智能体也会被大量通常并不重要的微小结果淹没。
因此,至关重要的是找到正确的方法,将代码分块或划分为范围恰当的组。每个组都应包含足够的必要上下文,并代表共同的语义含义。
固定大小分块为何不够
分块(chunking)是一种技术的通称:将要提供给智能体的内容划分成一组块。一种朴素的分块方法可以简单地将大文件按固定行数拆分成组。然而,如果采用这种方法,我们会发现得到的分组在语义上是错误的。不相关的代码片段会被归到一起,例如一条 import 语句和某段函数内容,从而导致检索时出错。
为了解决这个问题,我们可以利用一个事实:每个源文件都有相当明确的结构。以 Java 为例——import 通常位于文件顶部,随后是类定义,类头之前可能还有可选的文档注释。类中会包含字段和方法,而它们也可能有自己的文档注释。了解定义类结构的约定和规则,使我们能够进行更智能的分块,并实现周围信息的恰当平衡。
解析与结构感知分块
在过去 26 年里,我们 JetBrains 开发了足够智能的解析器,能够适应特定语言的各种怪癖、不规则之处、约定和细微差别。这些解析器与其他工具一起,构成了我们的内部平台 JetBrains Code Engine,JetBrains Context 正是在该平台上开发。在撰写本文时,JetBrains Context 支持九种主要语言的解析和结构感知分块:Kotlin、Java、Python、JavaScript、TypeScript、C#、PHP、Go 和 Rust。对于所有其他语言,我们的实现会简单地回退到朴素的、基于行的拆分,以确保任何语言或文档都能被索引和搜索。
解析器允许我们将源文件拆分为语法节点流,这些节点携带关于其代表内容的信息——注释、空白、修饰符列表等。分块算法随后消费该流,并应用逻辑来决定给定分块的范围。算法会根据节点的类型和大小以及其后代做出决定。如果某个节点超过大小阈值但没有子节点,就会回退到更原始的分割策略。
某些特定语言的结构即使超过首选大小,也会保留为单个切片。文档、注解、可见性修饰符和关键字等前缀会与声明保留在一起;后缀(通常是结束语法)则保持与其所闭合的结构相关联。此外还有一些特定语言的清理,例如会移除常见且语义上无意义的 Java 注解,如 @NotNull 或 @Override。
[LOADING...]
该算法与 Zhang 等人于 2025 年提出的 cAST 有一些相似之处。我们的实现与 cAST 都会保留能够容纳的最大语法单元,只细分过大的单元,并合并较小的相邻单元,以避免通常没有语义意义的微小分块。最大的区别在于,我们在实现中编码了更多语言语义,例如将 Python 装饰器与定义保留在一起,将 KDocs 紧邻 Kotlin 声明,等等。
分组之后,会执行分块归一化,包括:
- 去除首尾空白
- 删除空行
- 移除公共缩进,同时保留相对缩进
归一化流程之后,分块会连同元数据一起传递到下一步——嵌入,元数据由相对路径组成,该路径会与归一化后的分块内容一起被嵌入。
评估分块质量
很难具体回答嵌入模型的输入应该是什么样子。分块大小很重要,但如前所述,越大并不总是越好。此外,与代码一起嵌入的某些元数据可能有用,而另一些则可能引入噪声,最终降低搜索质量。
我们选择使用 LLM-as-a-judge 策略,在评估过程中检查分块。评判者会使用分块和源文件,判断边界是否合理。它会寻找意外产物,例如脱离的文档、孤立的结束语法,或被从有意义结构中截断的代码片段。此外,对源代码处理流水线的任何更改也会经过完整的端到端检索评估。我们将在本系列的下一部分中回到该评估流水线。## 向量化
在对源代码进行预处理之后,我们终于得到了文本块,希望它们的尺寸刚刚好,并且为语义检索进行了正确的分组。接下来的任务是以一种以后能够支持语义搜索的方式转换这些片段,这一过程称为向量化。
通过向量化,嵌入模型会读取一段文本,并输出一个固定长度的数字列表(即向量),这相当于数千维空间中的一个点。重要的是,模型经过训练,使得含义相似的文本彼此靠近。传统搜索可能会错过这种联系,但在这里,一个刷新缓冲写操作的函数和一个清空待处理队列的函数,尽管没有共同关键词,最终也可能彼此接近。因此,向量之间的距离就成为了相关性的度量。查询会被转换为同一空间中的一个位置,而结果就是离它最近的那些内容。 [LOADING...]
精打细算每一字节:为存储而优化
任何对大型代码库进行向量化的尝试,都必须同时考虑成本和性能。单个嵌入很便宜,但大型代码仓库会产生数百万个块,进而变成数百万个向量,这些向量必须被存储、驻留内存,并与每个到来的查询进行比较。一个数千维、使用32位浮点数的向量大约重16 KB,因此几百万个块在计入任何额外管理开销之前,索引就会达到几十 GB。在这种规模下,每个向量分配多少字节会受到成本限制,首要问题也迅速从“我们能有多准?”变成“每字节能换来什么?”换句话说,我们需要找到一种方法,在尽可能保持搜索质量的同时降低成本。
降低向量成本有两种方式。第一种是保留更少的维度。现代嵌入模型的训练方式,使得向量的前一段切片本身就能发挥作用。维度损失会同时施加在多个嵌套的前缀长度上,把最粗粒度的结构推向最早的维度中。这意味着你可以截短向量并重新归一化,它仍然能够检索。或者,你可以保留所有维度,但在每个维度上花费更少,方法是牺牲精度,从而为每个向量保留更少的字节。
这两种方案彼此独立,也可以组合,这意味着任何存储预算都可以通过维度数量和数值精度的不同组合来满足。真正的问题是,在相同字节数下,哪种组合的检索效果最好。这种取舍远非均衡。假设每个向量的预算是512字节。你可以把它花在128个维度、保持完整32位精度上,或者花在全部4,096个维度、每个维度只保留1位上。两者都恰好符合预算,但在测试中你会发现,第二种方案的检索效果明显更好。
为什么维度比精度更重要
要理解原因,可以把每个维度看作模型学会针对文本提出的一个小问题:这涉及错误处理吗?它会触及网络吗?它是测试代码吗?还有几千个类似但尚未命名的主题和问题。(真实维度比这更模糊,但这是一个有用的抽象。)
任何单一答案本身都没有多大意义。当两个块对其中许多问题的答案相同时,我们就认为它们相似。因此,我们评估向量时,应该看它覆盖了多少问题,而不是答案有多精确。
将全部4,096个维度各保留一位,可以保留对每个问题粗略的是/否答案。截断到128个维度,可以保留对3%问题的非常精确答案,却把其余问题丢掉;而且无论保留的维度精度多高,都无法恢复被丢弃维度所携带的信息。从某种意义上说,一份用勾选方式填完的长问卷,胜过一份只填到小数点后六位的短问卷。维度是你想要保留的;精度是你可以承受损失、并且之后更容易补偿的。
所以我们选择保留所有维度,并把精度压缩到极限,将每个向量降到1位,这比同样向量的32位浮点表示小32倍。量化本身出奇地简单。每个大于等于零的分量变成1,每个负分量变成0,而幅度被丢弃: [LOADING...]
改变表示方式,也就改变了度量方式。余弦相似度需要我们已经丢弃的幅度,因此二值向量改用汉明距离比较,也就是两个位模式不一致的位置数量。例如,比较10110100和10010110。它们有两个位置不同,所以距离为2。在完整长度下,计算同样简单。一个4,096位向量存为64个64位字,比较两个这样的向量,就是把每对字进行异或,两个向量不一致的位置会留下1,然后统计1的个数。CPU对每个字都能用一条指令完成这些操作,所以一次完整比较大约需要上百条指令,而对原始浮点数做余弦相似度则需要数千次乘法。
请注意,度量方式从来都不是一个单独的决定。我们为了节省存储而选择一位精度,而一旦每个分量都只是一个符号位,汉明距离就成了唯一合理的比较方式。选择了精度,也就选择了度量方式。
二值量化与未量化向量相比,仍然会损失几个百分点的召回率。我们接受了这一代价,因为考虑到消费结果的是一个推理型智能体。为智能体提供结果的代码搜索,更需要的是正确的邻域,而不是一个排序完美的前10名。当智能体询问会话令牌在哪里刷新时,重要的是相关的那几个文件出现在前十几个结果中。最佳块排第二还是第五并无影响,因为智能体无论如何都会打开候选并阅读它们。在这个循环中,一个在面向人类的三结果界面里会显而易见的排序下降,基本上是看不见的。
二值量化的局限
我们做出的取舍还有一个更隐蔽的代价,我们花了一段时间才理解。二值量化不仅牺牲准确率;它还压缩了相似度分数的范围。对于全精度向量,不相关的一对可能得分接近零,而近似重复项得分接近一,分布范围足够宽。符号位的行为不同。两个完全不相关的向量,大约有一半的位仍然会纯属偶然地一致,而强相关的一对可能达到三分之二的一致。所以索引中的每个分数,无论相关与否,都落在这个狭窄的区间里。
排序在压缩后仍然成立,因为相关结果得分仍高于不相关结果,但阈值判断不行。设想一个无需询问就主动提供相关代码的功能,比如在你输入时建议已有实现的面板。它最难的要求是知道何时保持沉默。要做出这种判断,它需要在“相关”和“不相关”分数之间有一个可用的间隔。二值向量没有留下这样的间隔。任何放在这个狭窄区间内的截断值,要么对所有内容触发,要么对任何内容都不触发。因此,当索引需要绝对的相关性判断而不是相对排序时,我们会保留16位浮点数,并为此支付存储成本。
嵌入范围
虽然索引和搜索使用同一个模型,但这两项工作截然不同。索引受吞吐量限制,需要异步处理数百万个块。GPU每批大约处理32个块就会饱和。另一方面,搜索需要快速且响应及时。如果最多几秒内没有给出结果,用户就会放弃。因此,在部署这些模型时,我们会相应地进行优化:一个最大化每秒处理的块数,另一个最小化首结果时间。
我们选择了一个指令遵循模型,它在训练时刻意在检索的两侧之间制造不对称。重要的是,这两侧由非常不同类型的文本表示。查询是自然语言中的简短问题,而文档是一块代码。在索引时,文档按原样嵌入。查询会被包装上描述检索任务的指令,类似“给定这个搜索查询,找出回答它的代码”,这告诉模型这段文本扮演什么角色。我们在推理时保留这种安排,因为这是模型学习到的形态。
为了让两侧更容易对齐,我们将每个块与其文件路径一起嵌入。路径提供了仅靠块本身所缺少的元数据:它位于哪个模块,以及这个文件是什么。不过,在单体仓库中,路径本身会成为一个问题。IntelliJ IDEA 单体仓库的文件数量超过一百万。那里源文件路径的中位数是九层目录深、91个字符长;接近10,000个源文件的路径超过150个字符,其中最长的为218个字符。这还没有算上前置的检出根路径。
这些字符大多用于结构嵌套,并不提供关于文件的有用信息。像 src/org/jetbrains/kotlin/idea/k2 这样一连串路径段重复了包层级,编译器需要它,而搜索不需要。与此同时,那条最长路径末尾的文件只有24行。如果我们只是把路径文本原样放在块旁边一起嵌入,就会发现路径有时会比代码本身占用更多空间。为了弥补这一点,路径在进入模型之前会被限制长度,规则是两端都要保留。开头的路径段告诉你位于哪个模块,而最后两段——直接父目录和文件名——告诉你这个文件是什么。中间部分是可以去掉的,而且只去掉截断所需的数量。保留仍然放得下的最长前缀,把中间省略为 ...,如果连父目录加文件名都太长,就只保留文件名本身。
[LOADING...]
当用户将搜索范围限定到子目录时,也适用同样的原则。显而易见的实现是元数据过滤:照常运行搜索,然后丢弃目录之外的结果。我们做法不同。范围会被渲染进查询文本本身,采用相同的形态、相同的缩写函数,以及索引块所使用的相同分隔符。如果一个块以 community/plugins/kotlin 的缩写形式进入索引,那么限定到该目录的查询会以完全相同的形式携带同一个字符串,因此查询向量会落在与它应该匹配的块相同的区域。
保护源代码
我们还考虑了最后一项设计因素。对我们来说,关注客户的隐私和安全关切非常重要。公司的源代码往往是其知识产权的核心。将其暴露给第三方云模型,甚至暴露给另一家公司,会增加无意中泄露敏感数据的风险,甚至可能让其他模型用它来训练。
为了确保解决这些关切,我们很早就决定遵循以下几项做法:
- 避免在我们的系统中存储代码:一个块包含集群引用、条目类型、文件路径、起始和结束偏移量、对向量的引用,以及一个可选的元数据字段。不保存任何内容,也不保存源代码本身的副本。搜索返回的是坐标,而你所看到的代码片段,是在你的机器上根据你的检出内容,用这些坐标组装出来的。服务器只知道某个相关的东西位于给定路径的第4,102--4,890字节处,而不知道它是什么。
- 不将数据用于训练:JetBrains Context 构建的每个代码索引,都由一个开放权重嵌入模型在我们运营的 GPU 上进行嵌入。没有任何嵌入请求会离开我们的基础设施——不会发给 OpenAI,不会发给 Google,也不会发给任何其他供应商。因此,我们可以保证没有任何数据会被用于训练任何东西。
这些自我施加的设计限制在检索质量方面没有带来任何代价。我们在自己的代码检索基准上,将开放权重候选模型与主要提供商托管的嵌入 API 进行了评估,结果我们的模型名列前茅。开放权重嵌入模型现在已经足够好,真正有趣的工程问题已经转移到你喂给它们什么、如何服务它们,以及你选择保留什么。
一段作为间奏的总结
在这篇博文中,我们介绍了检索流水线的最初几个阶段:从原始源文件到可供搜索的紧凑向量的旅程。
至此,我们有了数百万个二值向量,以及一种可以产生更多向量的方法。我们尚未解决的问题包括:如何高效地存储它们,如何创建一个能在毫秒级回答查询的系统,如何持续评估结果以确保我们做出了正确选择,以及如何让智能体真正使用我们这套闪亮的 RAG 装置。
这些主题以及更多内容,将是本系列接下来几篇的主题,我们会在未来几周内发布。一如既往,欢迎在评论中提出任何问题,或分享你在设计 RAG 解决方案时得到的深刻教训。我们很想了解你发现的不同而富有创意、且行之有效的方法!与此同时,欢迎查看 JetBrains Context,它目前处于公开预览阶段,并且已经包含在你的 JetBrains 许可证中 😀
下次见!