Claude Code を補強する RAG を自作してみた

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

mkdir myrag && cd myrag

# pyproject.toml を作る（依存関係を宣言）
cat > pyproject.toml <<'EOF'
[project]
name = "myrag"
version = "0.1.0"
requires-python = ">=3.11,<3.14"
dependencies = [
    "tree-sitter>=0.25",
    "tree-sitter-go>=0.25",
    "chromadb>=1.5",
    "dashscope>=1.25",
    "pathspec>=1.1",
]
EOF

# 仮想環境の作成と依存インストールを一括で
uv sync

# API キーは環境変数で渡す
export DASHSCOPE_API_KEY="sk-..."

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

# chunker.py
from tree_sitter import Language, Parser
import tree_sitter_go as tsgo

# Go 用のパーサ（他言語を足すときは tree_sitter_python などを import してパーサを増やす）
GO_LANGUAGE = Language(tsgo.language())
GO_PARSER = Parser(GO_LANGUAGE)

def chunk_go_file(path: str) -> list[dict]:
    """Go ファイルを関数・メソッド単位のチャンクに切る"""
    with open(path, "rb") as f:
        source = f.read()
    tree = GO_PARSER.parse(source)

    chunks = []

    def visit(node):
        # AST を深さ優先で走査し、関数/メソッドノードに当たったらチャンク化する
        # Go では関数（func Foo）とメソッド（func (r *R) Foo）で AST ノード名が違う
        if node.type in ("function_declaration", "method_declaration"):
            # 関数名を取得（匿名関数が混ざっても落ちないよう "anonymous" にフォールバック）
            name_node = node.child_by_field_name("name")
            name = name_node.text.decode() if name_node else "anonymous"
            chunks.append({
                "path": path,
                "symbol": name,
                # tree-sitter の行番号は 0 始まりなので +1 して人間が読みやすい形に
                "start_line": node.start_point[0] + 1,
                "end_line": node.end_point[0] + 1,
                "language": "go",
                # ノードのバイト範囲からソース該当部分を切り出してチャンク本文に格納
                "code": source[node.start_byte:node.end_byte].decode("utf-8", errors="replace"),
            })
            # 関数リテラル（クロージャ）は外側と内容が重複するため再帰しない
            return
        # 関数/メソッドでなければ子ノードを辿って関数宣言を探し続ける
        for child in node.children:
            visit(child)

    visit(tree.root_node)
    return chunks

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

# chunker.py（続き）
from pathlib import Path
import pathspec

# 行ベースチャンクのサイズ（関数・メソッド単位より粗いが、全体を拾える）
LINES_PER_CHUNK = 40

# インデックスから除外するパターン（.gitignore 記法）
EXCLUDE_PATTERNS = [
    # 依存・生成物
    "vendor/", "node_modules/", "bower_components/",
    ".git/", "__pycache__/", ".venv/", "dist/", "build/",
    "*.min.js", "*.min.css", "*.map", "*.generated.*", "*.pyc",
    # バイナリ・メディア
    "*.png", "*.jpg", "*.jpeg", "*.gif", "*.webp", "*.ico", "*.svg",
    "*.pdf", "*.zip", "*.tar", "*.gz", "*.bz2",
    "*.mp3", "*.mp4", "*.mov", "*.avi",
    "*.so", "*.dylib", "*.dll", "*.exe",
    # シークレット系
    ".env", ".env.*", "*.pem", "*.key", "*.p12", "*.pfx",
    "credentials.json", "serviceAccountKey.json", "*.secret", ".secrets/",
]


def chunk_by_lines(path: str) -> list[dict]:
    """ファイルを LINES_PER_CHUNK 行ずつの固定長チャンクに切る"""
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        lines = f.readlines()
    p = Path(path)
    # 拡張子がない場合（Makefile / Dockerfile 等）はファイル名を小文字で使う
    language = p.suffix.lstrip(".") or p.name.lower()
    chunks = []
    for i in range(0, len(lines), LINES_PER_CHUNK):
        block = lines[i:i + LINES_PER_CHUNK]
        chunks.append({
            "path": path,
            "symbol": "",  # 行ベースでは関数名は分からない
            "start_line": i + 1,
            "end_line": i + len(block),
            "language": language,
            "code": "".join(block),
        })
    return chunks


def chunk_repo(root: str) -> list[dict]:
    """リポジトリ配下を再帰的に走査してチャンクのリストを返す"""
    root_path = Path(root)
    # EXCLUDE_PATTERNS に加えて、リポ直下の .gitignore も読んで除外に反映
    patterns = list(EXCLUDE_PATTERNS)
    gitignore = root_path / ".gitignore"
    if gitignore.exists():
        patterns.extend(gitignore.read_text(errors="replace").splitlines())
    spec = pathspec.PathSpec.from_lines("gitwildmatch", patterns)

    chunks = []
    for path in root_path.rglob("*"):
        if not path.is_file():
            continue
        rel = str(path.relative_to(root_path))
        if spec.match_file(rel):
            continue
        if path.suffix == ".go":
            # Go は関数・メソッド単位で切る（精度が高い）
            chunks.extend(chunk_go_file(str(path)))
        else:
            # それ以外は行ベースに fallback（全体を拾う）
            chunks.extend(chunk_by_lines(str(path)))
    return chunks

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

# embedder.py
import os
import dashscope
from dashscope import TextEmbedding

# 国際版（DashScope Intl.）を使う場合は明示。中国本土版ならこの行は不要
dashscope.base_http_api_url = "https://dashscope-intl.aliyuncs.com/api/v1"

# text-embedding-v4 は 1 リクエストあたり最大 10 件の制限がある
_BATCH_SIZE = 10


def embed(texts: list[str]) -> list[list[float]]:
    """テキストのリストを埋め込みベクトルのリストに変換する"""
    results = []
    # API 側の上限（1 リクエスト 10 件）に合わせて分割して送る
    for i in range(0, len(texts), _BATCH_SIZE):
        batch = texts[i:i + _BATCH_SIZE]
        response = TextEmbedding.call(
            model="text-embedding-v4",
            input=batch,
            api_key=os.environ["DASHSCOPE_API_KEY"],
        )
        # output が None の場合はレート制限や認証エラー等。メッセージを付けて止める
        if response.output is None:
            raise RuntimeError(f"DashScope error: {response.code} {response.message}")
        # レスポンスから embedding ベクトルだけ取り出して結果に積む
        results.extend(item["embedding"] for item in response.output["embeddings"])
    return results

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

# index.py
from pathlib import Path
import chromadb
from chunker import chunk_repo
from embedder import embed

DEFAULT_DB_PATH = str(Path(__file__).parent / "chroma_db")

def build_index(repo_path: str, db_path: str = DEFAULT_DB_PATH):
    """リポジトリをチャンク化し、埋め込みベクトルと一緒に ChromaDB に格納する"""
    # ローカルファイルに永続化するクライアント
    client = chromadb.PersistentClient(path=db_path)
    # リネーム・削除で旧 ID が残らないように毎回コレクションを作り直す
    if any(c.name == "code_chunks" for c in client.list_collections()):
        client.delete_collection(name="code_chunks")
    # 類似度比較しやすいよう cosine 距離を指定（距離は 0〜2 の範囲で返ってくる）
    collection = client.create_collection(
        name="code_chunks",
        metadata={"hnsw:space": "cosine"},
    )

    # ① リポを走査してチャンク化 → ② コード本文だけ取り出し → ③ 埋め込みに変換
    chunks = chunk_repo(repo_path)
    codes = [c["code"] for c in chunks]
    vectors = embed(codes)

    # ベクトル・メタデータ・元コードをまとめて 1 件ずつコレクションに格納
    # ID は path + 開始行で一意にする（同じファイルの別チャンクを区別）
    collection.add(
        ids=[f"{c['path']}:{c['start_line']}" for c in chunks],
        embeddings=vectors,
        metadatas=[{
            "path": c["path"],
            "symbol": c["symbol"],
            "start_line": c["start_line"],
            "end_line": c["end_line"],
            "language": c["language"],
        } for c in chunks],
        documents=codes,
    )
    print(f"Indexed {len(chunks)} chunks")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

# search.py
from pathlib import Path
import chromadb
from embedder import embed

DEFAULT_DB_PATH = str(Path(__file__).parent / "chroma_db")

def search(query: str, top_k: int = 5, db_path: str = DEFAULT_DB_PATH) -> list[dict]:
    """クエリ文字列に近いチャンクを top_k 件返す"""
    # インデックス構築時に作ったコレクションを読み込む
    client = chromadb.PersistentClient(path=db_path)
    collection = client.get_collection(name="code_chunks")

    # クエリも埋め込みに変換してから近傍検索にかける
    query_vector = embed([query])[0]
    result = collection.query(
        query_embeddings=[query_vector],
        n_results=top_k,
    )

    # ChromaDB の返却は複数クエリ対応の 2 次元配列。今回は単一クエリなので [0] で展開
    # メタデータと距離を取り出し、Claude Code が扱いやすい辞書形式に整形する
    candidates = []
    for i in range(len(result["ids"][0])):
        meta = result["metadatas"][0][i]
        candidates.append({
            "path": meta["path"],
            "symbol": meta["symbol"],
            "start_line": meta["start_line"],
            "end_line": meta["end_line"],
            # ChromaDB の cosine 距離は `1 - cos(θ)`（範囲は 0〜2）。
            # 「0 が完全一致、2 が完全反対」を「1 が完全一致、0 が完全反対」に
            # 線形変換したいので、2 で割って 1 から引き 0〜1 のスコアにする。
            # 厳密なコサイン類似度ではなく、扱いやすい順位付けスコアとしての定義。
            "score": 1.0 - result["distances"][0][i] / 2.0,
        })
    return candidates

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

# cli.py
import argparse
import json
import sys
from index import build_index
from search import search

def main():
    parser = argparse.ArgumentParser(prog="myrag")
    sub = parser.add_subparsers(dest="command", required=True)

    # インデックス構築サブコマンド: build <repo_path>
    p_build = sub.add_parser("build", help="インデックスを構築")
    p_build.add_argument("repo_path")

    # 検索サブコマンド: search <query> [--top-k N]
    p_search = sub.add_parser("search", help="類似チャンクを検索")
    p_search.add_argument("query")
    p_search.add_argument("--top-k", type=int, default=5)

    args = parser.parse_args()

    if args.command == "build":
        build_index(args.repo_path)
    elif args.command == "search":
        results = search(args.query, top_k=args.top_k)
        # Claude Code が Bash ツール経由で受け取れるよう JSON を標準出力に書く
        json.dump({"candidates": results}, sys.stdout, ensure_ascii=False, indent=2)

if __name__ == "__main__":
    main()

1
2
3
4
5

# インデックス構築（最初の 1 回）
uv run cli.py build ./my-repo

# 検索
uv run cli.py search "ユーザー認証の実装はどこ" --top-k 5

1

alias myrag='uv run --project /absolute/path/to/myrag /absolute/path/to/myrag/cli.py'

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

{
  "candidates": [
    {
      "path": "internal/auth/handler.go",
      "symbol": "LoginHandler",
      "start_line": 42,
      "end_line": 78,
      "score": 0.87
    }
  ]
}