Amazon Bedrock ナレッジベースでメタデータを使ってフィルタリングする

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

https://aws.amazon.com/jp/about-aws/whats-new/2024/03/knowledge-bases-amazon-bedrock-metadata-filtering/

https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/kb-test-config.html#kb-test-config-filters

# やりたいこと

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

# 構築

## 前提

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

## 必要なファイルの準備

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

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

https://docs.aws.amazon.com/pdfs/cli/v1/userguide/aws-cli-v1.pdf

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 のドキュメントからダウンロード。

https://docs.aws.amazon.com/pdfs/cli/latest/userguide/aws-cli.pdf

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

```aws-cli-2.pdf.metadata.json
{
	"metadataAttributes": {
		"target": "aws-cli",
		"ver": 2
	}
}
```

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

次の通り、aws-cli 配下にアップロードする。
![スクリーンショット 2025-01-10 16.04.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/8757134b-5ee0-928d-0466-017c6aefcd96.png)

## Aurora にカラムを追加

クエリエディタよりずまずは、現在のカラムを確認します。この状態では、target と ver のカラムは存在しません。
![スクリーンショット 2025-01-10 17.02.10.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/aa3bc9c1-9cd6-8b72-002b-1fa4c5a5d1fe.png)

target と ver のカラムを追加します。
![スクリーンショット 2025-01-10 17.04.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/5eb30204-42b9-10c7-567b-86bbab718f85.png)

target と ver のカラムが追加されたことを確認しておきます。
![スクリーンショット 2025-01-10 17.04.18.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/d0061d44-be82-8c88-7c94-f5ebc18af29d.png)

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

ナレッジベースでデータを同期します。
![スクリーンショット 2025-01-10 17.08.52.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/406f5b90-14f1-7c1b-c049-a80a4c2ad494.png)

同期が完了しました。
![スクリーンショット 2025-01-10 17.29.23.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/d2775bfd-4841-b5b7-c138-302f2043fad0.png)

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

![スクリーンショット 2025-01-11 17.05.45.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/38789c06-4d3f-03eb-74a3-f634f87c7d6f.png)


![bedrock.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/b0076f29-c3a1-7969-17e9-3a50746f3181.gif)

### Azure CLI の実行結果

Azure CLI はナレッジベースにないので、promptTemplate に指示したとおり【参考情報】として返ってきました。
![スクリーンショット 2025-01-11 17.11.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/8603c17a-d599-b482-4e7a-fa2ff6b092c6.png)

# おまけ

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