Amazon Bedrock Knowledge Bases の構造化データ取得機能を掘り下げる

この記事は [**Amazon Bedrock Advent Calendar 2024**](https://qiita.com/advent-calendar/2024/amazon-bedrock-generative-ai) の 10 日目の記事です。

:::note info
2024/12/10 時点で Preview 提供されている機能について記載しています。将来的に動作や仕様が変更される可能性があります。
:::

## ナレッジベースの構造化データ取得機能
AWS re:Invent 2024 で発表されたナレッジベースの新機能です。本記事では re:Invent の[ブレイクアウトセッション](https://www.youtube.com/watch?v=SdKfdkec7SQ) の内容を元に、Text2SQL の精度を向上させるためにナレッジベースがどんな仕組みで動いているのか、ユーザー側でどのようなカスタイマイズオプションがあるのかを掘り下げます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/1b116186-4639-2378-18d6-a2b6d40e227a.png)

本機能では自然言語から SQL を生成し、対応するクエリエンジン経由でデータを取得、LLM が結果をまとめて応答します。ざっくりいうとデータベースに格納されている構造化データに対する RAG をマネージドで実現する機能です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/ca2e083c-ff56-3a55-9515-cfad69ac9b37.png)

Preview 時点でクエリエンジンは Amazon Redshift のみサポートしています。Redshift 上のデータの他、Glue Data Catalog との統合にも対応しており、SageMaker Lakehouse 経由で Redsfhit や S3 Tables のデータにアクセスすることも可能になっているようです。

## 何が嬉しいのか
ここでは大きく二つのユースケースで考えてみます。

### EC サイトのサポートチャットボット
ユーザーが「L サイズの T シャツの在庫はありますか？」「商品の配送ステータスは」どうなっていますか？というような問い合わせをするケースを考えてみます。ドキュメントに対するキーワード検索やセマンティック検索ではこういったリアルタイム性が高いデータを取得することは難しく、システム上で管理されているデータベースから直接データを参照したが方がよいでしょう。

### 社内情報を検索したいビジネスユーザー
SQL や BI ツールを使った分析はある程度前提知識が必要で、一般的にはハードルは高いと言えます。例えば人事担当者などのビジネスユーザーが「部門ごとの有給休暇の平均取得率を教えて」「AWS の資格保有者数を部門ごとに教えて」などいった形で自然言語により社内データからインサイトを得ることができるのは大きなメリットです。

## LLM における SQL 生成の課題
LLM は **構文的** に正しい SQL を作成することは得意ですが、実際のデータベース環境は以下に例をあげるように千差万別です。

* データベースやテーブルのスキーマ情報
    * 同じようなデータを格納していても環境によって異なることは想像に難くありません
* 値の登録のされ方
    * 国名 (Japan) として登録されているのか、カントリーコード (JP) で登録されているのかなど
* テーブル名や列名に社内でしか使われない独自用語が登録されている

このような環境固有のメタデータがないと、実環境で実行可能な SQL を生成することはできません。

```sql
-- 日本の顧客の総売上は？

-- LLM のみで生成した SQL
SELECT
    SUM(sales) AS total_sales
FROM
    customers 
WHERE
    country = 'Japan'; 

-- 環境のメタデータを踏まえて生成した SQL
SELECT
    SUM(sales) AS total_sales
FROM
    tbl_cust_sales
WHERE
    country = 'JP'; 
```

これまで同様の仕組みをユーザー側で構築にするには、別途データベースのメタデータを格納したベクターストアを用意し、LLM が SQL 生成する際のコンテキスト情報として渡す、など工夫をする必要がありました。ナレッジベースの構造化データ取得機能はこういった追加開発が不要で精度の高い SQL が生成される仕組みになっています。

## ナレッジベースにおける SQL の生成精度向上アプローチ
### 環境情報の取得 (自動)
構造化データストアをデータソースとして追加し、同期を行うとナレッジベース川に以下のような情報が取り込まれ、SQL 生成時のメタデータとして扱われます。

* データベース、テーブルのスキーマ情報
* SQL クエリの実行履歴

SQL の実行履歴を参照することで、その環境内におけるデータの扱われ方を理解し、精度向上に役立てているというのが個人的には目から鱗でした。前述の例でいうと、データが国名で登録されているのか、カントリーコードで登録されているのか、クエリの実行履歴を活用すれば実データを取得せずとも理解できるわけです。


### Query configurations (オプション)
生成される SQL の精度向上のため、ユーザー側で追加のコンテキスト情報を提供することもできます。

* Descriptions
* Inclusions/exclusions
* Curated queries

#### Descriptions
特定のテーブルまたは列名に対し、説明を付与できます。例えば以下の例では `orders` テーブルに ts1, ts2, ts3 という列名があり、それぞれが注文作成日時、発送日時、配送日時であるという説明を追加しています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/93329a77-a4e3-b57c-5404-8267624d0c6f.png)
[Unlocking power of structured data with Amazon Bedrock Knowledge Bases (AIM396)](https://www.youtube.com/watch?v=SdKfdkec7SQ) より引用

#### Inclusions/exclusions
SQL 生成の対象または除外にしたいテーブル名や列名を明示することで精度向上を図ることができます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/52872558-24f1-45b8-5ec4-6cee64316a30.png)

#### Curated queries
事前に想定される質問と期待される SQL の情報をキュレーションされたクエリとして登録しておくことができます。例えば社内用語など特殊な内容に関する質問と SQL のセットを登録しておくことで、精度向上に役立てることができます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/88f2784f-aa33-6a24-2739-db03d66dbd31.png)
[Unlocking power of structured data with Amazon Bedrock Knowledge Bases (AIM396)](https://www.youtube.com/watch?v=SdKfdkec7SQ) より引用

## GenerateQuery API
これまでの Retrieve API と RetrieveAndGenerate API に加え、構造化データストアでのみ利用可能な GenerateQuery API が追加されました。その名の通り、ユーザーの入力から生成された SQL のみを取得できる API です。

https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_GenerateQuery.html

Bedrock コンソール上ではナレッジベースのテストを行う際に、Generate SQL queries を選択します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/290811e3-6e98-d93e-eb61-5e398c4a567f.png)

実行結果の例

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/574f3699-fb80-d7fb-0f08-5c123193bfc3.png)

補足: 以下のサンプルデータを使用しています。

```sql
-- 社員マスタテーブルの DDL
CREATE TABLE employee_master (
    employee_id INTEGER PRIMARY KEY,
    first_name VARCHAR(50),
    last_name VARCHAR(50),
    email VARCHAR(100),
    hire_date DATE,
    department VARCHAR(50),
    job_title VARCHAR(100)
);

-- サンプルデータの挿入
INSERT INTO employee_master (employee_id, first_name, last_name, email, hire_date, department, job_title)
VALUES
    (1, 'Taro', 'Yamada', 'taro.yamada@example.com', '2020-04-01', 'IT', 'Software Engineer'),
    (2, 'Hanako', 'Sato', 'hanako.sato@example.com', '2019-09-15', 'HR', 'HR Manager'),
    (3, 'Ichiro', 'Tanaka', 'ichiro.tanaka@example.com', '2021-01-10', 'Sales', 'Sales Representative'),
    (4, 'Yuki', 'Nakamura', 'yuki.nakamura@example.com', '2018-07-22', 'IT', 'System Administrator'),
    (5, 'Akiko', 'Suzuki', 'akiko.suzuki@example.com', '2022-03-01', 'Finance', 'Financial Analyst');

-- 格保有情報テーブルの DDL：
CREATE TABLE certification_info (
    certification_id INTEGER PRIMARY KEY,
    employee_id INTEGER,
    certification_name VARCHAR(100),
    certification_date DATE,
    expiration_date DATE,
    FOREIGN KEY (employee_id) REFERENCES employee_master(employee_id)
);

-- サンプルデータの挿入
INSERT INTO certification_info (certification_id, employee_id, certification_name, certification_date, expiration_date)
VALUES
    (1, 1, 'AWS Certified Solutions Architect - Professional', '2021-06-15', '2024-06-15'),
    (2, 1, 'ITIL Foundation', '2020-11-30', '2023-11-30'),
    (3, 2, 'SHRM-CP', '2020-03-10', '2023-03-10'),
    (4, 3, 'Salesforce Certified Administrator', '2022-01-20', '2024-01-20'),
    (5, 4, 'Cisco Certified Network Associate', '2019-09-05', '2022-09-05'),
    (6, 4, 'CompTIA Security+', '2021-04-18', '2024-04-18'),
    (7, 5, 'CFA Level 1', '2022-05-01', '2025-05-01');
```

以下は Boto3 での実行例です。GenerateQuery API では特にモデル ID の指定は不要で SQL が生成される仕組みになっているようです。

```python
import boto3

def main():
    client = boto3.client("bedrock-agent-runtime")
    
    response = client.generate_query(
        queryGenerationInput={
            "text": "AWS Certified Solutions Architect - Professionalの取得者の氏名を教えて",
            "type": "TEXT"
        },
        transformationConfiguration={
            "mode": "TEXT_TO_SQL",
            "textToSqlConfiguration": {
                "knowledgeBaseConfiguration": {
                    "knowledgeBaseArn": "arn:aws:bedrock:<region>:<account-id>:knowledge-base/<knowledge-base-id>"
                },
                "type": "KNOWLEDGE_BASE"
            }
        }
    )
    
    print(response["queries"][0]["sql"])
    
if __name__ == "__main__":
    main()

```
実行結果の例

```text
$ python3 kb.py
SELECT e.first_name, e.last_name FROM public.employee_master e INNER JOIN public.certification_info c ON e.employee_id = c.employee_id WHERE c.certification_name = 'AWS Certified Solutions Architect - Professional';
```

## 料金
Preview 時点の情報ですが、1,000 クエリあたり $2 の料金が発生します。ただし、構造化データストアのナレッジベースは別途ベクターストアの用意が不要で、データソースのメタデータはサービス側で管理されます。利用規模にもよりますが、比較的コストはおさえることができるのではないかと思います。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/7a5ce923-6827-b1e6-6698-1af6c5f426de.png)


## 参考
本機能のドキュメント

https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base-build-structured.html

ナレッジベースが提供する機能ではありませんが、マテリアライズドビューの利用など開発者側で SQL 生成の精度を高めるためのコツが以下の AWS ブログで公開されています。

https://aws.amazon.com/jp/blogs/machine-learning/generating-value-from-enterprise-data-best-practices-for-text2sql-and-generative-ai/

以上です。
参考になれば幸いです。