tree-sitterを使ってコードをインデックスする際のチャンクを自動で最適化する
前回に引き続き、最近書いているOSSの話です。
前回↓
コードをインデックスする際の問題点
普通にコンテンツを読み込んでチャンクに分割する時は、指定されたチャンクサイズとチャンクオーバーラップの値を見て、単純な文字数で分割されます。この際、オーバーラップに指定した数値の分だけ隣り合うチャンクは文字が重複するように分割されます。
const chunks = chunkText(text, {
size: 50, // characters per chunk
overlap: 10, // overlapping characters
});
Original: "TypeScriptは静的型付けを提供するJavaScriptのスーパーセットです。
型システムにより、コンパイル時にエラーを検出でき、開発効率が向上します。
また、型推論機能により、明示的な型注釈なしでも型安全性を確保できます。"
↓
Chunks (size=50, overlap=10)
[1] "TypeScriptは静的型付けを提供するJavaScriptのスーパーセットです。型シ"
[2] "ーパーセットです。型システムにより、コンパイル時にエラーを検出でき、開発"
[3] "ラーを検出でき、開発効率が向上します。また、型推論機能により、明示的な型"
[4] "機能により、明示的な型注釈なしでも型安全性を確保できます。"
こうすることで、チャンク分割されてもembedding時に意味的に近いチャンクが生成され、検索時の関連性が高まります。この際、元の文章の20%程度のオーバーラップを指定すると良いとされています。
さて、コードをチャンクに分割する時はこれで良いのでしょうか？以下のようなコードがあるとします。
export class UserService {
private readonly cache = new Map<string, User>();
async findUser(id: string): Promise<User> {
if (this.cache.has(id)) {
return this.cache.get(id)!;
}
const user = await this.repository.findById(id);
if (!user) {
throw new NotFoundException('User not found');
}
this.cache.set(id, user);
return user;
}
async updateUser(id: string, data: UpdateUserDto): Promise<User> {
const user = await this.findUser(id);
Object.assign(user, data);
await this.repository.save(user);
this.cache.delete(id);
return user;
}
}
これを単純な文字数で分割すると、こうなります。
// chunkSize=200 orverlap=40
// Chunk 1: 文法的に不完全
"export class UserService {
private readonly cache = new Map<string, User>();
async findUser(id: string): Promise<User> {
if (this.cache.has(id)) {
return this.cache.get(i"
// Chunk 2: 開始が途中から、終了も途中
"che.get(id)!;
}
const user = await this.repository.findById(id);
if (!user) {
throw new NotFoundException('User not found');
}
this.cache.set(id, user"
// Chunk 3: 関数の境界を無視
");
return user;
}
async updateUser(id: string, data: UpdateUserDto): Promise<User> {
const user = await this.findUser(id);
Object.assi"
// Chunk 4: 意味不明な断片
"gn(user, data);
await this.repository.save(user);
this.cache.delete(id);
return user;
}
}"
このように、コードを単純分割してしまうとチャンクごとに意味がめちゃくちゃになってしまいます。
例えば キャッシュからユーザーを取得する処理 とクエリしても、embeddingが壊れているためコードを見つけられない可能性が高くなります。
そのため、コードをembeddingする際はそのコードの意味的な境界を保持することが良い、とされています。コードの意味的な境界とは、例えば {} ブロックや関数、クラス等のモジュール、変数名などです。
これらを破壊しない形でembeddingできれば、検索時にマッチさせられる確率が高まります。
tree-sitter CSTを用いた意味的境界レベルのチャンク生成
そこで、gistdexではASTレベルでのチャンク生成最適化を実装することにしました。ただし、gistdexで読み込めるコンテンツにはさまざまな言語があり、特定の言語のみをサポートするパーサーでは不十分なため、複数の言語をサポートできるかつ、統一的なアプローチで実装できる tree-sitter を用いることにしました。tree-sitterは正確にはASTではなくCSTという、具象構文木を扱うライブラリですが、具象構文木の説明は割愛します。
gistdexでは、読み込むファイルの拡張子を見て対応するtree-sitterのパーサーを選択し、CSTレベルでチャンク生成を行います。tree-sitterに依存しているのでtree-sitterがサポートしていない言語は非対応ですが、それでも以下の言語で対応できています。
- JavaScript/TypeScript (.js, .jsx, .ts, .tsx, .mjs, .mts, .cjs)
- Python (.py)
- Go (.go)
- Rust (.rs)
- Ruby (.rb)
- C/C++ (.c, .cpp, .h)
- Java (.java)
- HTML (.html)
- CSS/SCSS/Sass (.css, .scss, .sass)
- Bash/Shell (.sh, .bash)
仕組みとしては、language-node-types という各言語の意味的境界を持つNodeTypeの一覧を持ち、traverseする際に node.type と比較してマッチしたら配列にプッシュし、その配列を元にチャンクを生成するというものです。マークダウンファイルの場合は自前のパーサーでセクションレベルの分割をしています。
また、拡張子ごとのチャンクサイズ上限とオーバーラップも自動で決まるようにしています。意味的な境界が上限を超える場合は、その上限で分割されます。
これで、意味的な境界を保持しながら最適なチャンクサイズでembeddingできます。クエリすると、モジュールレベルの結果が返ってくることがわかります。
❯ pnpm dlx @ushironoko/gistdex query -k 3 "mcp index tool"
A CLI tool for indexing and searching content using vector databases (gistdex v1.1.4)
Searching for: "query mcp index tool"
Found 3 results
1. /Users/ushironoko/dev/gistdex/src/cli/utils/special-flags.ts
Score: 1.200
Type: file
| const mcpIndex = args.findIndex(isMcpFlag);
2. /Users/ushironoko/dev/gistdex/src/mcp/server.ts
Score: 1.200
Type: file
| import { handleIndexTool } from "./tools/index-tool.js";
3. /Users/ushironoko/dev/gistdex/src/mcp/server.ts
Score: 1.200
Type: file
| import { handleQueryTool } from "./tools/query-tool.js";
Search Statistics:
Average Score: 1.200
Score Range: 1.200 - 1.200
この機能は --preserve-boundaries 、もしくはショートハンドの -p を指定してインデックスすると使えます。また、デフォルトの chunkSize や chunkOrverlap を指定しない場合、コンテンツに応じた最適なこれらの値をセットする機能もあります。同時に使うこともできます。
微妙なところ
現在の実装だと、関数をコールしているところやimport宣言は引っかかりやすいですが、モジュールの宣言自体を検索することがあまりできていないかもしれません。hybrid search（セマンティックサーチとキーワードサーチを組み合わせるオプション）でも引っかからないので、実装漏れかもしれません。かなり実用性に影響が出るので、そのうち修正したいところです。
多分直りました
Discussion