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

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

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("
'text-align: center; font-size: 36px; font-weight: bold;'>テスト検索", 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"""
    現在のユーザーの質問は以下です：
    {user_input}
    """
    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})

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

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