Taste of Tech Topics

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

Amazon Bedrock Knowledge Base の構造化データ取得はどこまで複雑なクエリに対応できるか

AWS Bedrock 生成AI RAG

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

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

Amazon Bedrock Knowledge BaseのデータベースとしてRedshiftを指定して、構造化データを検索し、その結果を元に回答を生成することができるようになりました。 現在のところ、データベースとしてはRedshiftのみに対応しています。
Redshiftであれば、Serverless版、クラスタ版いずれも選択が可能です。

この記事ではRedshiftにサンプルとして用意されているTICKITデータセットを用いて、複雑なテーブルに対してどのようなSQLが生成できるのか、生成されたSQLを調整したい場合に何ができるか、を確認します。

今回試した構成

1. はじめに

1.1. Amazon Bedrock Knowledge Baseとは

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

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

acro-engineer.hatenablog.com

1.2. 従来のRAGの欠点

Knowledge Baseではデータベースとしてベクトルデータベースを使うことが多いです。

ベクトル検索は文章の意味を捉えるのが得意なため、それらしい内容をいくつかベクトルデータベースから取得し、回答を生成する、ということが上手くできます。

その反面、最も多い事象は何か、条件に一致するレコードが何件あるか、などデータの件数や正確な数値を扱うような情報を把握する必要がある質問には答えることができない、という欠点もあります。

しかし、今回の構造化データ検索を行うと、Redshift上のデータに対して、ユーザーの質問からSQLを作成し実行することで、上に書いたような質問にも適切に回答できるようになります。

2. [新機能]構造化データ検索とは

先日のre:Invent2024で発表された新機能として、Bedrock Knowledge BaseのデータベースとしてRedshiftを使用し、構造化データから取得した情報を元にRAGを行うことが可能になりました。

Bedrock Knowledge Baseの構造化データ検索は下記の流れで動きます。

  1. テーブル、データがあらかじめ投入されたRedshiftデータベースを用意する
  2. Bedrock Knowledge BaseにRedshift上のデータベースのテーブル、カラム情報を取り込む
  3. ユーザーが質問をする
  4. Bedrockが、ユーザーの質問に答えるために必要な情報をRedshiftのデータベースから取得するためのSQL文を生成する
  5. RedshiftでSQLを実行し、結果を取得する
  6. 取得した結果を元にユーザーの質問に回答する

従来のRAGと変わったのは、「4.」「5.」の部分で、SQL文で検索することで、「上位５件」や「合計件数」など、ベクトル検索では取得できない情報を取得できるようになりました。

3. 実施

今回は実際にどんなSQLが生成できるのか、テーブル／カラムの説明を追加することで、どれくらい精度を上げられるのか、を確認します。

下記手順で検証します。

  1. Redshiftデータベース準備
  2. Bedrock Knowledge Base作成
  3. SQL生成
  4. 説明を追加

  5. SQL再生成

    → より良いSQLが生成できていることを確認する

■使用した環境

No 項目 値
1 データベース Redshift
2 データ TICKIT データベース（※1）

※1 TICKITデータベースはRedshiftがサンプルとして用意しているデータベースで、スポーツやコンサートのチケットの売り上げ情報などを検索することができます。

扱う質問について

今回はKnowledge BaseのコンソールからSQLを生成させるために下記の質問を用意しました。

いずれもRedshiftでTICKITのデータをロードしたときにサンプルのクエリとして用意されていたSQLのようなものが生成されることを期待してます。

①「日付ごとの売り上げ上位5件を教えて」

■想定する参考SQL

SELECT d.caldate, sum(s.pricepaid) as sumpricepaid
FROM tickit.sales as s
JOIN tickit.date as d on s.dateid = d.dateid
GROUP BY d.caldate
ORDER BY sumpricepaid DESC
LIMIT 5;

■想定する検索結果

No caldate sumpricepaid
1 2008-04-14 407440
2 2008-09-10 402603
3 2008-07-27 400729

②「最も多くのチケットを購入した人は誰？」

■想定する参考SQL

SELECT firstname, lastname, total_quantity 
FROM   (SELECT buyerid, sum(qtysold) total_quantity
    FROM  tickit.sales
    GROUP BY buyerid
    ORDER BY total_quantity desc limit 10) Q, tickit.users
WHERE Q.buyerid = userid
ORDER BY Q.total_quantity desc;

■想定する検索結果

No firstname lastname total_quantity
1 Jerry Nichols 67
2 Armando Lopez 64
3 Kameko Bowman 64

3.1. Redshiftデータベース準備

こちらの手順に従って、サンプルデータをロードします。

docs.aws.amazon.com

次のようなノートブックが開いたら、準備完了です。

サンプルのSQLクエリが用意されている

3.2. Bedrock Knowledge Base作成

上記、Redshiftのデータベースを参照するKnowledge Baseを作成します。

  1. Bedrockコンソールのナレッジベースを開き「ナレッジベースを作成」>「Knowledge Base with structured data store」を押下する

  2. 「ナレッジベースの詳細を入力」画面でナレッジベース名など設定を行い「次へ」を押下する。

    後ほど必要になるため、ナレッジベースに紐づけるIAMサービスロール名をメモしておく。

  3. 「Configure query engine」画面で下記を設定して「次へ」を押下する。

    No 項目 値
    1 Query engine details.Connection options Redshift serverless
    2 Query engine details.Workgroup default-workgroup
    3 Authentication IAM Role
    4 Default storage metadata.Options Amazon Redshift databases
    5 Default storage metadata.Database sample_data_dev
    6 Query configurations 何も設定しない

  4. 設定を確認して、Knowledge Baseを作成する

  5. Redshiftのノートブック上で下記コマンドを実行し、Knowledge Baseに紐づけたIAMサービスロールがRedshiftにアクセスできるようにする

     CREATE USER "IAMR:{ロール名}" WITH PASSWORD DISABLE;

  6. Knowledge Baseのコンソールから「同期」を実行する

3.3. SQL生成

Bedrock Knowledge Baseのコンソールから、実際にSQLを生成した結果を以下に示します。

①「日付ごとの売り上げ上位5件を教えて」

SQLが生成された

■生成されたSQL

SELECT d.caldate AS "Date", SUM(s.qtysold * s.pricepaid) AS "Total Sales"
FROM tickit.sales s
JOIN tickit.date d ON s.dateid = d.dateid
GROUP BY d.caldate
ORDER BY "Total Sales" DESC
LIMIT 5;

■検索結果

No date total sales
1 2008-09-10 1036333
2 2008-04-14 1028582
3 2008-07-14 1019392

■問題点

  1. pricepaid は チケット単価 x チケット枚数 の値のため、改めて qtysold をかけてはいけない

上記の問題があるため、想定する検索結果とは全く異なる結果となりました。
チケット単価 x チケット枚数 ^ 2 を計算してしまっています。

②「最も多くのチケットを購入した人は誰？」

SQLが生成された

■生成されたSQL

SELECT u.firstname, u.lastname, SUM(s.qtysold) AS "Total Tickets Sold" 
FROM tickit.sales s
JOIN tickit.users u ON s.buyerid = u.userid 
GROUP BY u.firstname, u.lastname
ORDER BY "Total Tickets Sold" DESC 
LIMIT 1;

■検索結果（上位5件に編集して実行）

No firstname lastname total tickets sold
1 Charity Gamble 69
2 Jerry Nichols 67
3 Josephine Robbins 67

■問題点

  1. firstname + lastname の組は重複が存在するため、これだけでGROUP BYしてはいけない。

こちらは重複する可能性がある姓名でGROUP BYをしているため、異なる都市（BoiseとOverland Park）にそれぞれ存在するCharity Gamble が上位に来てしまいました。

続いて、この問題を解決するためにどのような設定が有効かを確認していきます。

3.4. SQL生成の精度を向上する

Bedrock Knowledge BaseではRedshiftのテーブル、カラムに対して、説明を追加する機能があります。

また、検索には使わないテーブル、カラムの指定、サンプルクエリを設定する機能もあり、 これらを組み合わせることで、SQL生成の精度を向上させることが可能です。

今回は、その中でも下記２点を確認しました。

  1. サンプルクエリを追加して上記問題を解決できるか
  2. カラム説明を追加して上記問題を解決できるか

3.4.1. サンプルクエリを追加

まずはTICKITのサンプルとして用意されていたクエリをBedrock Knowledge Baseの Curated queries に指定して、SQL生成の精度が上がるかを確認しました。

サンプルで用意されていたSQLクエリを例として追加

  1. Curated queries にサンプルSQLを設定して、「送信」を押下

  2. Query engine を「同期」

    Redshiftには変更を加えていないため、不要だと思うが、念のため

  3. 「3.3.」と同様にSQL生成を行う

■生成されたSQL

①「日付ごとの売り上げ上位5件を教えて」

SELECT d.caldate AS "Date", SUM(s.qtysold * s.pricepaid) AS "Total Sales"
FROM tickit.sales s
JOIN tickit.date d ON s.dateid = d.dateid
GROUP BY d.caldate
ORDER BY "Total Sales" DESC
LIMIT 5;

→ 変わりませんでした。。。

②「最も多くのチケットを購入した人は誰？」

SELECT u.firstname, u.lastname, SUM(s.qtysold) AS "Total Tickets Sold" 
FROM tickit.sales s
JOIN tickit.users u ON s.buyerid = u.userid 
GROUP BY u.firstname, u.lastname
ORDER BY "Total Tickets Sold" DESC 
LIMIT 1;

→ 変わりませんでした。。。

3.4.2. カラム説明を追加

誤った使い方をされていたカラムについて説明を追加

  1. Descriptions にテーブル名とカラム名、その説明を入力して、「送信」を押下

  2. Query engine を「同期」

    Redshiftには変更を加えていないため、不要だと思うが、念のため

  3. 「3.3.」と同様にSQL生成を行う

No Table name Column name Descriptions
1 sample_data_dev.tickit.sales pricepaid the total revenue. pricepaid / qtysold = price of a ticket. do not multiply pricepaid by qtysold when calculating total revenue.
2 sample_data_dev.tickit.users firstname name columns cannot be used to identify a user. there are users with same first and last names
3 sample_data_dev.tickit.users lastname name columns cannot be used to identify a user. there are users with same first and last names

■生成されたSQL

①「日付ごとの売り上げ上位5件を教えて」

SELECT d.caldate AS "Date", SUM(s.pricepaid) AS "Total Sales"
FROM tickit.sales s
JOIN tickit.date d ON s.dateid = d.dateid
GROUP BY d.caldate
ORDER BY "Total Sales" DESC
LIMIT 5;

→ ○ 売り上げの計算方法が訂正された

②「最も多くのチケットを購入した人は誰？」

SELECT u.firstname, u.lastname, SUM(s.qtysold) AS "Total Tickets Sold" 
FROM tickit.sales s
JOIN tickit.users u ON s.buyerid = u.userid 
GROUP BY u.firstname, u.lastname
ORDER BY "Total Tickets Sold" DESC 
LIMIT 1;

→ 変わりませんでした。。。

No 質問 期待するSQL 追加設定 生成されたSQL 結果 説明
1 日付ごとの売り上げ上位5件を教えて SELECT d.caldate, sum(s.pricepaid) as sumpricepaid
FROM tickit.sales as s
JOIN tickit.date as d on s.dateid = d.dateid
GROUP BY d.caldate
ORDER BY sumpricepaid DESC
LIMIT 5; 		なし SELECT d.caldate AS "Date", SUM(s.qtysold * s.pricepaid) AS "Total Sales"
FROM tickit.sales s
JOIN tickit.date d ON s.dateid = d.dateid
GROUP BY d.caldate
ORDER BY "Total Sales" DESC
LIMIT 5; 		× 売り上げの計算方法が間違っている
2 あり SELECT d.caldate AS "Date", SUM(s.pricepaid) AS "Total Sales"
FROM tickit.sales s
JOIN tickit.date d ON s.dateid = d.dateid
GROUP BY d.caldate
ORDER BY "Total Sales" DESC
LIMIT 5; 		○ 売り上げの計算方法を修正することができた
3 最も多くのチケットを購入した人は誰？ SELECT firstname, lastname, total_quantity
FROM (SELECT buyerid, sum(qtysold) total_quantity
  FROM tickit.sales
  GROUP BY buyerid
  ORDER BY total_quantity desc limit 10) Q, tickit.users
WHERE Q.buyerid = userid
ORDER BY Q.total_quantity desc; 		なし SELECT u.firstname, u.lastname, SUM(s.qtysold) AS "Total Tickets Sold"
FROM tickit.sales s
JOIN tickit.users u ON s.buyerid = u.userid
GROUP BY u.firstname, u.lastname
ORDER BY "Total Tickets Sold" DESC
LIMIT 1; 		× 集計時に「姓名」で判定してしまっている
4 あり SELECT u.firstname, u.lastname, SUM(s.qtysold) AS "Total Tickets Sold"
FROM tickit.sales s
JOIN tickit.users u ON s.buyerid = u.userid
GROUP BY u.firstname, u.lastname
ORDER BY "Total Tickets Sold" DESC
LIMIT 1; 		× 集計時の問題を是正できなかった

テーブルの中に類似の概念を持つカラム（イベントIDとイベント名など）が複数存在する場合に、うまく使い分けることができていないと言えると思います。

反対に必要な情報が複数テーブルにまたがっている場合などは適切にJOINなど使って情報を集計することができており、そこの精度は高く、再現性が高い、と感じました。
また、計算方法が間違っている部分に関してもカラムの説明を追加することで、是正できることが分かりました。

まとめ

今回は、Bedrock Knowledge Baseに追加されたRedshiftと連携する機能を使って、どこまでできるのかを確認しました。

簡単なテーブル・簡単な質問であれば、SQLを生成して正しい情報で回答を生成できますが、 複雑なテーブル・難しい質問をされた場合に、SQLの生成が上手くいかない（正しくないクエリを生成してしまう）ことが分かりました。

正しいSQLが生成できないときのカスタマイズとして、

  1. テーブル・カラムの説明を追加する
  2. サンプルのクエリを追加する

などの機能があり、これらを使って出力を調整できることが分かりました。

この機能を活用して、従来のRAGでは回答できなかったユーザーの質問に適切に回答できるシステムを構築したいです。

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