代码仓库索引调研--Cursor Repo Index 与基于 graph 的索引新方案

最近刷到一篇 paper，是关于代码定位（Code Localization）的，具体来说，是做“自然语言” -> “代码文件” 的映射。我正好在做基于 LLM 的代码生成功能，上个月把仿真脚本语言的 LLM 训练做完了，下一步想基于 RAG 技术提升代码生成的质量，正琢磨如何更高效地做代码索引和召回，感觉这篇文章在技术上有可借鉴之处，记录一下其中的技术细节，以及 Cursor 做代码 index 的方式。

paper：LocAgent: Graph-Guided LLM Agents for Code Localization

论文核心内容

LocAgent 是一个解决代码定位问题的框架，通过基于图的表示方法来处理这一挑战。将代码库解析为有向异构图，创建轻量级表示，捕获代码结构（文件、类、函数）及其依赖关系（导入、调用、继承），使 LLM 智能体能够通过多跳推理有效搜索和定位相关实体。

主要贡献

  1. 图表示方法：将代码库转换为有向异构图的新方法
  2. 多跳推理：LLM 智能体通过图结构进行有效的多跳推理
  3. 成本效率：使用微调的 Qwen-2.5-Coder-Instruct-32B 实现了与 SOTA 专有模型相当的结果
  4. 实际效果：在文件级定位上达到 92.7% 的准确率，下游 GitHub issue 解决成功率提高 12%（Pass@10）

图表示方法详解

异构图构造

LocAgent 将代码库解析为有向异构图（Directed Heterogeneous Graph），包含四类核心节点：

节点类型

  1. Directory 节点：代表目录/文件夹，反映代码的层次组织结构
  2. File 节点：代表代码文件
  3. Class 节点：代表类定义
  4. Function 节点：代表函数定义

边类型及构造方式

  1. Contains 关系

    • 表示包含关系（目录包含子目录和文件、文件包含类、类包含方法等）
    • 建立层次结构
    • 示例：src/ 目录包含 auth/ 子目录，auth/ 包含 login.py 文件

  2. Import 关系

    • 通过解析 import 语句和 from ... import ... 语句构建
    • 建立文件间的依赖关系
    • 示例：from utils import helper 创建 current_file → utils.py 的 import 边

  3. Invoke 关系

    • 通过 AST 分析识别函数调用和方法调用
    • 建立函数间的调用关系
    • 包括直接调用和间接调用

  4. Inheritance 关系

    • 解析类继承关系
    • 建立父类和子类之间的连接
    • 示例：class Child(Parent): 创建 Child → Parent 的继承边

代码到图的转换流程

1. 静态分析阶段

# 伪代码示例
def parse_codebase(codebase_path):
    graph = DirectedGraph()
    
    # 遍历所有文件
    for file_path in get_all_files(codebase_path):
        # 解析 AST
        ast_tree = parse_ast(file_path)
        
        # 创建文件节点
        file_node = create_file_node(file_path)
        graph.add_node(file_node)
        
        # 提取类和函数
        classes = extract_classes(ast_tree)
        functions = extract_functions(ast_tree)
        
        # 创建节点和边
        for cls in classes:
            cls_node = create_class_node(cls)
            graph.add_node(cls_node)
            graph.add_edge(file_node, cls_node, "contains")
            
        for func in functions:
            func_node = create_function_node(func)
            graph.add_node(func_node)
            graph.add_edge(parent_node, func_node, "contains")

2. 依赖关系提取

  • Import 分析：使用 AST 解析器识别所有导入语句
  • 调用关系分析：通过 AST 遍历识别函数调用和方法调用
  • 继承关系分析：解析类定义中的基类信息

3. 图优化

  • 噪声过滤：移除不相关的临时变量和内部函数
  • 权重赋值：根据调用频率和重要性为边分配权重
  • 层次压缩：合并过于细粒度的节点以提高效率

Qwen-2.5-Coder-Instruct-32B 微调详解

微调策略

LocAgent 采用了针对代码定位任务的专门微调方法：

1. 任务特化训练

  • 图遍历指令：训练模型理解图结构和遍历策略
  • 多跳推理：学习在图中进行多步推理
  • 代码理解：增强对代码语义和结构的理解

实验效果

主要指标

  • 文件级准确率：92.7%（在 SWE-bench 数据集上）
  • 成本降低：相比专有模型降低约 86%
  • 下游任务提升：GitHub issue 解决成功率提高 12%（Pass@10）

技术优势

  1. 轻量级表示：图表示相比完整代码更加紧凑
  2. 多跳推理能力：能够跨文件和模块进行推理
  3. 可扩展性：适用于不同规模的代码库
  4. 成本效率：显著降低了推理成本

对比 Cursor 的 Code RAG 技术

以现在风头正劲的 Cursor 为例，Cursor 的 repo index 技术是基于常规的 RAG 流程，通过以下步骤实现：

  1. Code Chunking and Processing
  2. Embedding Generation
  3. Storage and Indexing

Cursor 的实现细节参考了这篇文章，官方没有发布完整细节，有部分是推测的内容： How Cursor Indexes Codebases Fast

其中，Code Chunking 是最关键的，整个代码库索引的质量很大程度依赖于 chunk 的质量，不能简单根据 token 的数量来切分 chunk，需要使用能够理解代码结构的智能分割器，例如使用高级分隔符（例如，类和函数定义）在适当的语义边界处进行分割的递归文本分割器；另外，可以使用抽象语法树（AST）来实现 DFS 的代码遍历，从而实现代码的智能分割，一个好用的库是 tree-sitter，LocAgent 也使用了这个库来做 Class/Function 等节点的解析。

记录一个 tree-sitter 来实现代码的智能分割的文章：An attempt to build cursor's @codebase feature - RAG on codebases

Chunk 阶段还有一个小细节，使用了 merkle tree，通过计算 hash 值来实现代码的增量更新。

如何在推理时使用 embedding

上面的步骤中获得了整个 repo 的 embedding，在推理时，有以下几个步骤：

  1. query embedding
  2. Vector Similarity Search：使用向量检索来找到与问题最相关的 chunk，由于 cursor 的 embedding 是保存在官方服务器上的，所以这里查询的结果只用相关文件路径
  3. Local File Access：从上一步中获得相关代码的文件路径和行号，然后从本地文件中读取代码
  4. Context Assembly：将上一步中获得的代码作为 context 发送给 LLM
  5. Informed Response：LLM 根据 context 生成回答

embedding 模型

有开源的 embedding 模型，例如 sentence-transformers/all-MiniLM-L6-v2，看 HF 社区反馈，貌似中文支持一般。Cursor 很可能使用了 OpenAI 的 embedding 模型或者微调过的模型。

记录一个专门为 code 设计的 embedding 模型：microsoft/unixcoder-base

posted @ 2025-06-06 17:17  zion03  阅读(386)  评论(0)    收藏  举报