😸

CocoIndexを使ったリアルタイムコードベースインデックス作成

に公開
AI
PostgreSQL
Python
Rust
RAG
tech

CocoIndexは、AIワークロード向けの超高性能データ変換フレームワークで、Rustで書かれたコアエンジンを活用し、データの鮮度を重視したインクリメンタル処理を実現します。このチュートリアルでは、CocoIndexのビルトイン機能を使って大規模コードベースをインデックス化する方法を説明します。

なぜCocoIndexを使うのか？

従来のETLツールでは、コードの意味的チャンク化やリアルタイム更新が難しく、AIエージェント（例: ClaudeやCodex）のためのセマンティック検索が非効率です。CocoIndexはTree-sitterネイティブサポートにより構文構造に基づく精密な分割を可能にし、再計算を最小限に抑えてデータの一貫性を保ちます。これにより、コードレビューやRAG（Retrieval-Augmented Generation）アプリケーションの精度が大幅に向上します。

ユースケース

CocoIndexのコードベースインデックスは、多様なAI駆動のアプリケーションに活用可能です：

  • セマンティックコードコンテキストの提供：AIコーディングエージェント（Claude, Codex, Gemini CLI）向けに、常に最新のコードスニペットを供給し、正確なコード生成を支援します。CocoIndexのTree-sitterによる意味的チャンク化で、文脈の喪失を防ぎ、単純なテキスト分割より検索精度が向上します。

  • コードエディタのMCP統合：Cursor, Windsurf, VSCodeなどでコンテキスト認識検索を実現し、開発効率を高めます。CocoIndexのインクリメンタル更新により、ファイル変更時のみ再処理するため、大規模プロジェクトでも低遅延です。

  • コンテキスト認識コード検索：自然言語によるコード検索やレビューエージェントで活用。自動リファクタリングやSREワークフロー（例: IaCのルート原因分析）にも適し、変更影響評価を高速化します。

セットアップ

まず、Postgresをインストールし、CocoIndexをセットアップします。

pip install -U cocoindex

データベース接続 URLを環境変数に設定します：

export COCOINDEX_DATABASE_URL="postgresql://user:password@localhost:5432/dbname"

CocoIndexを使う利点は、プラグアンドプレイのビルディングブロック（ソース、変換、ストレージ）でパイプラインを構築できる点です。これにより、カスタムロジックを最小限に抑え、Rustコアの高性能で大規模コードベースを効率的に処理できます。

コードベースソースの追加と処理

CocoIndexのフロー定義で、LocalFileソースを使ってローカルファイルシステムからコードを読み込みます。例としてCocoIndex自身のコードベースを対象に、.py, .rs, .toml, .md, .mdxファイルを包含し、隠しファイルやtargetディレクトリを除外します。

import os
import cocoindex

@cocoindex.flow_def(name="CodeEmbedding")
def code_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
    data_scope["files"] = flow_builder.add_source(
        cocoindex.sources.LocalFile(
            path=os.path.join('..', '..'),
            included_patterns=["*.py", "*.rs", "*.toml", "*.md", "*.mdx"],
            excluded_patterns=[".*", "target", "**/node_modules"]
        )
    )
    code_embeddings = data_scope.add_collector()

ファイル拡張子を抽出するカスタム関数を定義します：

@cocoindex.op.function()
def extract_extension(filename: str) -> str:
    return os.path.splitext(filename)[1]

次に、SplitRecursivelyでTree-sitterを使ってコードをチャンク化します。言語パラメータ（拡張子から抽出）で構文解析し、chunk_size=1000, chunk_overlap=300を設定します。これにより、行ベースではなくAST（抽象構文木）に基づく意味的チャンクを作成し、AIの文脈理解を強化します。

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

埋め込み生成とエクスポート

チャンクごとにSentenceTransformerEmbed（モデル: sentence-transformers/all-MiniLM-L6-v2）で埋め込みを生成します。Transform Flowとして定義することで、インデックス作成とクエリ時の埋め込み一貫性を確保します。

@cocoindex.transform_flow()
def code_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
    return text.transform(
        cocoindex.functions.SentenceTransformerEmbed(
            model="sentence-transformers/all-MiniLM-L6-v2"
        )
    )

コレクターにファイル名、位置、コード、埋め込みを収集後、Postgresにエクスポートします。Cosine Similarityでベクトルインデックスを作成し、類似度検索を最適化します。

    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"]
        )

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

CocoIndexの利点は、状態管理によるインクリメンタル更新で、ソース変更時のみ再計算するため、生産環境でのメンテナンスコストを削減します。

インデックスクエリと実行

クエリ関数でSQLを実行し、クエリテキストの埋め込みとコサイン類似度を比較します：

from psycopg_pool import ConnectionPool

def search(pool: ConnectionPool, query: str, top_k: int = 5):
    query_embedding = code_to_embedding([query])[0]
    
    with pool.connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT filename, location, code, 
                       1 - (embedding <=> %s::vector) as similarity
                FROM code_embeddings
                ORDER BY embedding <=> %s::vector
                LIMIT %s
                """,
                (query_embedding, query_embedding, top_k)
            )
            return cur.fetchall()

メイン関数でループクエリを実行し、結果を表示します（スコア、ファイル名、コードスニペット）。

if __name__ == "__main__":
    pool = ConnectionPool(os.environ["COCOINDEX_DATABASE_URL"])
    
    while True:
        query = input("\n検索クエリを入力: ")
        if not query:
            break
            
        results = search(pool, query)
        for filename, location, code, similarity in results:
            print(f"\nスコア: {similarity:.3f}")
            print(f"ファイル: {filename}:{location}")
            print(f"コード: {code[:200]}...")

インデックス更新は以下のコマンドで実行します：

cocoindex update main

CocoInsightを使うと、フローを視覚化・デバッグできます：

cocoindex server main.py -ci

これにより、開発者がパイプラインの動作をステップバイステップで検証でき、Tree-sitterのチャンク化がAI検索の精度をどう向上させるかを確認できます。

まとめ

CocoIndexを使うことで、以下の利点が得られます：

  1. Tree-sitterによる意味的チャンク化：ASTベースの分割で、AIエージェントがコードの文脈を正確に理解
  2. インクリメンタル更新：変更ファイルのみ再処理し、大規模コードベースでも高速動作
  3. Rustコアの高性能：データ変換が超高速で、リアルタイムアプリケーションに最適
  4. プラグアンドプレイのビルディングブロック：カスタムETLロジックを書かずに、簡単にパイプライン構築
  5. RAGアプリケーションに最適：セマンティック検索でAIエージェントのコード生成精度を向上

このチュートリアルで紹介した手法を使えば、自社のコードベースをAIエージェント向けに効率的にインデックス化でき、開発生産性を大幅に向上させることができます。

Discussion