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

# 内容

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

![rag50.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2523652/5393dee3-416c-4aad-9b2a-b7a32683c67b.png)

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

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

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

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

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

![rag51.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2523652/8c3762bc-e05c-4b9a-8461-3b77be57f46e.png)

# Streamlitのプログラム

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

```python
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`に格納されます。

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

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

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

# 検索

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

![rag52.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2523652/7f635915-49e2-4a31-9786-7027ede47a8e.png)