Amazon Bedrock ナレッジベースのメタデータ

やりたいこと

手動メタデータフィルタリングで、バージョンにあったドキュメントを検索し回答する。

構築

前提

ナレッジベースを Aurora で構築済

必要なファイルの準備

aws-cli-1.pdf と aws-cli-1.md.metadata.json

aws-cli-1.pdf は AWS のドキュメントからダウンロード。

aws-cli-1.pdf.metadata.json は次の通りとする。

aws-cli-1.pdf.metadata.json

{
	"metadataAttributes": {
		"target": "aws-cli",
		"ver": 1
	}
}

aws-cli-2.pdf と aws-cli-2.md.metadata.json

aws-cli-2.pdf は AWS のドキュメントからダウンロード。

aws-cli-2.pdf.metadata.json は次の通りとする。

aws-cli-2.pdf.metadata.json

{
	"metadataAttributes": {
		"target": "aws-cli",
		"ver": 2
	}
}

準備したファイルを S3 へアップロード

次の通り、aws-cli 配下にアップロードする。

Aurora にカラムを追加

クエリエディタよりずまずは、現在のカラムを確認します。この状態では、target と ver のカラムは存在しません。

target と ver のカラムを追加します。

target と ver のカラムが追加されたことを確認しておきます。

ナレッジベースでデータを同期

ナレッジベースでデータを同期します。

同期が完了しました。

API の作成

今回は、streamlit を使っています。

ポイント1: メタデータ
- AWS CLI のバージョンはユーザーに選択できるようする
ポイント2: パラメータ
- ユーザーで調節可能にする。（一貫性を持たせた回答にしたい場合は、埋め込みで良い）
ポイント3: promptTemplate
- ナレッジベースに検索結果が存在する場合は、参照したドキュメントのタイトルとどのページに記載があったかを記載する
- ナレッジベースに検索結果が存在しない場合は、【参考情報】と記載して回答する

retrieve_and_generate.py

import os
import boto3
import streamlit as st

def initialize_bedrock_client():
    region = "us-west-2"
    kb_id = "<ナレッジベースId>"
    client = boto3.client("bedrock-agent-runtime", region_name=region)
    return client, kb_id

def query_knowledge_base(client, kb_id, input_text, ver, temperature, top_p):
    filter = {
        "equals": {"key": "ver", "value": ver}
    }

    response = client.retrieve_and_generate(
        input={"text": input_text},
        retrieveAndGenerateConfiguration={
            "knowledgeBaseConfiguration": {
                "generationConfiguration": {
                    "promptTemplate": {
                        "textPromptTemplate": """
    
                        以下の検索結果を参考に回答してください：

                        '$search_results$'

                        回答フォーマット：
                        ---
                        【参照ドキュメント】
                        - 参照したドキュメントのタイトルを記載
                        - 参照したドキュメントのページ

                        【回答】
                        具体的な回答内容
                        ---

                        注意事項：
                        - 検索結果が存在する場合は、必ず参照したドキュメントのタイトルとどのページに記載されているかを記載すること
                        - 検索結果が存在しない場合は、「【参考情報】」と記載して回答すること
                        - 回答は上記のフォーマットに従って構造化すること
                        """
                    },
                    "inferenceConfig": {
                        "textInferenceConfig": {
                            "maxTokens": 4000,
                            "temperature": temperature,
                            "topP": top_p
                        }
                    }
                },
                "knowledgeBaseId": kb_id,
                "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-5-sonnet-20241022-v2:0",
                "retrievalConfiguration": {
                    "vectorSearchConfiguration": {
                        "filter": filter,
                        "numberOfResults": 5,
                    }
                },
            },
            "type": "KNOWLEDGE_BASE",
        },
    )
    
    return response["output"]["text"]

def main():
    st.title("AWS Bedrock Knowledge Base Query")
    
    # Initialize Bedrock client
    client, kb_id = initialize_bedrock_client()
    
    # Sidebar configurationså
    with st.sidebar:
        st.header("メタデータ")
        ver = st.selectbox("AWS CLI バージョン:", options=[2, 1], index=0)
        
        st.subheader("パラメータ")
        temperature = st.slider("温度", min_value=0.0, max_value=1.0, value=0.1, step=0.1)
        top_p = st.slider("トップ P", min_value=0.0, max_value=1.0, value=0.9, step=0.1)
    
    # Main area
    input_text = st.text_area("質問を入力してください:")
    
    if st.button("検索"):
        with st.spinner("回答を生成中..."):
            try:
                response_text = query_knowledge_base(client, kb_id, input_text, ver, temperature, top_p)
                st.write("### 回答")
                st.write(response_text)
            except Exception as e:
                st.error(f"エラーが発生しました: {str(e)}")

if __name__ == "__main__":
    main()

実行

次の通り実行します。

python3 -m streamlit run bedrock.py --server.port 8080

AWS CLI の実行結果

AWS CLI のインストール方法がしっかり返ってきました。ただし、どのページに記載があったかの記載は若干曖昧でもっと他の点で考慮が必要な気がします。

Azure CLI の実行結果

Azure CLI はナレッジベースにないので、promptTemplate に指示したとおり【参考情報】として返ってきました。

おまけ

retrieve_and_generate_stream 版

retrieve_and_generate_stream.py

import os
import boto3
import streamlit as st

def initialize_bedrock_client():
    region = "us-west-2"
    kb_id = "<ナレッジベース Id>"
    client = boto3.client("bedrock-agent-runtime", region_name=region)
    return client, kb_id

def query_knowledge_base(client, kb_id, input_text, ver, temperature, top_p):
    filter = {
        "equals": {
            "key": "ver", 
            "value": str(ver),
        }
    }

    response = client.retrieve_and_generate_stream(
        input={"text": input_text},
        retrieveAndGenerateConfiguration={
            "knowledgeBaseConfiguration": {
                "generationConfiguration": {
                    "promptTemplate": {
                        "textPromptTemplate": """
                        以下の検索結果を参考に回答してください：

                        '$search_results$'

                        回答フォーマット：
                        ---
                        【参照ドキュメント】
                        - 参照したドキュメントのタイトルを記載
                        - 参照したドキュメントのページ

                        【回答】
                        具体的な回答内容
                        ---

                        注意事項：
                        - 検索結果が存在する場合は、必ず参照したドキュメントのタイトルとどのページに記載されているかを記載すること
                        - 検索結果が存在しない場合は、「【参考情報】」と記載して回答すること
                        - 回答は上記のフォーマットに従って構造化すること
                        """
                    },
                    "inferenceConfig": {
                        "textInferenceConfig": {
                            "maxTokens": 4000,
                            "temperature": temperature,
                            "topP": top_p
                        }
                    }
                },
                "knowledgeBaseId": kb_id,
                "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-5-sonnet-20241022-v2:0",
                "retrievalConfiguration": {
                    "vectorSearchConfiguration": {
                        "filter": filter,
                        "numberOfResults": 5,
                    }
                },
            },
            "type": "KNOWLEDGE_BASE",
        },
    )
    
    return response.get("stream")

def main():
    st.title("AWS Bedrock Knowledge Base Query")
    
    # Initialize Bedrock client
    client, kb_id = initialize_bedrock_client()
    
    # Sidebar configurations
    with st.sidebar:
        st.header("メタデータ")
        ver = st.selectbox("AWS CLI バージョン:", options=[2, 1], index=0)
        
        st.subheader("パラメータ")
        temperature = st.slider("温度", min_value=0.0, max_value=1.0, value=0.1, step=0.1)
        top_p = st.slider("トップ P", min_value=0.0, max_value=1.0, value=0.9, step=0.1)
    
    # Main area
    input_text = st.text_area("質問を入力してください:")
    
    if st.button("検索"):
        with st.spinner("回答を生成中..."):
            try:
                stream = query_knowledge_base(client, kb_id, input_text, ver, temperature, top_p)
                if stream:
                    placeholder = st.empty()
                    full_response = ""
                    for event in stream:
                        if "output" in event:
                            chunk = event['output']['text']
                            full_response += chunk
                            placeholder.markdown(full_response)
            except Exception as e:
                st.error(f"エラーが発生しました: {str(e)}")

if __name__ == "__main__":
    main()