RAGで過去コードを自然言語検索する
X-tech推進本部 小林過去に作ったプログラムを、後から検索したくなることはないでしょうか。私は開発をしていると「似た機能を作ったことがあるな」「何のプロジェクトだったっけ」と思うことがよくあります。
もちろん、grepで検索もできます。しかし、grepは基本的に「文字列一致」です。例えば、「LLMにテキスト分類させるプロジェクト」のように、やりたいことベースで探したい場合、コード中にそのままの文言が存在しないことも多くあります。
そんな時に便利なのがRAG（Retrieval-Augmented Generation: 検索拡張生成）です。この記事では、「自然言語で過去コードを検索できる個人用RAGツール」をPythonで実装します。
RAGとは
RAGとは、LLM（大規模言語モデル）に外部情報を参照させて回答を生成する手法です。これによりLLMが学習済みの情報だけでなく、特定のドキュメントの情報をもとに回答を生成できます。例えば最近では「社内ドキュメントをもとに質問に回答するAI」という文脈でRAGという言葉を耳にすることが多くなりました。
ただし、多くの場合ドキュメントはそのままLLMで扱えるわけではありません。検索しやすくするため、テキストを分割し、意味ベクトル（意味を表す数値の並び）へ変換してデータベースへ登録します。
そのため、RAGというと大規模な社内ナレッジ整備やタグ付けといった、大掛かりな仕組みを想像しがちです。しかし、実際には個人レベルのドキュメント検索にもRAGは利用可能です。特に、プログラマーは既にRAG向きのテキスト資産を多く持っています。例えば、プログラムのソースコード、Markdownのメモなどです。以下では、そのような情報を使ってプロジェクトを検索可能にするツールを実装してみます。
実装
ディレクトリ構成
今回は次のような構成にします。
code_rag/
sources.yaml
build_index.py
ask.py
chroma_db/ # このディレクトリはbuild_index.py実行時に自動作成されるsources.yaml
まず、検索対象にしたいプロジェクトを設定します。今回は例としてPythonのプロジェクトを何件か参照します。
sources:
- project: local_llm_project
path: "D:/experiment/python/local_llm"
- project: notation_consistency_project
path: "D:/experiment/python/notation_consistency_checker"
- （略）ここで指定したディレクトリ内のファイル（今回は.pyと.md）がRAGへの登録対象になります。ポイントは、「既存プロジェクトをそのまま参照すればよい」ということです。専用フォーマットへ変換したり、タグ付けしたりする必要はありません。
build_index.py
次に、ソースコードを意味ベクトルに変換し、検索できる形でデータベースへ登録します。このことを「インデックス作成」と呼びます。このスクリプトでは次の処理を行います。
- ファイル探索
- chunk分割
- embedding
- ChromaDBへ登録
# build_index.py
from pathlib import Path
import chromadb
import yaml
from sentence_transformers import SentenceTransformer
BASE_DIR = Path(__file__).resolve().parent
CONFIG_PATH = BASE_DIR / "sources.yaml"
DB_DIR = BASE_DIR / "chroma_db"
TARGET_EXTENSIONS = {".py", ".md"}
embedder = SentenceTransformer("intfloat/multilingual-e5-small")
client = chromadb.PersistentClient(path=str(DB_DIR))
collection = client.get_or_create_collection("code_knowledge")
def load_sources():
with CONFIG_PATH.open(encoding="utf-8") as f:
return yaml.safe_load(f)["sources"]
def iter_target_files(root: Path):
for path in root.glob("*"):
if path.is_file() and path.suffix.lower() in {".py", ".md"}:
yield path
def split_text(text: str, max_chars: int = 3000):
return [
text[i:i + max_chars]
for i in range(0, len(text), max_chars)
]
def main():
existing = collection.get()
if existing["ids"]:
collection.delete(ids=existing["ids"])
for source in load_sources():
project = source["project"]
root = Path(source["path"])
for path in iter_target_files(root):
text = path.read_text(encoding="utf-8")
for i, chunk in enumerate(split_text(text)):
embedding = embedder.encode(chunk).tolist()
collection.add(
ids=[f"{path}-{i}"],
documents=[chunk],
embeddings=[embedding],
metadatas=[{"project": project, "path": str(path)}],
)
if __name__ == "__main__":
main()ファイル探索
iter_target_files()関数ではRAGで扱う対象となるファイルを探します。今回は処理をシンプルにするため、sources.yamlで示されたディレクトリ直下の.pyと.mdのみを探しています。
実運用では、必要に応じて、ディレクトリ内を再帰的に探索したり、特定のフォルダ（.venv、.gitなど）を除外する設定を入れるとよいでしょう。
chunk分割
長いコードをそのまま扱うと、意味が混ざって検索精度が落ちることがあります。そのため、RAGでは通常「chunk」と呼ばれる小さな単位に分割したデータを扱います。今回はシンプルさを優先し、コード内のsplit_text()関数で示したように、数千文字ごとの単純な文字数分割にしました。
実際には、「関数単位」「class単位」「Markdown見出し単位」など、「意味のまとまり」で分割したほうが検索品質は上がることがあります。ただ、今回のような個人用検索ツールであれば、単純な文字数分割でも十分実用的でした。
embedding
embeddingとは、文章を「意味ベクトル」に変換することです。意味の近い文章は近いベクトルになるため、「文字列一致」ではなく「意味の近さ」で検索できるようになります。意味ベクトルへの変換にはsentence-transformersというライブラリを使います。
また、今回はembedding modelとしてintfloat/multilingual-e5-smallを利用しています。これは多言語対応の軽量embeddingモデルで、日本語にも対応しています。
from sentence_transformers import SentenceTransformer
# （中略）
embedder = SentenceTransformer("intfloat/multilingual-e5-small")ChromaDBへ登録
ChromaDBは、embeddingで作った意味ベクトルを保存するデータベースです。これにより質問時に「意味の近いコード」を検索できるようになります。
import chromadb
# （中略）
client = chromadb.PersistentClient(path=str(DB_DIR))
collection = client.get_or_create_collection("code_knowledge")chunk分割したコードを意味ベクトル化しデータベースに登録
RAGの本体はここです。ここまでで作ったsplit_text()（chunk分割）、embedder（意味ベクトル作成）、collection（ChromaDB）を使ってRAG用のデータベースを作ります。
for i, chunk in enumerate(split_text(text)):
embedding = embedder.encode(chunk).tolist()
collection.add(
# （中略）
embeddings=[embedding],
# （中略）
)ask.py
次に、自然言語で検索するためのスクリプトです。
# ask.py
import sys
from pathlib import Path
import chromadb
import requests
from sentence_transformers import SentenceTransformer
BASE_DIR = Path(__file__).resolve().parent
DB_DIR = BASE_DIR / "chroma_db"
API_BASE = "http://localhost:11434/v1"
MODEL_NAME = "qwen2.5-coder:7b"
embedder = SentenceTransformer("intfloat/multilingual-e5-small")
client = chromadb.PersistentClient(path=str(DB_DIR))
collection = client.get_collection("code_knowledge")
def ask_llm(prompt: str) -> str:
res = requests.post(
f"{API_BASE}/chat/completions",
json={
"model": MODEL_NAME,
"messages": [{"role": "user", "content": prompt}],
}
)
res.raise_for_status()
return res.json()["choices"][0]["message"]["content"]
def answer(query: str) -> str:
query_embedding = embedder.encode([query]).tolist()[0]
result = collection.query(
query_embeddings=[query_embedding],
n_results=5,
)
contexts = []
for doc, meta in zip(result["documents"][0], result["metadatas"][0]):
contexts.append(
f"project: {meta['project']}\npath: {meta['path']}\n\n{doc}"
)
context = "\n\n---\n\n".join(contexts)
prompt = f"""
以下の検索結果だけを根拠に質問へ答えてください。
# 質問
{query}
# 検索結果
{context}
# 出力形式
- 該当しそうなプロジェクト
- 根拠
- 参照ファイル
"""
return ask_llm(prompt)
if __name__ == "__main__":
query = " ".join(sys.argv[1:])
print(answer(query))検索部分
まず質問を意味ベクトルに変換し、近いコードを検索します。文字列一致ではなく、「意味の近さ」で検索できます。なお、「自然言語で関連コードを検索する」だけであれば、このcollection.query()の時点でほぼ実現できています。今回はさらに、「検索結果の理由」を説明させるために次のステップでLLMを利用します。
query_embedding = embedder.encode([query]).tolist()[0]
result = collection.query(
query_embeddings=[query_embedding],
n_results=5,
)LLMへ渡す
つづいて、検索結果として得られたコード断片を1つのテキストcontextに整形します。
contexts = []
for doc, meta in zip(result["documents"][0], result["metadatas"][0]):
contexts.append(
f"project: {meta['project']}\npath: {meta['path']}\n\n{doc}"
)
context = "\n\n---\n\n".join(contexts)そして質問と一緒に1つのプロンプトとしてLLMに投げます。
prompt = f"""
以下の検索結果だけを根拠に質問へ答えてください。
（中略）
"""
return ask_llm(prompt)ask_llm()関数は作成したプロンプトをLLMへ送信し、回答文を取得する関数です。今回は例としてOllamaを利用しqwen2.5-coderのローカルモデルを呼び出しましたが、OpenAI互換APIを利用しているため、他のLLM環境へも比較的簡単に切り替えられます。
Ollamaの使い方については「ローカルLLMでできること：軽量モデルを業務自動化に組み込む方法」という記事で紹介しています。
実行
まず必要なライブラリをインストールします。
pip install chromadb sentence-transformers pyyaml requestsつづいて、build_index.pyを実行しインデックスを作成します。初回実行時にはembedding modelをHugging Face Hubからダウンロードするため、少し時間がかかります。
python build_index.pyこれでコードを検索できる準備が整いました。例えば、次のように質問付きでask.pyを実行します。
python ask.py "LLMでテキスト分類を行ったプロジェクトはありますか？"すると下記のような回答が得られ、該当プロジェクトを見つけることができました。
- 該当しそうなプロジェクト: local_llm_project
- 根拠: 既存のテキスト分類APIを使用して文章をカテゴリーに分類する機能が実装されています。
- 参照ファイル:
- `D:\experiment\python\local_llm\ollama_api.py`
- `D:\experiment\python\local_llm\main.py`
これらのファイルは、特定のテキストをPOSTリクエストでLLMリソースに送信し、そのレスポンスを解析して分類結果を取得しています。まとめ
今回は、シンプルなRAG検索ツールを作ってみました。RAGというと大規模システムのようなイメージがあるかもしれませんが、個人レベルで持っているテキスト資産だけでも便利に使えます。
特に、今回試したように「タグ付けや専用フォーマットを作ったり、ファイルを特定の場所にアップロードしたりしなくても、とりあえず既存プロジェクトへの参照さえ設定すれば自然言語検索できる」というのは、個人用途ではかなり便利に活用できそうです。
手軽にできる自然言語でのコード検索方法として試してみてはいかがでしょうか。