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

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

https://qiita.com/does_not_exist/private/0f8962442a4c4df101fb

# コード

```python
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のクエリ結果がリストのリスト形式になる場合があるため、以下のようにしています。

```python
    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上で質問した結果

![スクリーンショット 2024-12-02 21.26.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2889192/dfcd69d8-c40e-b43e-221a-6d85ee0ec3a9.png)

![スクリーンショット 2024-12-02 22.07.12.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2889192/6b238c11-4988-994f-a76d-1012e5144f0e.png)


...以上！

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