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

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

前回の記事ではデータベースの自動更新を行い、Geminiからの回答結果が変わることを同じ質問で試しました。
更新するデータはスクリプト内に記述していましたが、今回はSlackで回答を受け取った際にGoogleスプレッドシート経由で更新できるように改良してみました。

スプレッドシート更新前と更新後の回答の変化

• あなたはだれ

• ジャイアンはどんな人ですか

データベースが更新されて、質問に関連する情報を回答してくれていることが確認できました。
これで、質問に対する回答精度がいまいちだった際に、求める回答をシートに記載するだけで次に活きます。

コード

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 = "シート1!A2:A4"

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 []

        documents = [row[0] for row in values]
        return documents

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

def add_documents_to_chroma(documents, ids):
    try:
        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()

注意点

• SAMPLE_RANGE_NAMEはファイル名ではなく、シート名とセル範囲を指定
  ファイル名を指定して沼りました

• シートのアクセス権限許可があることを確認

• credentials.jsonをsample.pyと同じ階層に置く

おわりに

無事スプレッドシート経由でデータベースを更新することに成功しました！
次は登録するデータをQA形式にしてみたいです