7-12. Gemini API / OpenAI API ＋ Chroma DB によるローカルRAGチャットボット

社内テキストドキュメントをベクトルデータベース（Chroma DB）に保存し、問い合わせに応じて文脈を検索・付与して回答生成するRAGシステムです。

Python

import chromadb
import google.generativeai as genai

# Chroma DBの初期化（第5章）
chroma_client = chromadb.Client()
collection = chroma_client.get_or_create_collection(name="company_knowledge")

# ドキュメントの登録
collection.add(
    documents=[
        "リモートワークの交通費精算は月額最高20,000円まで支給されます。",
        "有給休暇の申請は社内申請システム（Workflow）から希望日の3日前までに行ってください。"
    ],
    ids=["doc_1", "doc_2"]
)

def answer_query(user_query: str) -> str:
    # 類似テキストの抽出
    results = collection.query(query_texts=[user_query], n_results=1)
    context = results['documents'][0][0] if results['documents'] else ""
    
    # プロンプトの組み立て
    prompt = f"以下の参考情報をもとに質問に回答してください。\n\n【参考情報】\n{context}\n\n【質問】\n{user_query}"
    
    # LLMの呼び出し（例: Gemini API）
    # genai.configure(api_key=os.environ["GEMINI_API_KEY"])
    # model = genai.GenerativeModel('gemini-1.5-flash')
    # response = model.generate_content(prompt)
    # return response.text
    return f"（生成AI回答プレビュー）\nコンテキスト: '{context}' を使用して回答を生成しました。"

print(answer_query("有給はいつまでに申請すればいいですか？"))