リアルタイムのコードベースインデックス作成
このブログでは、CocoIndex を使ってRAGのためのコードベースのインデックス作成方法をご紹介します。
CocoIndexは、ネイティブのTree-sitterサポートとリアルタイム更新機能を備えた、コードベースのチャンク化を標準でサポートしています。
内部的にインクリメンタル処理（変更があった部分のみ再処理）を採用しているため、インデックスはほぼリアルタイムで更新されます。
たとえばコードエディタのような、変更通知を直接プッシュするソースと連携すれば、完全なリアルタイム化も可能です。
この仕組みはコード検索アプリケーションや、低遅延が求められるAIコード生成システムに関連コンテキストを提供する用途など、シームレスな開発体験に最適です。
もしこのチュートリアルが役に立ったと感じたら、CocoIndexのGitHub にスター⭐を付けていただけると嬉しいです。
Tree-sitter
CocoIndexは、Tree-sitterの機能を活用して、単なる行区切りではなく実際の構文構造に基づいてコードをインテリジェントにチャンク化します。
このように構文的に一貫性のあるチャンクを使うことで、RAGシステム向けにより効果的なインデックスを構築でき、より精度の高いコード検索やコンテキストの保持が可能になります。
Tree-sitterはパーサジェネレーターツールであり、インクリメンタルパース用のライブラリです。
Rust 🦀 で利用できます。
このプロジェクトの全コードはこちらで公開しています。
インデックス作成部分はたった約50行のPythonで実装されています。ぜひチェックしてみてください :rocket:!
前提条件
Postgresがインストールされていない場合は、インストールガイドをご覧ください。
CocoIndexは、インクリメンタル処理のためにPostgresを使用します。
インデックス作成フローの定義
フローの設計
フローのダイアグラムは、コードベースをどのように処理するかを示しています:
- ローカルファイルシステムからコードファイルを読み込む
- Tree-sitterが解析するためのコードの言語を取得するために、ファイル拡張子を抽出する
- Tree-sitterを使用してコードをセマンティックチャンクに分割する
- 各チャンクの埋め込みを生成する
- ベクトルデータベースに保存し、検索に利用する
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, content) を持つテーブルを作成します。詳細はドキュメントをご覧ください。
2. 各ファイルを処理し、情報を収集する
2.1 ファイル名の拡張子を抽出する
Tree-sitterでコードを解析するために、言語（または拡張子）を渡す必要があります。
各ファイルを処理する際に、ファイル名の拡張子を抽出する関数を定義します。
カスタム関数のドキュメントはこちらをご覧ください。
@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フィールドに保存します。
2.2 ファイルをチャンクに分割する
Tree-sitterを使用してコードをチャンクに分割します。
SplitRecursively関数を使用して、ファイルをチャンクに分割します。
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では、12kモデルがサポートされています。
お気に入りのモデルを選択できます。
@cocoindex.transform_flow()
def code_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
"""
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コレクターに収集します。
@cocoindex.transform_flow() は、インデックス作成とクエリ間で変換を共有するために必要です。
ベクトルインデックスを構築し、それに対してクエリを実行します。
埋め込み計算は、インデックス作成とクエリ間で一貫性が必要です。詳細はドキュメントをご覧ください。
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_indexes=[cocoindex.VectorIndex("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
Cosine Similarityを使用して、クエリとインデックス化されたデータの類似度を測定します。
Cosine Similarityの詳細はWikiをご覧ください。
3. インデックスに対してクエリを実行する
ユーザーが提供したテキストに対して、インデックス作成フローで使用した埋め込み操作を再利用して、SQLクエリで一致させます。
def search(pool: ConnectionPool, query: str, top_k: int = 5):
# Get the table name, for the export target in the code_embedding_flow above.
table_name = cocoindex.utils.get_target_storage_default_name(code_embedding_flow, "code_embeddings")
# Evaluate the transform flow defined above with the input query, to get the embedding.
query_vector = code_to_embedding.eval(query)
# Run the query and get the results.
with pool.connection() as conn:
with conn.cursor() as cur:
cur.execute(f"""
SELECT filename, code, embedding <=> %s::vector AS distance
FROM {table_name} ORDER BY distance LIMIT %s
""", (query_vector, top_k))
return [
{"filename": row[0], "code": row[1], "score": 1.0 - row[2]}
for row in cur.fetchall()
]
ターミナルでクエリを実行するためのメイン関数を定義します。
@cocoindex.main_fn()
def _run():
# Initialize the database connection pool.
pool = ConnectionPool(os.getenv("COCOINDEX_DATABASE_URL"))
# 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
# Run the query function with the database connection pool and the query.
results = search(pool, query)
print("\nSearch results:")
for result in results:
print(f"[{result['score']:.3f}] {result['filename']}")
print(f" {result['code']}")
print("---")
print()
except KeyboardInterrupt:
break
@cocoindex.main_fn() デコレーターは、環境変数から設定を読み込んでライブラリを初期化します。
詳細はドキュメントをご覧ください。
インデックスのセットアップと更新
🎉 これですべての準備ができました！
インデックスのセットアップと更新を行うには、以下のコマンドを実行します。
python main.py cocoindex setup
python main.py cocoindex update
ターミナルでインデックスの更新状態を確認できます。
クエリのテスト
この時点で、CocoIndexサーバーを起動し、データに対してRAGランタイムを開発できます。
インデックスをテストするには、以下のコマンドを実行します。
python main.py
プロンプトが表示されたら、検索クエリを入力できます。たとえば: spec
ターミナルで検索結果を確認できます。
返された結果 - 各エントリには、スコア（コサイン類似度）、ファイル名、一致したコードスニペットが含まれています。
サポート
私たちは常に改善しています。また、機能や例がさらに追加されています。
この記事が気に入ったら、GitHubで⭐スターをお願いします。
ありがとうございました！
Discussion