BedrockとStreamlitでメタデータによるドキュメントを絞るRAG検索

内容

Amazon BedrockとStreamlitを使用したRAGシステムを構築します。既に検索したいドキュメントが分かっている場合、全てのドキュメントから検索するよりもドキュメントを絞った方が検索性能が向上する可能性があります。

この機能を実装するため、検索対象のドキュメントにtitleというメタデータを付与します。メタデータの値を使用して検索できる仕組みを追加します。

ベクトルDBに列を追加

ベクトルDBとしてAuroraを使用しています。以下のSQLを実行してメタデータ用の列をテーブルに追加します。

ALTER TABLE bedrock_knowledgebase.dodeuce
ADD COLUMN title VARCHAR(128);

メタデータファイルの追加

ベクトルDBに挿入するドキュメントをS3上に保存します。例としてdocument01.pdfやdocument02.pdfを作成していますがこの際、メタデータ用のファイルを追加します。メタデータファイル名はドキュメント名.metadata.jsonになります。このファイルの中にtitleというメタデータを追加しておきます。

Streamlitのプログラム

下記がプログラム全文になります。

import streamlit as st
import boto3
import json

BEDROCK_MODEL = 'anthropic.claude-3-haiku-20240307-v1:0'
KNOWLEDGEBASE_ID = '******'
MODEL_ARN = "arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-haiku-20240307-v1:0"

st.set_page_config(page_title="テスト")
if "messages" not in st.session_state:
    st.session_state["messages"] = []

st.markdown("<p style='text-align: center; font-size: 36px; font-weight: bold;'>テスト検索</p>", unsafe_allow_html=True)

with st.sidebar:
    st.markdown("### 検索対象ドキュメントを選択してください")
    doc_titles = [
        "ドキュメントA",
        "ドキュメントB",
        "ドキュメントC",       
    ]
    selected_title = st.selectbox("検索対象ドキュメントのタイトル", doc_titles)

for message in st.session_state["messages"]:
    with st.chat_message(message["role"]):
        st.markdown(message["text"])

user_input = st.chat_input("質問を入力してください")
if user_input:
    st.session_state["messages"].append({"role": "user", "text": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    bedrock_agent = boto3.client("bedrock-agent-runtime", region_name='ap-northeast-1')

    context_prompt = f"""
    現在のユーザーの質問は以下です：
    <question>{user_input}</question>
    """
    response = bedrock_agent.retrieve_and_generate_stream(
        input={"text": context_prompt},
        retrieveAndGenerateConfiguration={
            "type": "KNOWLEDGE_BASE",
            "knowledgeBaseConfiguration": {
                "knowledgeBaseId": KNOWLEDGEBASE_ID,
                "modelArn": MODEL_ARN,
                "retrievalConfiguration": {
                    "vectorSearchConfiguration": {
                        "filter": {
                            "equals": {
                                "key": "title",
                                "value": selected_title
                            }
                        },
                    }
                }
            }
        }
    )

    stream = response.get("stream")
    bot_response = ""
    # citations = []

    if stream:
        with st.chat_message("assistant"):
            bot_area = st.empty()

            for event in stream:
                if "output" in event:
                    bot_response += event['output']['text']
                    bot_area.markdown(bot_response)

    st.session_state["messages"].append({"role": "assistant", "text": bot_response})

下記でサイドバーにリストを追加しています。選択したタイトルはselected_titleに格納されます。

with st.sidebar:
    st.markdown("### 検索対象ドキュメントを選択してください")
    doc_titles = [
        "ドキュメントA",
        "ドキュメントB",
        "ドキュメントC",       
    ]
    selected_title = st.selectbox("検索対象ドキュメントのタイトル", doc_titles)

selected_titleを使ってフィルタリングを行い、対象ドキュメントに対してのみ検索が実行されます。

"vectorSearchConfiguration": {
    "filter": {
        "equals": {
            "key": "title",
            "value": selected_title
        }
    },
}

検索

「ドキュメントA」を選択した状態で「ラーメン屋さん」と検索すると、「ドキュメントA」内からのみ該当箇所が抽出されます。