用 Tree-Sitter 构建 CodeRAG

CodeRAG的概念正在被广泛用于许多开发工具，如Cursor、Windsurf和Copilot。然而，这些流程的实现作为用户项目较少可用，因此我决定通过学习研究论文、博客和其他资源，构建一个CodeRAG系统，通过LLM与您的代码库聊天。我在开发过程中发现特别有帮助的一篇研究论文是CAST：通过抽象语法树增强代码检索增强生成的结构化分块。

1、我在为代码库构建RAG系统时遇到的最大挑战

代码文件与PDF和文本文档等通用文件有根本性的不同。随机分块策略（固定大小）会破坏所有分块的含义。

函数、类和代码块等代码元素具有依赖性，除非我们了解在此函数内部调用的类，否则我们无法完全理解任何函数。

代码库文件中的导入也会创建文件间依赖关系，为了理解这些导入的用途，LLM需要这些导入的上下文。

我如何应对这些挑战：

使用Tree-sitter进行上下文丰富的分块。
使用Tree-sitter创建依赖图来解析所有类型的调用（函数、类）和导入。
有了依赖图，我检索最佳匹配的分块及其邻居，这有助于LLM提供更好的结果和完整的上下文。

下面是Neo4j中代码库的依赖图：

此应用程序现已完全部署并开源，欢迎探索。

这篇博文介绍了我的CodeRAG系统的第2版。我通过学习新方法不断改进它，添加其他很酷的功能（如代理行为）。该工作流由两个主要管道组成：摄取管道（用于处理和存储代码库）和检索管道（用于回答查询）。

3、摄取管道

3.1 克隆仓库

首先，我从客户端获取仓库URL，并通过Git使用唯一的会话ID将其克隆到本地。

git.Repo.clone_from(repo_url, local_path)

3.2 遍历仓库

这会逐个文件遍历仓库，通过扩展名字典提取每个文件的语言。如果是代码文件，我们将其发送到一个函数，从该代码文件中提取每个chunk_node。

for root, dirs, files in os.walk(repo_path):
    if ext in LANGUAGES:
        nodes = extract_nodes_from_file(file_path, language)

33 构建解析器

我使用tree-sitter-language-pack Python库为特定语言构建解析器，以解析代码并构建抽象语法树（AST）。

# get_parser方法从tree-sitter库导入
parser = get_parser(language)

3.4 提取导入

对于任何文件，一种提取导入的方法是在树遍历期间。如果存在类型为import_statement或import_declaration（取决于语言）的节点，我们可以将它们存储为文件的导入。

另一种方法是使用正则表达式在源代码上提取每种语言的所有文件导入。我使用了特定的结构来存储导入，字段包括imports_from、module等。

3.5 提取分块并创建文件内依赖关系

抽象语法树节点有几个字段，帮助我们理解任何代码的详细信息：

字段：

type（function_definition、class_declaration、identifier等）
is_named
text（源代码）
start_point（源代码第一个字符的行、列）
end_point（源代码最后一个字符的行、列）
parent（访问父节点的属性）
children（访问所有子节点的数组-like属性）

子节点只是当前节点表示的源代码的子组件。

4、树遍历（算法）

我使用DFS（深度优先搜索）算法遍历树。以下是DFS遍历期间发生的情况：

从抽象语法树的根节点开始：

1) 如果代码字符串大小小于MIN_CHUNK_SIZE，则为此树节点创建一个chunk_node，因为平衡的分块大小对于更好的检索很重要。以下是我们如何创建chunk_node：

创建唯一的node_id
解析此源代码的名称
将树节点的特定值存储在chunk节点中（如文件路径、节点类型、start_point、end_point等）
从该节点提取调用（我们有两种方法，下面讨论）
维护每个节点的兄弟关系
在此chunk节点中存储父节点ID

2) 否则，如果节点的代码字符串大小大于MIN_CHUNK_SIZE，我将该节点进一步分解并遍历其子节点。

以下是chunk_node结构：

{
    "id": node_id,
    "name": name,
    "code_str": code_str,
    "ast_type": ast_type,
    "file": file_path,
    "language": language,
    "start_line": start_line,
    "end_line": end_line,
    "start_byte": start_byte,
    "end_byte": end_byte,
    "size": size,
    "relationships": {
        "belongs_to": [],
        "parent": [],
        "sibling": [],
        "function_call": [],
        "class_call": [],
        "implements": [],
        "extends": [],
        "imports_from": []
    },
    "metadata": {
        "depth": depth,
        "calls": [],
        "type_references": [],
        "is_definition": false,
        "definition_type": null
    }
}

4.1 MIN_CHUNK_SIZE

抽象语法树可以将我们带到仅包含字面量、运算符和关键字的叶节点，这些单个令牌可能对理解文件的代码流没有价值。将这些存储为分块可能会导致数据库中填充垃圾分块。

因此，决定平衡的MIN_CHUNK_SIZE对于存储更好的上下文很重要。

决定因素：

您的LLM令牌大小（成正比）
您的向量嵌入维度（成正比）
您的依赖图的关系密度（成反比）

4.2 创建节点ID

我可以在这里创建一个随机节点ID，但有意义的实例具有更好的价值，因此我使用此结构作为任何节点ID：

node_id = f"{file_path}:{node.start_point[0]}:{node.type}"

4.3 提取名称

每种语言都有不同的节点架构，导致不同的节点类型和字段名称。为了解决这个问题，我们有一些方法：

a. 使用标准化映射： 为每种语言创建字段及其值的字典映射。

b. 使用正则表达式： 在每种语言的源代码上使用正则表达式来提取名称。

4.4 提取调用

这里我们也有几种方法从节点的源代码中提取调用：

a. 通过更深入的遍历： 当我们遍历到此节点的子节点或孙节点时，我们可能会找到表示函数调用的节点，如foo(a, b)，名称如call_expression、function_call等。通过提取标识符，我们可以识别此代码库中的调用。

b. 使用正则表达式： 因为任何语言中任何调用的源代码都遵循特定的结构和语法，所以使用此方法提取调用名称很容易，但可靠性较低。

4.5 其他考虑因素

我们可以根据需要在此处提取和存储其他信息，如类继承、接口实现等，因为改进没有限制。

5、解析调用

现在我们有了每个节点的调用名称，为了准备有用的关系，我们必须将这些调用与该函数或类调用的实际节点ID解析。这是构建依赖图所必需的。

要解析任何chunk节点的调用，我们遍历所有chunk节点并检查是否有任何chunk节点与调用同名。

示例： 如果任何节点有调用foo()，我搜索名称为foo的节点，并在关系类别中添加chunk_node_id。

"relationships": {
    "belongs_to": [],
    "parent": [],
    "sibling": [],
    "function_call": ["file_path.ext:start_point:function_declaration"],
    "class_call": []
}

现在我们有了每个文件的源代码作为分块节点及其已解析的文件内依赖关系。

5.1 解析文件间导入

由于我已经创建了每个文件的完整导入详细信息，包含模块名称和extract_from字段（第4点），现在为了将这些导入与实际的node_id解析，我们首先从每种语言的imports_from字段中识别文件。然后我们搜索该文件的chunk节点，看看它是否有该模块。如果找到，我们使用chunk_node_id解析该导入。

在此步骤中，我们为仓库的所有代码文件创建了完整的依赖图。

5.2 处理非代码文件

在仓库中，会有许多文件没有直接编写代码，但它们对于更好地理解代码库至关重要，如README文件、文本文件等。

我使用了用于文档的一般分块方法。以下是分块它们的一些方法：

固定大小分块： 最简单和最基本的方法，我们将文档分成恒定的预定义大小，忽略文本的自然结构，并有一些重叠以便为LLM提供更好的上下文。

语义分块： 在这里，文档首先被分解成更小的单元，如句子，然后为每个单元生成嵌入。我们将具有更高语义相似度的单元分组并将它们创建为一个分块。

递归分块： 对于README和文档等文件，这是最佳的分块策略。在这里，我们尝试使用一系列分隔符来分隔文本，从最大的结构单元开始，然后回退到较小的单元，如首先主要标题（#或##），然后子标题或段落。

6、将依赖图存储在图形数据库中

我使用Neo4j，因为它还提供在节点中存储向量嵌入的能力，并提供基于余弦相似度提取节点的功能。我们为每个chunk节点使用相同的会话ID存储在数据库中。

6.1 代码嵌入

我们准备每个chunk节点的源代码列表，并将其发送到嵌入模型。我使用gemini-embedding-001模型进行代码嵌入。

根据以下因素决定向量维度：

您的MIN_CHUNK_SIZE（成正比）
您的数据库约束

6.2 在Neo4j中存储Chunk节点

向每个对应的chunk节点添加嵌入，并删除任何不重要的字段。然后为您的数据库创建一个客户端，并使用chunk_node_id将节点存储在图形数据库中，这将帮助我们创建关系。

6.3 存储关系

通过遍历所有chunk节点，为每个关系创建一个包含source_id、target_id和关系类型的列表，并将这些关系存储在图形数据库中。

现在我们的摄取管道已完成，我们有了获取整个依赖图的会话ID。

7、检索管道

7.1 使用会话ID获取用户查询

当客户端为此仓库上的用户查询创建请求时，我们同时获取查询和会话ID，并将它们发送以进行进一步处理。

7.2 增强用户查询

也许客户端很懒，或者作为人类，我们总是有一些假设，所以用户可能会问"定义loginController"。为了解决这个问题，我们向LLM发送一条消息，其中包含user_query和自定义指令以增强此用户查询。以下是增强查询的示例：

原始查询： "定义loginController"

增强查询： "提供代码库中loginController的清晰解释或定义，包括其目的、相关函数以及它如何处理身份验证或登录逻辑"

这不是一个昂贵的解决方案，因为每个查询的令牌很少，但它会引入一些额外的延迟。最好的一面是它将大大提高通过向量嵌入检索更好上下文的机会。这是一个可选步骤。

7.3 创建增强用户查询的向量嵌入

使用与摄取期间相同的嵌入模型和相同维度，创建增强用户查询的向量嵌入。

7.4 数据库节点检索

由于我们需要从数据库中检索分块及其依赖关系，我们分两步处理：

检索前K个节点： 使用向量嵌入，提取与查询嵌入具有最高余弦相似度的前K个chunk节点。
检索所有相关chunk_nodes： 获取通过函数调用、类调用、导入等关系与这些前K个chunk_nodes相关的所有chunk_nodes。

7.5 上下文准备

现在我们有了相关的源代码及其依赖关系的源代码。从chunk节点中提取对上下文有帮助的字段，并按相似度降序连接它们。

7.6 LLM响应

现在准备一个包含用户查询、代码上下文和良好指令的提示。LLM响应用户查询的漂亮答案。将其发送回客户端。

这完成了我们的检索管道。

8、结束语

进一步改进：

添加代理行为： 我正在实现代理行为，LLM将决定哪些上下文相关。如果它对检索到的上下文不满意，LLM可以重新获取并迭代搜索更好的上下文，使用TAO（任务、操作、观察）原则提供最佳可能的结果。
改进的调用和名称提取： 目前，我通过正则表达式提取调用和名称，这有局限性。我计划通过在AST中更深入的节点遍历来提取详细信息，使其更可靠且与语言无关。
实现合并策略： 在我的学习过程中，我读到一篇研究论文（CAST：通过抽象语法树增强代码检索增强生成的结构化分块），其中他们合并了低于某个阈值的较小分块。我可以在这里实现类似的解决方案来优化分块大小并提高上下文质量。
多语言支持扩展： 此CodeRAG目前支持JavaScript、TypeScript、JSX、TSX和Python，每种语言都有专用的解析器。由于Tree-sitter支持多种语言，我可以通过添加相应的解析器配置来扩展支持任何语言。
路由特定查询： 有些用户查询LLM无法使用此检索方法回答，如"给我这个仓库的摘要"或"这个项目中使用了哪些技术？"对于此类查询，我计划为代理创建专门的工具，可以处理仓库级分析、聚合统计和高级概述，而不是特定于代码的上下文检索。

原文链接: How I Built CodeRAG with Dependency Graph Using Tree-Sitter

汇智网翻译整理，转载请标明出处