Amazon Bedrock Knowledge Basesの ⾃動クエリフィルタ機能とは

# 内容

AWS re:Invent 2024 で発表があった Amazon Bedrock Knowledge Basesの ⾃動クエリフィルタ機能に対応したことについてシンプルに纏めます。

https://aws.amazon.com/jp/about-aws/whats-new/2024/12/amazon-bedrock-knowledge-bases-auto-generated-query-filters-improved-retrieval/

手動で複雑なフィルタを作成することなく、ドキュメントのメタデータに基づいて自動で関連する検索結果を絞り込むことが可能になる

（例）
「ワシントン州で申請する方法」というクエリでは、自動的に「ワシントン州」というフィルターが適用され、関連する文書のみを取得できるイメージ


# 利用可能な LLM

- Anthropic Claude 3.5 sonnet


# 利用可能なリージョン

- 米国東部 (バージニア北部)
- 米国西部 (オレゴン)
- アジアパシフィック (東京)
- アジアパシフィック (ソウル)
- 欧州 (フランクフルト)
- 欧州 (チューリッヒ)
- AWS GovCloud (米国西部) 


# 設定

`implicitFilterConfiguration` で行えば良さそう。

https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html

- 暗黙的なメタデータフィルタリング

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


# 動作確認

## データソースの配置

S3 に 2 つのファイルを配置します。

`s3://<S3バケット名>/data/aws-cli-1.pdf`

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

`s3://<S3バケット名>/data/aws-cli-2.pdf`

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

## メタデータの配置

`s3://<S3バケット名>/aws-cli//aws-cli-1.pdf.metadata.json`

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

`s3://<S3バケット名>/aws-cli/aws-cli-2.pdf.metadata.json`

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

## ナレッジベースへのデータ同期

データソースから「同期」します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/f03d5020-7b30-7a2f-6911-b644fc282eca.png)

## IMA ポリシーの追加

Claude 3.5 Sonnet で自動クエリフィルターを生成するめに、Knowledge Bases の IAM ロールの IAM ポリシーに追加が必要です。(図の 10 行目)

![スクリーンショット 2025-01-13 16.52.28.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/c153a9d9-45a8-8bcc-b48b-511769b5254e.png)

## API の作成

基本は以前に作成した以下の記事の API とほぼ同じで、　`vectorSearchConfiguration` で動作させるようにしています。

https://qiita.com/leomaro7/items/f123bf18e0c7b03d08ab

```bedrock.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, temperature, top_p):
    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-20240620-v1:0",
                "retrievalConfiguration": {
                    "vectorSearchConfiguration": {
                        "numberOfResults": 5,
                        'implicitFilterConfiguration': {
                            'metadataAttributes': [
                                {
                                    'description': 'AWS CLI ver',
                                    'key': 'ver',
                                    'type': 'STRING'
                                },
                            ],
                            'modelArn': 'arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-5-sonnet-20240620-v1:0'
                        },
                    }
                },
            },
            "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.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, temperature, top_p)
                st.write("### 回答")
                st.write(response_text)
            except Exception as e:
                st.error(f"エラーが発生しました: {str(e)}")

if __name__ == "__main__":
    main()
```

## 検証結果

### Prompt「AWS CLI  のインストール手順」

「AWS Command Line Interface バージョン 2 用ユーザーガイド」が参照されました。

![スクリーンショット 2025-01-13 16.58.34.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/1234ea71-c5b2-8a51-943b-1a805a79f760.png)

ちょっと結果だけではわかりにくいのでログを確認してみましょう。filter は null のようです。

```json
    "output": {
        "outputContentType": "application/json",
        "outputBodyJson": {
            "output": {
                "message": {
                    "role": "assistant",
                    "content": [
                        {
                            "text": "```json\n{\n  \"query\": \"AWS CLI のインストール手順\",\n  \"filter\": null\n}\n```"
                        }
                    ]
                }
```

### Prompt「AWS CLI 1 のインストール手順」

「AWS Command Line Interface バージョン 1 用ユーザーガイド」が参照されました。

![スクリーンショット 2025-01-13 17.01.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280929/d82071bd-1d02-b304-3194-57a5283799d9.png)

こちらもログを確認してみましょう。 ver 1 を filter してくれているように見えました。

```json
        "outputBodyJson": {
            "output": {
                "message": {
                    "role": "assistant",
                    "content": [
                        {
                            "text": "```json\n{\n  \"query\": \"AWS CLI インストール手順\",\n  \"filter\": {\n    \"eq\": {\n      \"key\": \"ver\",\n      \"value\": \"1\"\n    }\n  }\n}\n```"
                        }
                    ]
                }
            },
```