セマンティックコード検索のためのRAGパイプラインの構築：開発者の日記とフィールドノート

JetBrains AI 多くのJetBrains製品内でAI搭載機能を使用してツールを強化する フォローする フォロー: RSS RSS さらに探す エージェント型AI AI セマンティックコード検索のためのRAGパイプラインの構築：開発者の日記とフィールドノート Adam Malek Ashot Kazaryan パート1：解析、チャンキング、ベクトル化 しばらく前、私たちは可能な限り最高のセマンティックコード検索プラットフォーム、つまりLLMエージェントにgrepが偶然見つけるもの以上の、正確で引用可能な証拠を実際のレポジトリから提供するRAGパイプラインを構築することに着手しました。最終的なソリューションはAir Contextでした。私たちはそれを機能させ、本番環境に投入し、その過程で多くの苦労を経験しました。この一連の記事では、初日に誰かに教えてほしかった部分を共有します。 コーディングエージェントは、間違いなく私たちの時代のソフトウェア開発における最大の技術的飛躍です。エージェントとフロンティアモデルは、表面上は解決不可能なコードの複雑さに直面して、表面的には信頼性の高いコードを生成する能力を証明しています。しかし、ますます多くの開発プロセスがエージェント駆動型になるにつれて、エージェントの効率性と生成されるコードの品質はますます重要になります。問題は、エージェントがタスクを完了できるかどうかというよりも、生産グレードの結果を生成するためにどれだけの時間、労力、および指示が必要かということです。特に大規模なコードベースの場合、エージェントは、作業中の機能に関連するコードの関連部分を検索し、それをコンテキストに引き込むためにかなりの時間を費やすことになります。 セマンティック検索が重要な理由 適切なコードスニペットを見つけようとするとき、エージェントはキーワード検索やgrepのような従来のコード検索ツールに頼ることになります。しかし、これらのツールは、エージェントが検索する正確なテキストを事前に知っている必要があるという点で制限があります。たとえば、セッショントークンがリフレッシュされる場所を探しているエージェントは、コードに「refresh」という単語が含まれていることを期待することはできません。抽象的なドメインを推論するために、エージェントは意味によるコード検索、つまりセマンティック検索の能力を必要とします。ここで、検索拡張生成（RAG）が登場します。ソースコードをそのセマンティクスを捉える方法でインデックス化し、エージェントがオンデマンドで関連部分をフリーテキスト検索を使用して取得できるようにすれば、エージェントの強みを活かせるインターフェースを作成できます。 プロトタイプから本番環境へ エージェント時代の多くの優れたアイデアと同様に、ネイティブなプロトタイプ実装は非常にシンプルです。評価の高い本番グレードのソリューションは、間違いなくそうではありません。この一連のブログ記事では、効果的なRAGシステムを構築するために必要なこと、および私たち自身のシステムであるAir Contextを作成するまでの道のりで犯した間違いを共有したいと思います。前処理からストレージ、エージェント統合まで、各段階を扱い、より技術的なコンテキストとアドバイスを提供します。このシリーズの最初のパートでは、パイプラインの初期段階、つまり生のソースファイルを適切にスコープされた単位に分割する解析とチャンキング、およびそれらの単位をセマンティック検索をサポートする表現に変換するベクトル化について説明します。 解析とチャンキングの洗練されたAST 解析とチャンキングは、優れたRAGソリューションにおける重要な前処理ステップですが、しばしば見過ごされがちです。LLMがソースコードを埋め込むか、その他の方法でインデックス化できるようにするには、まず生のコード行を供給する必要があります。これは些細なことのように聞こえるかもしれませんが、小規模なデモプロジェクトではおそらくそうでしょう。しかし、本番グレードのシステムには数千のファイルが含まれており、それらはさらに数百または数千行に及びます。もしあれば、エージェントは問題を悪化させています。なぜなら、エージェントは熱心なライターの傾向があり、コードベースをさらに膨張させるからです。各ファイルには、関連性の度合いが異なる多数のクラス、フィールド、メソッドが含まれている可能性があります。適切なチャンクサイズを見つける これらの巨大なコードファイルを全体として埋め込みモデルに収めることができたとしても、その高価な作業は最終的に自己破壊的になります。ファイル全体が単一の単位として埋め込まれるため、検索はファイル全体を返します。これは、エージェントのコード探索とナビゲーションの目標、つまり特定の関数、シンボル、またはコードスニペットを見つけることに主に焦点を当てていることとは逆効果です。一方、他の極端なアプローチを取り、個々のコード行をそれぞれ埋め込むと、別の種類の問題に直面することになります。これらの個々の行は、周囲のコンテキストなしでは意味的に無意味になる可能性があります。一般的な関数名やコメントは埋め込む価値がなく、誤った検索結果を生み出します。ある意味では、木を見て森を見ることができず、エージェントは多数の、しばしば無意味なマイクロ結果で過負荷になります。したがって、コードを適切にスコープされたグループにチャンクまたは分割するための適切な方法を見つけることが不可欠です。各グループには、必要なコンテキストが十分に含まれ、一般的な意味論的意味を表している必要があります。 固定サイズのチャンキングがうまくいかない理由 チャンキングは、エージェントに供給されるコンテンツを取得し、一連のチャンクに分割する技術の一般的な名前です。チャンキングの単純なアプローチは、大きなファイルを単純に固定行数のグループに分割することかもしれません。しかし、そのアプローチを取った場合、結果のグループ化は意味論的に間違っていることがわかります。たとえば、インポートステートメントと一部の関数コンテンツがグループ化され、検索中に間違いが生じる可能性があります。この問題を解決するために、各ソースファイルには非常に明確に定義された構造があるという事実を利用できます。たとえばJavaを例にとると、インポートは通常ファイルの先頭にあり、その後にクラス定義があり、オプションでドキュメントコメントがヘッダーの前に付きます。クラスにはフィールドとメソッドが含まれ、それら自体も独自のドキュメントコメントを持つ場合があります。クラス構造を定義する規則と規約を知ることで、よりスマートなチャンキングを実行し、周囲の情報との適切なバランスを達成できます。 解析と構造を意識したチャンキング 過去26年間、JetBrainsでは、さまざまな言語のさまざまな奇妙な点、不規則性、規約、ニュアンスに対応できるスマートなパーサーを開発してきました。他のツールと並んで、これらのパーサーは、Air Contextが開発されている内部JetBrains Code Engineプラットフォームを形成しています。この記事の執筆時点では、Air Contextは9つの主要言語（Kotlin、Java、Python、JavaScript、TypeScript、C#、PHP、Go、Rust）の解析と構造を意識したチャンキングをサポートしています。他のすべての言語については、実装は単純に線ベースの分割にフォールバックし、どの言語やドキュメントでもインデックス化して検索できるようにします。パーサーにより、ソースファイルを、コメント、空白、修飾子のリストなど、それが表すものに関する情報を持つ構文ノードのストリームに分解できます。次に、チャンキングアルゴリズムがそのストリームを消費し、特定のチャンクのスコープを決定するロジックを適用します。ノードのタイプとサイズ、およびその子孫に基づいて、アルゴリズムが決定を下します。ノードがサイズ制限を超えているが子がない場合、より基本的な分割戦略にフォールバックします。一部の言語固有の構造は、推奨サイズを超えていても単一のスライスとして保持されます。ドキュメント、アノテーション、可視性修飾子、キーワードなどのプレフィックスは、宣言とともに保持されます。接尾辞（通常は終了構文）は、それが閉じる構造に関連付けられたままになります。また、たとえば@NotNullや@Overrideのような一般的で意味論的に無意味なJavaアノテーションが削除されるなど、言語固有のクリーニングも行われます。アルゴリズムは、2025年にZhangらによって作成されたcASTといくつかの類似点があります。私たちの実装とcASTの両方とも、適合する最大の構文単位を保持し、大きすぎる単位のみを分割し、小さすぎる隣接する単位をグループ化して、tinを回避します。