【RAG】データを自動更新して回答結果が変わることを確認したい

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

https://qiita.com/does_not_exist/private/ac915f050dbff18974ef

Slackから質問を投げるとChromaDBに検索をかけ、関連情報をもとにGeminiから回答が返ってくるところまで確認ができています。

次は事前データを更新しやすくしたいので改良してみました。

# 事前データ(更新前)

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


# 質問と回答結果

![スクリーンショット 2024-12-02 22.28.50.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2889192/9d2d477c-9031-d7d6-2b02-1c9552951ca1.png)


# 事前データ(更新後)

```
new_documents = [
    "ドラえもんは未来から来たロボット猫です。",
    "ジャイアンは実はやさしいです",
]
```

# 質問と回答結果

![スクリーンショット 2024-12-02 22.29.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2889192/6b3574b8-1510-030c-d028-155fa03679d4.png)



# コード

```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
import schedule

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

def add_documents_to_chroma(documents, ids):
    try:
        collection.add(documents=documents, ids=ids)
        print("データがChromaDBに追加されました。")
    except Exception as e:
        print(f"エラーが発生しました: {e}")

add_documents_to_chroma(documents, ids)

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

def update_data_periodically():
    new_documents = [
        "ドラえもんは未来から来たロボット猫です。",
        "ジャイアンは実はやさしいです",
    ]
    new_ids = ["id3", "id4"]

    add_documents_to_chroma(new_documents, new_ids)

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

schedule.every(1).minutes.do(update_data_periodically)

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

if __name__ == "__main__":
    from threading import Thread
    threading_thread = Thread(target=run_scheduled_tasks)
    threading_thread.daemon = True
    threading_thread.start()

    SocketModeHandler(app, SLACK_APP_TOKEN).start()
```

# 事前データの更新処理

今回はすぐ検証したかったので1分おきに自動更新をしています

```python
schedule.every(1).minutes.do(update_data_periodically)
```

1分経つと以下のようなメッセージがターミナルに表示されます。

![スクリーンショット 2024-12-02 22.18.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2889192/b353df91-ac3e-1d9d-9709-42d855a60355.png)

さらに1分経つとこのようなメッセージが。

![スクリーンショット 2024-12-02 22.19.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2889192/d30ffc0e-eaed-73ea-1672-45ffbfc65839.png)


# おわりに
データベースを自動更新できるようになりました！
次は事前データの登録をSlackから行えるようにしたいですね:smirk: