AIにコードを理解させたいですか？このツールが役に立つそうです

このブログでは、CocoIndex を使用して RAG のコードベースをインデックスする方法を説明します。CocoIndex は、データのインデックス作成とクエリを支援するツールです。独自のデータ パイプラインを構築するためのフレームワークとして使用できるように設計されています。CocoIndex は、ネイティブの Tree-sitter サポートとともに、コード ベースのチャンク化の組み込みサポートを提供します。

木に座る人

Tree-sitterはパーサー ジェネレーター ツールと増分解析ライブラリで、Rust 🦀 - GitHubで利用できます。CocoIndex には、Tree-sitter との Rust 統合が組み込まれており、さまざまなプログラミング言語のコードを効率的に解析し、構文ツリーを抽出します。

コードベースのチャンク化は、コードベースをより小さく意味的に意味のあるチャンクに分割するプロセスです。CocoIndex は Tree-sitter の機能を活用して、任意の改行ではなく実際の構文構造に基づいてコードをインテリジェントにチャンク化します。これらの意味的に一貫したチャンクは、RAG システムのより効果的なインデックスを構築するために使用され、より正確なコード取得とより優れたコンテキスト保持が可能になります。

ファストパス🚀 - 完全なコードはここにあります。RAG パイプラインの Python コードはわずか 50 行程度ですが、ぜひチェックしてください🤗!

私たちの仕事が気に入ったら、 Github の CocoIndex にスターを付けて応援してください。温かいココナッツのハグで本当にありがとう 🥥🤗。

前提条件

Postgres がインストールされていない場合は、インストール ガイドを参照してください。CocoIndex はデータ インデックスの管理に Postgres を使用していますが、進行中のものも含め、他のデータベースをサポートするロードマップに含まれています。他のデータベースに興味がある場合は、 GitHub の問題またはDiscordを作成してお知らせください。

cocoIndexフローを定義する

コードベースから読み取り、RAG 用にインデックスを作成するための cocoIndex フローを定義しましょう。

上記のフロー図は、コードベースを処理する方法を示しています。

1. ローカルファイルシステムからコードファイルを読み取る
2. ファイル拡張子を抽出する
3. Tree-sitterを使用してコードを意味的なチャンクに分割する
4. 各チャンクの埋め込みを生成する
5. 検索用にベクターデータベースに保存する

このフローを段階的に実装してみましょう。

1. コードベースをソースとして追加します。

 @cocoindex.flow_def(name="CodeEmbedding") def code_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope): """ Define an example flow that embeds files into a vector database. """ data_scope["files"] = flow_builder.add_source( cocoindex.sources.LocalFile(path="../..", included_patterns=["*.py", "*.rs", "*.toml", "*.md", "*.mdx"], excluded_patterns=[".*", "target", "**/node_modules"])) code_embeddings = data_scope.add_collector()

この例では、ルート ディレクトリから cocoindex コードベースをインデックスします。インデックスを作成するコードベースへのパスは変更できます。拡張子が.py 、 .rs 、 .toml 、 .md 、 .mdxのすべてのファイルをインデックスし、 . で始まるディレクトリ、target (ルート内)、および node_modules (任意のディレクトリの下) はスキップします。

flow_builder.add_source次のサブフィールドを持つテーブルを作成します。こちらのドキュメントを参照してください。

• filename (キー、型: str ): ファイルのファイル名、例: dir1/file1.md
• content （型： binaryがFalseの場合はstr 、それ以外の場合はbytes ）：ファイルの内容

2. 各ファイルを処理して情報を収集します。

2.1ファイル名の拡張子を抽出する

まず、各ファイルを処理するときにファイル名の拡張子を抽出する関数を定義しましょう。カスタム関数のドキュメントは、こちらにあります。

 @cocoindex.op.function() def extract_extension(filename: str) -> str: """Extract the extension of a filename.""" return os.path.splitext(filename)[1]

次に、各ファイルを処理して情報を収集します。

 # ... with data_scope["files"].row() as file: file["extension"] = file["filename"].transform(extract_extension)

ここでは、ファイル名の拡張子を抽出し、 extensionフィールドに保存します。たとえば、ファイル名がspec.rsの場合、 extensionフィールドは.rsになります。

2.2 ファイルをチャンクに分割する

次に、ファイルをチャンクに分割します。ファイルをチャンクに分割するには、 SplitRecursively関数を使用します。この関数のドキュメントは、こちら にあります。

CocoIndex は Tree-sitter の組み込みサポートを提供しているため、言語をlanguageパラメータに渡すことができます。サポートされているすべての言語名と拡張機能を確認するには、こちらのドキュメントを参照してください。Python、Rust、JavaScript、TypeScript、Java、C++ など、すべての主要言語がサポートされています。指定されていない場合、または指定された言語がサポートされていない場合は、プレーンテキストとして扱われます。

 with data_scope["files"].row() as file: # ... file["chunks"] = file["content"].transform( cocoindex.functions.SplitRecursively(), language=file["extension"], chunk_size=1000, chunk_overlap=300)

2.3チャンクを埋め込む

チャンクを埋め込むには、 SentenceTransformerEmbed関数を使用します。関数のドキュメントは、こちら にあります。🤗 Hugging Faceでは 12,000 のモデルがサポートされています。お気に入りのモデルを選択するだけです。

 def code_to_embedding(text: cocoindex.DataSlice) -> cocoindex.DataSlice: """ Embed the text using a SentenceTransformer model. """ return text.transform( cocoindex.functions.SentenceTransformerEmbed( model="sentence-transformers/all-MiniLM-L6-v2"))

次に、各チャンクに対して、 code_to_embedding関数を使用して埋め込み、その埋め込みをcode_embeddingsコレクターに収集します。

transform(cocoindex.functions.SentenceTransformerEmbed(...)) を直接呼び出す代わりに、この code_to_embedding 関数を抽出します。

これは、インデックス フローの構築とクエリ ハンドラーの定義の間でこれを共有したいためです。あるいは、よりシンプルにするためです。この追加機能を回避して、その場で直接処理することもできます。少しコピーして貼り付けるのは大したことではありません。クイックスタートプロジェクトではこれを行いました。

 with data_scope["files"].row() as file: # ... with file["chunks"].row() as chunk: chunk["embedding"] = chunk["text"].call(code_to_embedding) code_embeddings.collect(filename=file["filename"], location=chunk["location"], code=chunk["text"], embedding=chunk["embedding"])

2.4埋め込みを収集する

最後に、埋め込みをテーブルにエクスポートしましょう。

 code_embeddings.export( "code_embeddings", cocoindex.storages.Postgres(), primary_key_fields=["filename", "location"], vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])

3. インデックスのクエリハンドラーを設定する

SimpleSemanticsQueryHandlerを使用してインデックスをクエリします。code_to_embedding 関数code_to_embedding query_transform_flowパラメータに渡す必要があることに注意してください。これは、クエリ ハンドラがフローで使用されるものと同じ埋め込みモデルを使用するためです。

 query_handler = cocoindex.query.SimpleSemanticsQueryHandler( name="SemanticsSearch", flow=code_embedding_flow, target_name="code_embeddings", query_transform_flow=code_to_embedding, default_similarity_metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)

クエリ ハンドラーを実行するためのメイン関数を定義します。

 @cocoindex.main_fn() def _run(): # Run queries in a loop to demonstrate the query capabilities. while True: try: query = input("Enter search query (or Enter to quit): ") if query == '': break results, _ = query_handler.search(query, 10) print("\nSearch results:") for result in results: print(f"[{result.score:.3f}] {result.data['filename']}") print(f" {result.data['code']}") print("---") print() except KeyboardInterrupt: break if __name__ == "__main__": load_dotenv(override=True) _run()

@cocoindex.main_fn() デコレータは、環境変数から読み込まれた設定を使用してライブラリを初期化します。詳細については、初期化のドキュメントを参照してください。

インデックスのセットアップと更新を実行する

🎉 これで準備完了です！

インデックスを設定および更新するには、次のコマンドを実行します。

 python main.py cocoindex setup python main.py cocoindex update

ターミナルにインデックス更新の状態が表示されます

クエリをテストする

この時点で、cocoindex サーバーを起動し、データに対して RAG ランタイムを開発できます。

インデックスをテストするには、次の 2 つのオプションがあります。

オプション 1: ターミナルでインデックス サーバーを実行する

python main.py

プロンプトが表示されたら、検索クエリを入力できます。例: spec。

 Enter search query (or Enter to quit): spec

検索結果はターミナルで確認できます

返される結果 - 各エントリには、スコア (コサイン類似度)、ファイル名、一致するコード スニペットが含まれます。cocoindex では、 cocoindex.VectorSimilarityMetric.COSINE_SIMILARITYを使用して、クエリとインデックス データ間の類似度を測定します。他のメトリックに切り替えて、すぐにテストすることもできます。

Consine Similarity の詳細については、 Wiki を参照してください。

オプション 2: CocoInsight を実行してデータ パイプラインとデータインデックスを理解する

CocoInsight は、データ パイプラインとデータ インデックスを理解するのに役立つツールです。データを保持せずにローカルの CocoIndex サーバーに接続します。

CocoInsight は現在早期アクセス中です (無料) 😊 見つけていただきありがとうございます! CocoInsight に関する 3 分間の簡単なビデオ チュートリアル: YouTube でご覧ください。

CocoIndexサーバーを実行する

python main.py cocoindex server -c https://cocoindex.io

サーバーが起動したら、ブラウザでCocoInsightを開きます。ローカルの CocoIndex サーバーに接続し、データ パイプラインとインデックスを調べることができます。

右側には、定義したデータフローが表示されます。

左側のデータプレビューでは、データインデックスを確認できます。

任意の行をクリックすると、コード チャンクの完全な内容とその埋め込みを含む、そのデータ エントリの詳細が表示されます。

コミュニティ

コミュニティからのご意見をお待ちしています。Github とDiscordで私たちを見つけることができます。

この投稿と私たちの仕事が気に入ったら、 Github で CocoIndex に⭐ スターを付けて応援してください。温かいココナッツのハグで感謝します 🥥🤗。