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()