代码检索RAG：使用 tree-sitter 按函数分块而不是512个Token

当代码库助手针对诸如“重试延迟是在哪里计算的？”这样的查询检索到错误的上下文时，罪魁祸首通常是简单粗暴的分块器。标准的长文本固定大小分块器像切文章一样每隔 512个 token 切割代码，硬生生把函数定义切成两半，迫使大语言模型对着残缺不全的代码盲目猜测。要解决这个问题，我们需要使用抽象语法树（AST）在代码的自然语义边界处进行分块。

高性能的人工智能工程需要稳定可靠的基础设施。无论您是在构建内部开发者工具、IDE 集成还是企业级代码检索系统，高效管理大模型与嵌入模型管道都至关重要。像 n1n.ai 这样的平台为开发者提供了统一的 API 聚合层，能够以极低的延迟在各大顶尖模型之间智能路由请求。

固定大小代码分块的痛点

固定窗口分块器完全没有语法概念。如果一个函数正好开始于 512 token 窗口的第470个 token，那么第一个块将包含装饰器和签名，而下一个块则包含没有关联任何名称或标识符的主体。当检索引擎将查询与这些碎片化的数据进行匹配时，效果往往差强人意。

卡耐基梅隆大学与 Augment Code 在 cAST 论文中的研究证实了这一问题：“现有的基于行的分块启发式方法经常破坏语义结构，分割函数或合并无关代码。”他们基于 AST 的方法在 RepoEval 检索上取得了 4.3个百分点的 Recall@5 提升，在 SWE-bench 生成上取得了 2.67个百分点的 Pass@1 提升。

使用 tree-sitter 实现 AST 分块器

为了尊重代码的固有结构，我们使用 tree-sitter 在函数和类定义处断开文件，并将文件路径与封闭符号作为元数据进行携带。

首先，安装必要的依赖项，包括 py-tree-sitter 和嵌入模型：

pip install tree-sitter tree-sitter-python sentence-transformers numpy

接下来，实现 chunkers.py 来处理基于 AST 的解析以及作为基准对比的固定窗口分块：

# chunkers.py
from dataclasses import dataclass
from typing import Callable

import tree_sitter_python as tspython
from tree_sitter import Language, Parser

PARSER = Parser(Language(tspython.language()))
DEFS = \{"function_definition