【RAG】Slack上で質問してGeminiから回答を受け取りたい

こちらの記事の続きです！
カレンダーの投稿順ずれていたので内容前後しています

コード

from slack_bolt import App
from slack_bolt.adapter.socket_mode import SocketModeHandler
import os
from dotenv import load_dotenv
import chromadb
import google.generativeai as genai
import time

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')

app = App(token=SLACK_BOT_TOKEN)

documents = [
    "おれはジャイアン",
    "おまえのものはおれのもの、おれのものはおれのもの",
]

ids = ["id1", "id2"]

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

def check_data_in_collection():
    all_documents = collection.get()
    print("現在のコレクションの内容:")
    for doc in all_documents["documents"]:
        print(doc)

check_data_in_collection()

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)
    
    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}")
    else:
        say("関連する情報が見つかりませんでした。")

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

Slack上で質問するとChromaDBに検索をかけるように実装していますが、
ChromaDBのクエリ結果がリストのリスト形式になる場合があるため、以下のようにしています。

    if search_results and 'documents' in search_results:
       する
        flat_results = [item for sublist in search_results['documents'] for item in (sublist if isinstance(sublist, list) else [sublist])]
        
        # 結果を文字列として結合
        context = "\n".join(flat_results)

リストのリスト形式の場合

# 文書が分割されている場合
{
    "documents": [
        ["document_1_part_1", "document_1_part_2"],
        ["document_2_part_1", "document_2_part_2"]
    ],
    "distances": [0.1, 0.2]
}

リスト形式の場合

{
    "documents": [
        "document_1",
        "document_2"
    ],
    "distances": [0.1, 0.2]
}

Slack上で質問した結果

...以上！

おわりに

次はデータベースの登録周りをもう少し実用的にしてみたいです。