【RAG】スプレッドシート経由でChromaDBにQAデータを登録したい

こちらの記事の続きです！

今回はデータベースに登録する形式を単一の文章ではなく、QA形式にしてみたいと思います。

単一形式とQA形式での回答の違い

• あなたはだれ

• 一人アドベントカレンダーはどれくらい記事かかないといけないの？

アドベントカレンダーについての回答が、単一形式だとGeminiの学習内容をもとにしているのに対し、
QA形式だとデータベースに登録したanswerをもとにしていることが確認できました。

コード

sample.py

import os
import time
import chromadb
import google.generativeai as genai
import schedule
from dotenv import load_dotenv
from slack_bolt import App
from slack_bolt.adapter.socket_mode import SocketModeHandler
from google.oauth2.credentials import Credentials
from google.auth.transport.requests import Request
from googleapiclient.discovery import build
from threading import Thread

load_dotenv()

genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

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

SLACK_BOT_TOKEN = os.getenv('SLACK_BOT_TOKEN')
SLACK_APP_TOKEN = os.getenv('SLACK_APP_TOKEN')

SHEET_ID = os.getenv("SAMPLE_SPREADSHEET_ID")
SAMPLE_RANGE_NAME = "シート2!A2:B5"

app = App(token=SLACK_BOT_TOKEN)

def get_sheets_data():
    creds = None
    if os.path.exists("token.json"):
        creds = Credentials.from_authorized_user_file("token.json", ["https://www.googleapis.com/auth/spreadsheets.readonly"])
    
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            flow = InstalledAppFlow.from_client_secrets_file("credentials.json", ["https://www.googleapis.com/auth/spreadsheets.readonly"])
            creds = flow.run_local_server(port=0)
        with open("token.json", "w") as token:
            token.write(creds.to_json())

    try:
        service = build("sheets", "v4", credentials=creds)
        sheet = service.spreadsheets()
        result = sheet.values().get(spreadsheetId=SHEET_ID, range=SAMPLE_RANGE_NAME).execute()
        values = result.get("values", [])

        if not values:
            print("No data found.")
            return []

        qa_pairs = [{"question": row[0], "answer": row[1]} for row in values if len(row) > 1]
        return qa_pairs

    except Exception as err:
        print(f"Error retrieving data from Google Sheets: {err}")
        return []

def add_documents_to_chroma(qa_pairs, ids):
    try:
        documents = [f"質問: {qa['question']}\n回答: {qa['answer']}" for qa in qa_pairs]
        collection.add(documents=documents, ids=ids)
        print("データがChromaDBに追加されました。")
    except Exception as e:
        print(f"エラーが発生しました: {e}")

def update_chroma_from_sheets():
    documents = get_sheets_data()
    if documents:
        new_ids = [f"id{index+1}" for index in range(len(documents))]
        add_documents_to_chroma(documents, new_ids)

def search_similar_documents(query_text):
    try:
        results = collection.query(query_texts=[query_text], n_results=2)
        return results
    except Exception as e:
        print(f"Error in searching documents: {e}")
        return None

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 in generating Gemini response: {e}")
        return "Sorry, I couldn't fetch an answer at the moment."

@app.message("")
def handle_message(message, say):
    user_message = message.get('text', '')

    search_results = search_similar_documents(user_message)
    
    spreadsheet_link = f"https://docs.google.com/spreadsheets/d/{SHEET_ID}/edit"

    if search_results and 'documents' in search_results:
        # search_results['documents'] がリストのリストである場合があるためフラット化する
        flat_results = [item for sublist in search_results['documents'] for item in (sublist if isinstance(sublist, list) else [sublist])]
        
        # 結果を文字列として結合
        context = "\n".join(flat_results)
        
        gemini_response = get_gemini_response(user_message, context)
        
        say(f"質問: {user_message}\n回答: {gemini_response}\nスプレッドシートリンク: {spreadsheet_link}")
    else:
        say(f"関連する情報が見つかりませんでした。\nスプレッドシートリンク: {spreadsheet_link}")

schedule.every(1).minute.do(update_chroma_from_sheets)

def run_scheduled_tasks():
    while True:
        schedule.run_pending()
        time.sleep(1)

threading_thread = Thread(target=run_scheduled_tasks)
threading_thread.daemon = True
threading_thread.start()

if __name__ == "__main__":
    SocketModeHandler(app, SLACK_APP_TOKEN).start()

修正点

• シートの項目名を変更し、読み込めるように修正

qa_pairs = [{"question": row[0], "answer": row[1]} for row in values if len(row) > 1]

• データベースにQA形式で登録できるように修正

documents = [f"質問: {qa['question']}\n回答: {qa['answer']}" for qa in qa_pairs]

おわりに

次はスプレッドシート内の各シートからデータ登録できるようにしたり、他の拡張子でも登録できるようにしたいです。