【RAG】ChromaDBを使ってRAGを試したい

はじめに

前回はGUIからRAGを試してみましたが、今回はコードを実装して試してみました！

ChromaDBのチュートリアルはこちらです。

コード

sample.py

import chromadb
import os
from dotenv import load_dotenv
import google.generativeai as genai

load_dotenv()

genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
chroma_client = chromadb.Client()

collection = chroma_client.create_collection(name="japanese_documents")

documents = [
    "みかんは美味しい果物だゾ",
    "りんごは赤い",
    "スマイルカンパニーのシステムにログインするためには笑顔認証が必要です",
    "スマイルカンパニーのシステムはスマホ対応していません。",
]

ids = ["id1", "id2", "id3", "id4"]

collection.add(documents=documents, ids=ids)

def search_similar_documents(query_text):
    results = collection.query(
        query_texts=[query_text],
        n_results=2
    )
    return results

def get_gemini_response(question, context):
    try:
        model = genai.GenerativeModel("gemini-1.5-flash")
        response = model.generate_content(
            f"以下の情報をもとに質問に答えてください:\n{context}\n質問: {question}"
        )
        return response.text
    except Exception as e:
        print(f"Error: {e}")
        return "Sorry, I couldn't fetch an answer at the moment."

if __name__ == "__main__":
    query_with_info = "美味しい果物について教えてください"
    query_without_info = "バナナの栄養について教えてください"
    
    results_with_info = search_similar_documents(query_with_info)
    results_without_info = search_similar_documents(query_without_info)

    sorted_results_with_info = sorted(zip(results_with_info["documents"], results_with_info["distances"]), key=lambda x: x[1], reverse=True)
    sorted_results_without_info = sorted(zip(results_without_info["documents"], results_without_info["distances"]), key=lambda x: x[1], reverse=True)

    similar_docs_with_info = "\n".join([doc if isinstance(doc, str) else "\n".join(doc) for doc, _ in sorted_results_with_info]) if sorted_results_with_info else "関連する情報が見つかりませんでした。"
    similar_docs_without_info = "\n".join([doc if isinstance(doc, str) else "\n".join(doc) for doc, _ in sorted_results_without_info]) if sorted_results_without_info else "関連する情報が見つかりませんでした。"

    gemini_response_with_info = get_gemini_response(query_with_info, similar_docs_with_info)
    gemini_response_without_info = get_gemini_response(query_without_info, similar_docs_without_info)

    print("質問:", query_with_info)
    print("Geminiの回答:", gemini_response_with_info)

    print("\n質問:", query_without_info)
    print("Geminiの回答:", gemini_response_without_info)

GEMINI_API_KEYは.envファイルに記載し、sample.pyと同階層に置いてください。

起動

$ python3 -m venv .venv
$ source .venv/bin/activate
$ python3 sample.py

事前にデータを登録する

"みかんは美味しい果物だゾ",
"りんごは赤い",
"スマイルカンパニーのシステムにログインするためには笑顔認証が必要です",
"スマイルカンパニーのシステムはスマホ対応していません。",

質問①

"美味しい果物について教えてください"
"バナナの栄養について教えてください"

結果①

質問: 美味しい果物について教えてください
Geminiの回答: テキストからは、みかんが美味しい果物であることがわかります。


質問: バナナの栄養について教えてください
Geminiの回答: 提示された情報からは、バナナの栄養について何も読み取れません。  みかんと笑顔認証システムの情報は、バナナの栄養とは全く関係ありません。  バナナの栄養については、別の情報源を参照する必要があります。

事前に登録したデータに対して関連する内容が見つかる場合は、適切な回答をしていますね！

次は事前データに"似ている"内容を質問してみます。

質問②

"スマイルカンパニーのシステムにログインする方法について教えてください"
"スマイルカンパニーシステムの管理者ページにアクセスする方法を教えてください"

結果②

質問: スマイルカンパニーのシステムにログインする方法について教えてください
Geminiの回答: スマイルカンパニーのシステムにログインするには、笑顔認証が必要です。  ただし、システムはスマホには対応していませんので、パソコンなどの対応機器から笑顔認証を行う必要があります。


質問: スマイルカンパニーシステムの管理者ページにアクセスする方法を教えてください
Geminiの回答: スマイルカンパニーのシステムにアクセスするには、笑顔認証が必要です。  しかし、スマホではアクセスできません。  管理者ページへのアクセス方法については、笑顔認証に対応したパソコン等からアクセスする必要があります。具体的なアクセス方法は記載されていませんので、それ以上の情報は提供できません。

"管理者ページ"のアクセス方法は登録していないのですが、笑顔認証を行う必要があると回答されてしまっていますね。

おわりに

今後は登録するデータを増やしたときに回答精度がどうなるか検証してみたいです