【ぴよログでRAG: 第8回】質問応答の核を実装：RAGの検索フェーズ

こんにちは、たねやつです。

これまでのステップで、私たちは育児ログをきれいに整形し（前処理）、意味のある塊に分け（チャンキング）、AIが理解できる数値ベクトルに変換し（エンベディング）、そしてそれらを高速に検索できる知識の基地（ベクトルデータベース）に格納しました。

準備は万端です。いよいよ今回は、このシステムに命を吹き込む最初のステップ、ユーザーからの質問を受け取り、知識ベースから関連情報を探し出す「検索（Retrieval）」フェーズを実装します。これは、RAG (Retrieval-Augmented Generation) の "R" の部分にあたり、AIアシスタントが的確な答えを返すための根幹となる機能です。

前の記事
この記事でできること
事前に必要なもの
RAGの検索フェーズとは？
実装：質問応答システム（検索部分）
実行結果の確認
最後に
次の記事

この記事でできること

RAGにおける検索（Retrieval）フェーズの役割と処理の流れがわかる。
ユーザーが入力した自然言語の質問を、検索用のベクトルに変換する方法を実装できる。
ChromaDBを使い、質問ベクトルと意味的に最も近いテキストチャンク（知識）をデータベースから取得する方法を習得できる。

事前に必要なもの

Pythonの実行環境
関連ライブラリ: chromadb, sentence-transformers, torch
前回の成果物: ChromaDBにデータが投入された chroma_db ディレクトリと、使用したエンベディングモデル。

RAGの検索フェーズとは？

検索フェーズの仕事は非常にシンプルです。

質問の受付: ユーザーから「昨日のミルクの合計量は？」といった自然言語の質問を受け取ります。
質問のベクトル化: 受け取った質問の文字列を、知識ベースのチャンクをベクトル化した時と全く同じエンベディングモデルを使って、質問ベクトルに変換します。
類似ベクトル検索: 生成された質問ベクトルをクエリとして、ChromaDBに「このベクトルに最も近いチャンクをいくつかください」と問い合わせます。
コンテキストの取得: ChromaDBから返された、関連性の高いテキストチャンク（複数）を取得します。これが、後のLLMが回答を生成するための「参考資料（コンテキスト）」となります。

重要なのは、知識のエンベディングと質問のエンベディングに必ず同じモデルを使うことです。異なるモデルを使うと、同じ「意味」でも全く異なるベクトルが生成されてしまい、正しく類似度を比較できません。

実装：質問応答システム（検索部分）

それでは、実際に検索フェーズを実装していきましょう。rag_retrieval.py という新しいファイルを作成します。

# rag_retrieval.py
import chromadb
from sentence_transformers import SentenceTransformer

class SimpleRAG:
    def __init__(self, model_name='intfloat/multilingual-e5-large', db_path="./chroma_db", collection_name="piyolog_rag_collection"):
        """
        RAGシステムの初期化
        """
        print("モデルとデータベースを初期化しています...")
        # 1. エンベディングモデルの読み込み
        self.model = SentenceTransformer(model_name)

        # 2. ChromaDBクライアントの初期化とコレクションの取得
        self.client = chromadb.PersistentClient(path=db_path)
        self.collection = self.client.get_collection(name=collection_name)
        print("初期化が完了しました。")

    def search(self, query, k=3):
        """
        質問を受け取り、関連性の高いチャンクを検索して返す
        :param query: ユーザーからの質問文字列
        :param k: 取得するチャンクの数
        :return: 関連性の高いチャンクのリスト
        """
        print(f"\n質問: '{query}'")

        # 1. 質問をベクトル化
        print("質問をベクトル化しています...")
        query_embedding = self.model.encode(query, normalize_embeddings=True)

        # 2. ChromaDBで類似ベクトルを検索
        print(f"ChromaDBで類似チャンクを{k}件検索しています...")
        results = self.collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=k
        )

        # 3. 検索結果からドキュメント（テキストチャンク）を抽出して返す
        retrieved_chunks = results['documents'][0]
        return retrieved_chunks

# --- 実行部分 ---
if __name__ == '__main__':
    # RAGシステムをインスタンス化
    rag_system = SimpleRAG()

    # 検索を実行
    test_query = "9月1日のミルクの量は？"
    retrieved_context = rag_system.search(test_query, k=2)

    # 結果の表示
    print("\n--- 検索結果 ---")
    for i, chunk in enumerate(retrieved_context):
        print(f"【関連チャンク {i+1}】")
        print(chunk)
        print("--------------------")

このスクリプトでは、RAGの機能を SimpleRAG というクラスにまとめています。

__init__: 初期化時に、エンベディングモデルとChromaDBへの接続を準備します。
search: ユーザーの質問（query）を引数として受け取ります。
- model.encode(query, ...): 質問をベクトル化します。
- collection.query(...): ChromaDBの最も重要なメソッドです。query_embeddings に質問ベクトルを渡すと、それに最も近いドキュメントを n_results で指定した数だけ返してくれます。
- results['documents'][0]: query はリストで結果を返すため、最初の要素 [0] に目的のチャンクが入っています。

実行結果の確認

上記のスクリプトを実行すると、どうなるでしょうか。（※皆さんの piyolog.txt の内容によって結果は変わります）

モデルとデータベースを初期化しています...
初期化が完了しました。

質問: '9月1日のミルクの量は？'
質問をベクトル化しています...
ChromaDBで類似チャンクを2件検索しています...

--- 検索結果 ---
【関連チャンク 1】
日付: 2023年09月01日
--------------------
時刻: 05:45, イベント: ミルク, 詳細: 240.0ml, メモ: 全部飲んだ！
時刻: 06:30, イベント: うんち, 詳細: メモ: ちょっとゆるめ
時刻: 21:00, イベント: 寝る
... (以下略)
--------------------
【関連チャンク 2】
日付: 2023年09月02日
--------------------
時刻: 00:55, イベント: 起きる
... (以下略)
--------------------

見事に、「9月1日」と「ミルク」というキーワードに最も関連が深いであろう「9月1日のチャンク」が1番目の結果として取得できました！2番目には、日付が近い9月2日のチャンクが来ていますね。

これで、AIがユーザーの質問に答えるための強力な「参考資料」を手に入れることができました。