Taste of Tech Topics

Acroquest Technology株式会社のエンジニアが書く技術ブログ

Amazon Bedrock Knowledge Baseのクエリフィルター自動生成で検索の精度を向上させる

こんにちは、機械学習チーム YAMALEX の駿です。
YAMALEX は Acroquest 社内で発足した、会社の未来の技術を創る、機械学習がメインテーマのデータサイエンスチームです。
(詳細はリンク先をご覧ください。)

この記事は Amazon Bedrock Advent Calendar 2024 17日目の投稿です。

re:Invent2024 ではたくさんの新機能が発表されて、あんなこともできるようになった、こんなことも……と興奮が止まらない日々です。

今回はそんなre:Inventで発表されたAmazon Bedrock Knowledge Baseの新機能の一つである、クエリフィルター自動生成を試してみました。

aws.amazon.com

今回試した構成

1. はじめに

1.1. Amazon Bedrock Knowledge Baseとは

Amazon Bedrock Knowledge Base (以下、Knowledge Base)はAmazonが提供する、RAGを簡単に構築するためのサービスです。

概要についてはこちらの記事で説明しているので、RAGやKnowledge Baseになじみのない方はまずはこちらをご一読ください。

acro-engineer.hatenablog.com

1.2. ベクトル検索の欠点

Knowledge Baseではデータベースとしてベクトルデータベースを使うことが多いです。
(ただし、アップデートで構造化されたデータを扱えるRedShiftをデータベースに使用可能になっていましたね)

ベクトル検索は文章の意味を捉えるのが得意なため、多少の表記ゆれや同義語を含んだ文章などでも近しい文章として取得することができる、という利点があります。

その反面、製品番号などで検索をしたい場合は、

  1. 製品番号が完全に一致する製品の情報のみ検索してほしい
  2. 製品番号は一文字違うだけで別の製品になってしまうため、一言一句正確に検索してほしい

などの要望が出てきますが、ベクトル検索は上に書いた利点が仇となり、上記の要望を満たせません。

しかし、今回のクエリフィルター自動生成機能をベクトル検索と組み合わせることで、意味的に近しい文章を取得しながら、 製品番号など正確に検索したい部分は正確に検索することが可能になりました。

2. [新機能]クエリフィルター自動生成とは

Bedrock Knowledge Baseにはドキュメントの同期時に、そのドキュメントの属性をメタデータとして付与する仕組みが2つ用意されています。

  1. 取り込みファイル名.metadata.json という命名メタデータ付与ファイルを用意しておく
  2. 同期時に実行されるLambda内でメタデータを付与する

クエリフィルター自動生成では、上記いずれかの方法でベクトルデータベースに取り込まれたメタデータに対して、絞り込みを行うためのクエリを生成AIを使って自動で生成します。

現在、クエリフィルター自動生成には Anthropic Claude 3.5 Sonnet モデルのみが使用可能です。

たとえば

次のような構造のレコードがBedrock Knowledge Baseで作成したOpenSearch Serverlessのインデックスに保存されているとします。

{
  "AMAZON_BEDROCK_TEXT": """{"製品番号": "A2-B3-C4", "分類": "テレビ", "製品名": "プラズマテレビ", "説明": "鮮やかな色と深い黒を実現するプラズマパネル", "特筆すべき内容": "プラズマパネル", "その他": "高コントラスト"}""",
  "item_id": "A2-B3-C4",
  "category": "テレビ"
}

ここでフィルターに使用できるメタデータitem_idcategory のふたつです。

クエリフィルター自動生成は、ユーザー入力にitem_idcategoryに相当する部分が存在するか否かを判断し、存在した場合、その項目をクエリフィルターとして、検索条件に追加します。

ユーザーが「製品番号:A42-B43-C44の特長を教えて」と聞いた場合、 item_id に相当する項目があるため、生成AIは下記のようなクエリを生成します。

{
  "query": "製品番号:A42-B43-C44の特長を教えて",
  "filter": {
    {"eq": {"key": "item_id", "value": "A42-B43-C44"}}
  }
}

このクエリは、item_idA42-B43-C44と一致するレコードを抽出した後に、「製品番号:A42-B43-C44の特長を教えて」でベクトル検索を行うことを意味するため、 確実に製品番号が一致する検索結果のみが返されることが保証されます。

あくまでベクトル検索に追加するフィルターという立ち位置のため、 ベクトル検索ではヒットしなかったレコードが、この機能を使うことでヒットするようになるわけではありません。

3. 実施

実際に上のような挙動をするか、試してみました。

下記手順で検証します。

  1. データ作成
  2. メタデータ付与Lambda作成
  3. Bedrock Knowledge Base作成
  4. 同期実行
  5. 検索

■使用した環境

No 項目
1 データベース OpenSearch Serverless
2 データ 家電のカタログを想定して生成したCSV(後述)
3 メタデータ 製品番号, 分類
4 クエリフィルター生成モデル Anthropic Claude 3.5 Sonnet V1

3.1. データ作成

データ作成にはAmazonが発表した新しい生成AIモデルである Nova Lite V1を使用しました。

家電のカタログを検索するユースケースを想定し、下のようなデータを生成しました。

製品番号,分類,製品名,説明,特筆すべき内容,その他
A1-B2-C3,テレビ,スマートテレビ,4K解像度で美しい映像を体験できる最新モデル,HDR対応,Bluetooth機能付き
A4-B5-C6,洗濯機,全自動洗濯機,大容量で洗浄力抜群、忙しい朝もラクラク,除菌機能,静音設計
A7-B8-C9,冷蔵庫,冷蔵冷凍庫,新鮮さを保つ冷却技術で食材の鮮度を長持ち,自動製氷機能,エコ運転モード
...

「製品番号で絞り込みたい」、「すぐ温まるドライヤーを調べたいのに、電子レンジが検索結果に含まれて困る」などのケースに対応できるかを確認するため、 「製品番号」、「分類」を列として用意してあります。

作成したCSVはBedrock Knowledge Baseで取り込むため、S3にアップロードしておきます。

3.2. メタデータ付与Lambda作成

次に、上記CSVの一行一行にフィルターするためのメタデータを付与する方法を説明します。

今回は、Knowledge Baseの同期時にLambda関数を実行してメタデータを付与する方法を使用しました。

今回はCSVの1行を1チャンクとし、 その「製品番号」、「分類」列を抽出し、それぞれ item_idcategory としてメタデータを付与します。

下はそのコード中のメタデータ付与を行っている関数です。

def build_contents(file_content: dict) -> list:
    """行ごとに分割、メタデータを付与する"""
    content_metadata = file_content["contentMetadata"]
    content_body = file_content["contentBody"]
    content_type = file_content["contentType"]

    blob = StringIO(content_body)
    reader = csv.reader(blob)
    header = next(reader)
    new_contents = []
    for row in reader:
        # 一行ごとにJSON文字列としてレコードを作成する
        d = {k: v for k, v in zip(header, row)}
        body = json.dumps(d, ensure_ascii=False)

        # 各レコードに対して、製品番号と分類をメタデータとして付与する
        m = copy.deepcopy(content_metadata)
        m["item_id"] = d["製品番号"]
        m["category"] = d["分類"]
        file_content = {
            "contentMetadata": m,
            "contentBody": body,
            "contentType": content_type,
        }
        new_contents.append(file_content)
    return new_contents

このLambdaの中身についてはこのブログの範囲外となるため、詳細は下記ドキュメントをご参照ください。

docs.aws.amazon.com

3.3. Bedrock Knowledge Base作成

下記手順で上記Lambdaを同期時に実行するKnowledge Baseを作成します。

  1. Bedrockコンソールのナレッジベースを開き「ナレッジベースを作成」>「Knowledge Base with vector store」を押下する
  2. 「ナレッジベースの詳細を入力」画面でナレッジベース名など設定を行い「次へ」を押下する
  3. 「データソースを設定」

    1. S3 URIには「3.1.」でCSVをアップロードした先のS3 URIを指定する
    2. Transformation functionとして「3.2.」でデプロイしたLambdaを指定する

      Transformation functionを設定する

    3. 「次へ」を押下する

  4. 組み込みモデルは「Titan Text Embeddings v2」、ベクトルデータベースは「新しいベクトルストアをクイック作成」を選択し、「次へ」を押下する

3.4. 同期実行

「3.3.」で作成されたデータソースを選択し、「同期」を実行します。

同期時に、クローリング後、インデクシング前に Transformation functionに設定したLambdaが実行されます。

ステータスが「Available」になるのを待ってください。

投入したデータはOpenSearch Serverlessのインデックスから確認できます。

item_idcategory がそれぞれ付与されているのが分かります。

{
  "took": 77,
  "timed_out": false,
  "_shards": {
    "total": 0,
    "successful": 0,
    "skipped": 0,
    "failed": 0
  },
  "hits": {
    "total": {
      "value": 35,
      "relation": "eq"
    },
    "max_score": 1,
    "hits": [
      {
        "_index": "bedrock-knowledge-base-default-index",
        "_id": "1%3A0%3Apbf3kpMB_yTemhSJpjEn",
        "_score": 1,
        "_source": {
          "AMAZON_BEDROCK_TEXT": """{"製品番号": "A4-B5-C6", "分類": "洗濯機", "製品名": "全自動洗濯機", "説明": "大容量で洗浄力抜群、忙しい朝もラクラク", "特筆すべき内容": "除菌機能", "その他": "静音設計"}""",
          "item_id": "A4-B5-C6",
          "category": "洗濯機"
        }
      },
      {
        "_index": "bedrock-knowledge-base-default-index",
        "_id": "1%3A0%3Ap7f3kpMB_yTemhSJpjEn",
        "_score": 1,
        "_source": {
          "AMAZON_BEDROCK_TEXT": """{"製品番号": "A10-B11-C12", "分類": "電子レンジ", "製品名": "オーブンレンジ", "説明": "高速加熱で時短調理が叶う、便利な機能満載", "特筆すべき内容": "フラットテーブル", "その他": "自動メニュー"}""",
          "item_id": "A10-B11-C12",
          "category": "電子レンジ"
        }
      },
...

3.5. 検索

Knowledge Baseの準備ができたので、実際に検索をしてみましょう。

検索は下記2パターンで、Pythonのboto3を使って試しました。

  1. クエリフィルター自動生成を使わない場合

     response = bedrock_agent.retrieve(
         knowledgeBaseId=KNOWLEDGE_BASE_ID,
         retrievalQuery={"text": query},
     )
    
  2. クエリフィルター自動生成を使う場合

     response = bedrock_agent.retrieve(
         knowledgeBaseId=KNOWLEDGE_BASE_ID,
         retrievalQuery={"text": query},
         retrievalConfiguration={
             "vectorSearchConfiguration": {
                 "implicitFilterConfiguration": {
                     "metadataAttributes": [
                         # ユーザークエリから抽出するメタデータの候補を記述します
                         {
                             "key": "item_id",
                             "type": "STRING",
                             "description": "id of the item.  製品番号. e.g. A45-B46-C47, A2-B3-C4, D1-E2-F3, ..."
                         },
                         {
                             "key": "category",
                             "type": "STRING",
                             "description": "category of the item. enum: テレビ, 洗濯機, 冷蔵庫, 電子レンジ, 掃除機, エアコン, 炊飯器, コーヒーメーカー, オーブンレンジ, ドライヤー"
                         },
                     ],
                     "modelArn": "anthropic.claude-3-5-sonnet-20240620-v1:0"
                 }
             }
         }
     )
    

3.4. 結果

  1. 製品番号でフィルター出来る「製品番号:A42-B43-C44の特長を教えて」

    1. 使わない場合

       [
        {
         "content": {"text": "{\"製品番号\": \"A39-B40-C41\", \"分類\": \"冷蔵庫\", \"製品名\": \"冷蔵冷凍庫\", \"説明\": \"新鮮さを保つ冷却技術で食材の鮮度を長持ち\", \"特筆すべき内容\": \"自動製氷機能\", \"その他\": \"エコ運転モード\"}"},
         "metadata": {"category": "冷蔵庫", "item_id": "A39-B40-C41"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A41-B42-C43\", \"分類\": \"掃除機\", \"製品名\": \"コードレス掃除機\", \"説明\": \"軽量で使いやすく、様々な床材に対応\", \"特筆すべき内容\": \"HEPAフィルター\", \"その他\": \"コードレス\"}"},
         "metadata": {"category": "掃除機", "item_id": "A41-B42-C43"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A2-B3-C4\", \"分類\": \"テレビ\", \"製品名\": \"プラズマテレビ\", \"説明\": \"鮮やかな色と深い黒を実現するプラズマパネル\", \"特筆すべき内容\": \"プラズマパネル\", \"その他\": \"高コントラスト\"}"},
         "metadata": {"category": "テレビ", "item_id": "A2-B3-C4"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A42-B43-C44\", \"分類\": \"エアコン\", \"製品名\": \"ルームエアコン\", \"説明\": \"快適な温度を保つ高効率モデル\", \"特筆すべき内容\": \"除湿機能\", \"その他\": \"リモコン操作\"}"},
         "metadata": {"category": "エアコン", "item_id": "A42-B43-C44"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A1-B2-C3\", \"分類\": \"テレビ\", \"製品名\": \"スマートテレビ\", \"説明\": \"4K解像度で美しい映像を体験できる最新モデル\", \"特筆すべき内容\": \"HDR対応\", \"その他\": \"Bluetooth機能付き\"}"},
         "metadata": {"category": "テレビ", "item_id": "A1-B2-C3"}
        }
       ]
      
    2. 使った場合

       [
        {
         "content": {"text": "{\"製品番号\": \"A42-B43-C44\", \"分類\": \"エアコン\", \"製品名\": \"ルームエアコン\", \"説明\": \"快適な温度を保つ高効率モデル\", \"特筆すべき内容\": \"除湿機能\", \"その他\": \"リモコン操作\"}"},
         "metadata": {"category": "エアコン", "item_id": "A42-B43-C44"}
        }
       ]
      

    使わない場合は、欲しかった製品番号の製品が4番目に来ているのに対し、クエリフィルター自動生成を使った場合は、欲しい商品のみが取得できています。

  2. カテゴリでフィルターできる「忙しい朝でも早く乾かせるドライヤー」

    1. 使わない場合

       [
        {
         "content": {"text": "{\"製品番号\": \"A36-B37-C38\", \"分類\": \"ドライヤー\", \"製品名\": \"ヘアドライヤー\", \"説明\": \"速乾で髪のダメージを軽減、使いやすい設計\", \"特筆すべき内容\": \"イオン機能\", \"その他\": \"軽量設計\"}"},
         "metadata": {"category": "ドライヤー", "item_id": "A36-B37-C38"}
         }
        },
        {
         "content": {"text": "{\"製品番号\": \"A28-B29-C30\", \"分類\": \"ドライヤー\", \"製品名\": \"ヘアドライヤー\", \"説明\": \"速乾で髪のダメージを軽減、使いやすい設計\", \"特筆すべき内容\": \"イオン機能\", \"その他\": \"軽量設計\"}"},
         "metadata": {"category": "ドライヤー", "item_id": "A28-B29-C30"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A46-B47-C48\", \"分類\": \"ドライヤー\", \"製品名\": \"ヘアドライヤー\", \"説明\": \"速乾で髪のダメージを軽減、使いやすい設計\", \"特筆すべき内容\": \"イオン機能\", \"その他\": \"軽量設計\"}"},
         "metadata": {"category": "ドライヤー", "item_id": "A46-B47-C48"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A30-B31-C32\", \"分類\": \"電子レンジ\", \"製品名\": \"オーブンレンジ\", \"説明\": \"高速加熱で時短調理が叶う、便利な機能満載\", \"特筆すべき内容\": \"フラットテーブル\", \"その他\": \"自動メニュー\"}"},
         "metadata": {"category": "電子レンジ", "item_id": "A30-B31-C32"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A40-B41-C42\", \"分類\": \"電子レンジ\", \"製品名\": \"オーブンレンジ\", \"説明\": \"高速加熱で時短調理が叶う、便利な機能満載\", \"特筆すべき内容\": \"フラットテーブル\", \"その他\": \"自動メニュー\"}"},
         "metadata": {"category": "電子レンジ", "item_id": "A40-B41-C42"}
        }
       ]
      
    2. 使った場合

       [
        {
         "content": {"text": "{\"製品番号\": \"A36-B37-C38\", \"分類\": \"ドライヤー\", \"製品名\": \"ヘアドライヤー\", \"説明\": \"速乾で髪のダメージを軽減、使いやすい設計\", \"特筆すべき内容\": \"イオン機能\", \"その他\": \"軽量設計\"}"},
         "metadata": {"category": "ドライヤー", "item_id": "A36-B37-C38"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A28-B29-C30\", \"分類\": \"ドライヤー\", \"製品名\": \"ヘアドライヤー\", \"説明\": \"速乾で髪のダメージを軽減、使いやすい設計\", \"特筆すべき内容\": \"イオン機能\", \"その他\": \"軽量設計\"}"},
         "metadata": {"category": "ドライヤー", "item_id": "A28-B29-C30"}
        },
        {
         "content": {"text": "{\"製品番号\": \"A46-B47-C48\", \"分類\": \"ドライヤー\", \"製品名\": \"ヘアドライヤー\", \"説明\": \"速乾で髪のダメージを軽減、使いやすい設計\", \"特筆すべき内容\": \"イオン機能\", \"その他\": \"軽量設計\"}"},
         "metadata": {"category": "ドライヤー", "item_id": "A46-B47-C48"}
        }
       ]
      

    使わなかった場合は、電子レンジなどドライヤーが関係のない製品が検索結果に含まれていますが、クエリフィルター自動生成のおかげて、ドライヤーのみが検索結果となっています。

まとめ

今回は、Amazon Bedrock Knowledge Baseの新機能である、クエリフィルター自動生成を試して、どれほど検索の精度が上がるのかを確認しました。

今までベクトル検索が苦手としていたキーワードでの一字一句一致する検索を検索時の設定を追加するだけで解決できるので非常に有用だと思います。

注意事項は下記2点です。

  1. メタデータは事前に付与しておく必要があります。

    ドキュメント毎に.metadata.json で付与するか、今回のようにLambdaを実行して付与する必要があります。

  2. 生成AIが生成したクエリは検索のレスポンスから知ることができません。

    精度が出ないときに実際に使われたクエリフィルターの内容が確認できないのは、デバッグする上で不便ですが、 使用しているプロンプトは公開されているため、そちらで別途出力させて確認が可能です。

re:Invent2024では、この他にも生成AI周りで多くのアップデートが発表されているため、情報キャッチアップして、取り入れていきたいですね。

Acroquest Technologyでは、キャリア採用を行っています。
  • Azure OpenAI/Amazon Bedrock等を使った生成AIソリューションの開発
  • ディープラーニング等を使った自然言語/画像/音声/動画解析の研究開発
  • マイクロサービス、DevOps、最新のOSSクラウドサービスを利用する開発プロジェクト
  • 書籍・雑誌等の執筆や、社内外での技術の発信・共有によるエンジニアとしての成長
少しでも上記に興味を持たれた方は、是非以下のページをご覧ください。 www.wantedly.com