【RAG】SlackからPDFデータをChromaDBに登録して関連情報を回答に含めたい

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

結果

• アップロードするPDF

• 仕様書という名のPDFをChromaDBに登録する前

• 仕様書という名のPDFをChromaDBに登録した後

しっかりPDFの内容について回答してくれていますね！

コード

import os
import requests
import time
import chromadb
from slack_bolt import App
from slack_sdk import WebClient
import PyPDF2
from io import BytesIO
from dotenv import load_dotenv
import google.generativeai as genai
from threading import Thread

load_dotenv()

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

client = WebClient(token=SLACK_BOT_TOKEN)

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

genai.configure(api_key=GEMINI_API_KEY)

app = App(token=SLACK_BOT_TOKEN)

@app.event("file_shared")
def handle_file_shared(event, say):
    try:
        file_id = event['file_id']
        response = client.files_info(file=file_id)
        file_info = response['file']
        file_url = file_info['url_private']

        file_content = download_file_from_slack(file_url)
        if not file_content:
            say("ファイルのダウンロードに失敗しました。")
            return

        text_content = extract_text_from_pdf(file_content)
        if not text_content:
            say("PDFファイルの内容を抽出できませんでした。")
            return

        document_id = file_id
        add_document_to_chroma(text_content, document_id)

        say(f"PDFファイルの内容をChromaDBに登録しました。")
    except Exception as e:
        say(f"エラーが発生しました: {str(e)}")

def download_file_from_slack(file_url):
    headers = {'Authorization': f'Bearer {SLACK_BOT_TOKEN}'}
    response = requests.get(file_url, headers=headers)
    return response.content if response.status_code == 200 else None

def extract_text_from_pdf(file_content):
    try:
        pdf_reader = PyPDF2.PdfReader(BytesIO(file_content))
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text()
        return text
    except Exception as e:
        print(f"Error extracting PDF text: {e}")
        return ""

def add_document_to_chroma(text_content, document_id):
    try:
        collection.add(documents=[text_content], ids=[document_id])
        print(f"Document {document_id} added to ChromaDB.")
    except Exception as e:
        print(f"Error adding document to ChromaDB: {e}")

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

def get_gemini_response(question, context):
    try:
        prompt = f"""
        以下の情報をもとに、質問に対する原因と対応内容をそれぞれ明確に答えてください:
        {context}
        質問: {question}
        原因:
        対応内容:
        """
        model = genai.GenerativeModel("gemini-1.5-flash")
        response = model.generate_content(prompt)
        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:
        flat_results = [item for sublist in search_results for item in sublist]

        context = "\n".join(flat_results)

        gemini_response = get_gemini_response(user_message, context)
        
        cause = ""
        solution = ""
        if '原因:' in gemini_response and '対応内容:' in gemini_response:
            try:
                cause = gemini_response.split('原因:')[1].split('対応内容:')[0].strip()
                solution = gemini_response.split('対応内容:')[1].strip()
            except IndexError:
                cause = "情報が不足しています。"
                solution = "情報が不足しています。"
        else:
            cause = "原因が記載されていません。"
            solution = "対応内容が記載されていません。"

        formatted_message = f"""
        問い合わせ内容: 
        {user_message}

        **考えられる原因:** {cause}  
        **考えられる対応内容:** {solution}  

        関連する過去の情報:
        {context.strip()}
        """

        say(formatted_message)
    else:
        say("関連情報が見つかりませんでした。")
        
if __name__ == "__main__":
    from slack_bolt.adapter.socket_mode import SocketModeHandler
    SocketModeHandler(app, SLACK_APP_TOKEN).start()

※プロンプトはこれまでの内容そのままです

おわりに

次は画像データのテキスト化などにもチャレンジしたいです。