大規模なコードベースに対する検索拡張生成（RAG）で私が使っているアプローチを共有し、同じような問題を考えている方々からフィードバックをいただきたいと思いました。
問題 素朴なコードベースRAGは通常、ファイルをテキストの断片に分割し、それを埋め込みにして類似検索します。これがコードではうまくいきません。というのも、断片レベルでの意味的な類似性は構造的な関係を捉えられないからです。たとえば、ファイルAの関数がファイルCで定義された型を呼び出していても、埋め込みの近さだけではその依存関係が表面化しません。
アプローチ: AST由来の型付きグラフ 断片化の代わりに、Tree-sitterを使って各ファイルをそのASTとしてすべて解析し、型付きのノード／エッジグラフを抽出します:
- ノード: 関数、クラス、インターフェイス、型、モジュール
- エッジ: インポート、エクスポート、呼び出し関係、継承、合成
これをSQLiteに永続的なグラフとして保存します。解析コストはプロジェクトごとに1回だけです。
検索: グラフノードに対するBM25 クエリ時には、埋め込みの類似度ではなく、ノードのメタデータ（名前、シグネチャ、ドキュメンテーション文字列、ファイルパス）に対してBM25スコアリングを実行します。上位スコアのノードをLLMに渡します。グラフ構造のおかげで、取得した関数はエッジをたどることで、その直接の依存関係を自動的に取り込みます。
経験的に、中〜大規模のコードベースでは、素朴なフルコンテキスト方式で必要になる約100Kトークンの代わりに、1クエリあたりおよそ5Kトークンに収まります。
複雑なクエリのための階層的フォールバック 複数ファイルにまたがる推論タスクの場合:
- 完全なグラフのMermaid図は、常にコンテキスト内にある永続的なアーキテクチャマップとして機能する
- BM25によるノード検索で、目的に合った参照（ターゲットのルックアップ）を行う
- コンテキスト容量の70%に達する時点で、主要モデルに渡す前に、重要度の低いノードを素早く圧縮する
ここでBM25を埋め込みより優先する理由 コード識別子（関数名、型名、モジュールパス）は語彙的に非常に特徴的です。BM25は、コードクエリにおける支配的な検索パターンである「完全一致〜ほぼ完全一致の識別子マッチング」において、埋め込みの類似度を上回ります。埋め込みは、自然言語のドキュメント文字列に対するクエリではより有効かもしれませんが、その比較についてはまだ厳密にベンチマークしていません。
まだ考えている未解決の問い:
- グラフに対するより良いエッジ重み付け戦略 — 現在はすべてのエッジが重みなし
- BM25単体よりも、クロスエンコーダによる再ランキングで精度が有意に改善するかどうか
- 静的に完全解決できないためコールグラフを完全に構築できない動的言語への対応
コードベース規模のRAGに対して、別のアプローチで取り組んだ方はいますか？ 特に、実際のコードベースに対して定量的なベンチマークを行い、ASTグラフ方式を埋め込みベースの断片取得と比較した例があるかどうかが気になります。
[link] [comments]
