在 IDE 中实现自然语言搜索代码：RAG 策略的设计与落地

5月前

最近，我们为 AutoDev 的 VSCode 版本中重新引入了先前设计的自然语言搜索代码功能。即，你可以使用自然语言提出问题，AutoDev 将搜索代码，并使用现有代码库作为上下文，来辅助你理解现有的代码库。你可以在 VSCode 的应用商店中搜索 AutoDev 或者 GitHub 上下载与安装最新版本。（由于精力所限，加现有的架构设计不够轻量，暂时没有支持 Intellij 平台的计划。）

TL;DR：我们在去年的自然语言搜索代码策略做了一系列优化，以更好地适应不同命名实体、不同代码库下的搜索需求。诸如于，添加领域名词的解析与支持、更好的 chunk 机制、检索降级、多种搜索策略、新的 ReRank 策略（实现中）等等。

为什么需要自然语言搜索代码？

相比于开开心心编写新代码，接手别人留下的屎山往往才是现实。在现实的软件开发中，你往往很容易交接到一个别人离职的代码或者系统。在这样的系统中，大量的代码是没有文档的，而且代码的命名实体也是千奇百怪的，诸如于 yyds、cxkdlq 等等。这就导致了哪怕是有文档，正常人也无法理解这些代码。

所以，在缺乏明显的分层架构与文档时，我们往往会从 API 入口或者数据库入口开始，逐步理解整个系统。因此，在旧时代里，我们构建了 ArchGuard 这样的架构治理平台，来可视化现有的系统，以帮助开发者更好地理解现有的代码库，并帮助架构师更好地治理软件架构。

在新时代里，我们也有了更多的可能性，也有了更好的技术基础设施。

什么是自然语言搜索代码？

简单来说，你可以使用自然语言提出任何关于代码库的问题，而后系统将会告诉你，诸如此类的信息：

代码实现逻辑
对应问题的代码信息
可选的代码的依赖链路、调用信息

还能基于上述的信息，告诉你的新代码应该如何实现，甚至于直接生成对应的 patch，让你直接应用到代码库中。诸如于 Bloop 就是一个非常不错的开源工具：

而事实上，当你使用技术相对比较领先的 AI 辅助工具时，其背后默默地进行围绕你问题的分析和处理。诸如于 GitHub Copilot Chat 提供了一种轻量级的实现，以支持你的代码搜索与理解：

使用高质量模型（如 GPT 4）将你的问题描述成自然语言与关键信息
使用 RAG（检索增强生成）/NLP（自然语言处理）技术进行代码搜索
使用长上下文模型（如 GPT 3.5）对检索后的代码进行分析与解释

经过一系列的机制，就可以帮助你更好地理解现有的代码库。

AutoDev 是如何实现自然语言搜索代码？

与我们去年设计时相比，现今已经有大量的开源工具与技术可以支持我们实现自然语言搜索代码。诸如，我们在设计 AutoDev for VSCode 时，底层的代码使用的是 Continue 的实现，主要是在 RAG （检索增强生成）策略上进行了一系列的优化，以更好地支持不规范开发场景下的代码问题。在实现上，你需要：

引入合适的向量存储数据库，以支持代码的向量化存储
构建其它上下文补充数据能力，诸如于全文搜索、远程 API 等
设计适用于不同场景的 RAG 策略

与我们先前在服务端实现的自然语言搜索代码不同，IDE 插件是客户端，大部分计算需要运行在本地，因此需要考虑到性能与体验，诸如于包体验、搜索速度等等。并且，由于 AutoDev 是一个开源项目，要尽可能考虑到用户的定制化需求，即他们可以自定义自己的向量化模型、自定义自己的搜索策略等等。因此，在设计上与服务端的工具差异比较大。

AutoDev RAG 技术栈选型

在去年设计的版本中，我们使用了 Chocolate Factory 作为 LLM SDK，以 Chapi 作为多语言代码解析引擎，ElasticSearch 作为代码搜索引擎。而在新的版本中，我们参考了 Continue 在 VSCode 中实现的技术栈：

LanceDB 作为向量搜索数据库
SQLite 作为文本搜索数据库
TreeSitter 作为代码解析引擎
ONNX Runtime 作为本地模型推理引擎
Sentence-Transformers 的 All-MiniLM-L6-v2 作为本地向量化模型

当然了，用户也可以选择自己的云端向量化模型，以提升检索质量。

RAG 策略设计与实现

在 RAG 的架构与实现上，几大 IDE 插件在设计上差距并不大。根据所解决问题的不同，差异主要体现在细节上：

Indexing 阶段：全量代码还是部分代码。根据想解决用户问题的不同，即只是针对当前代码上下文还是代码库做上下文，有的插件只会索引最近编辑的代码，有的插件会索引全量代码。其会导致实现上出现差异，诸如于全量代码时，我们需要引入本地增量索引 + 本地数据库，以支持全量代码的检索。而如果只是针对当前代码上下文，我们可以结合 NLP 技术，直接在内存中检索。
Chunk 机制：代码块的划分。在 RAG 的实现上，我们需要考虑到代码块的划分，以支持更好的检索。通常来说，会结合多种策略，诸如于基于语法块的划分、基于文本的划分等等。由于在实现上，主流的 VSCode 插件都是基于 TreeSitter 构建二次语法解析，所以大部分依旧会基于语法块的划分，即将代码块划分为函数、类等等。
Query 阶段：查询改写。我们需要考虑到用户的查询可能是不规范的，诸如于拼写错误、语法错误等等。除此，为了支持我们更好地进行代码检索，还需要结合搜索策略，生成不同的查询条件，诸如适用全文搜索的关键词，适用代码搜索的假设性代码等等。
ReRank 机制：检索结果排序。在最后，我们还需要考虑到检索结果的有效性以及排序，以支持更好的检索体验。在场景足够重要的情况下，我们可以结合 LLM 进行二次判断，以剔除无效的 chunk。同时，在数据不足的情况下，我们可以从全文搜索中、 Git 历史中获取更多的信息，以支持更好的排序。

考虑到不同场景下的不同需求，在当前的 AutoDev 版本中，我们引入了两种不同的搜索策略：

基于关键词生成的搜索策略：即用户输入的自然语言问题，会由 LLM 生成对应的问题和关键词，再进行对应的代码检索，其通常会以全文搜索为主。
基于假设性代码生成的搜索策略：即用户输入的自然语言问题，会由 LLM 生成对应的假设性代码和问题，再进行对应的代码检索，其通常会以语义化搜索为主。

随后，再根据返回的数据是否足够多，数据不足时会在 Git 历史中搜索，以提供更多的信息。其次，再根据 ReRank 策略，对检索结果进行排序和处理。最后，再由 LLM 生成对应的理解文本，以支持用户更好地理解代码。

需要注意的是：参数小的大语言模型理解能力并不好，所以在转换用户问题到关键词、代码等时，效果并不是特别理想。特别是，当用户的问题描述不够准确时，其生成的关键词、代码也会出现问题。

领域语言的设计与实现

在去年的旧版本中，由于服务端限制的原因，使得我一直加入的自定义领域名词的功能，并没有得到很好的支持。而在 VSCode 版本中，由于我们可以直接读取本地的文件，因此我们可以直接读取项目中的数据，以支持自定义领域名词的功能。因此，在实现上，我们会读取项目中的 team_terms.csv，以支持自定义领域名词的功能。在查询时，我们会将用户的问题与领域名词进行匹配，以支持更好的检索体验。

除此，我们还会将其作为 TF-IDF 的一部分，以支持更好的检索体验。

未来

尽管，每一次检索可能要花费大量的 token，但是对于遗留系统来说，这是一个非常好的补充。在未来的版本中，我们将会继续优化搜索策略，如果你有任何的建议或者想法，欢迎在 GitHub 上提出 Issue，或者加入我们：https://github.com/unit-mesh/auto-dev-vscode 。

微信扫码关注该文公众号作者

来源：phodal