Amazon AuroraとAmazon Bedrockを統合してRAG化検証

AuroraがBedrockに対応したブログを見て衝撃を受けました。

https://dev.classmethod.jp/articles/amazon-aurora-postgresql-bedrock-ml2/

```sql
SELECT aws_bedrock.invoke_model (
    'amazon.titan-text-express-v1',
    'application/json',
    'application/json',
    '{"inputText": "what is orange"}'
    );
```

**何このSQL www**

と思ったのですが、どういうふうに活用するのかいまいちイメージがつかなかったので、実際に検証してみました。


## 準備

このような環境を構築します。詳細な構築方法はこちらを参照ください。

https://moritalous.pages.dev/75793555a5e1fe2657d0

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/a382fdaf-d418-9a78-9410-6100edd54cb3.png)

AuroraからBedrockを使用するために必要なのはこのあたりの設定です。

1. AuroraからBedrockにアクセスするIAMロールを作成する。
    手順を参考に実施してください。

    https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/postgresql-ml.html#postgresql-ml-setting-up

1. Auroraに機械学習拡張をインストールする。
    https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/postgresql-ml.html#postgresql-ml-aws_ml-install

    ```sql
    CREATE EXTENSION IF NOT EXISTS aws_ml CASCADE;
    ```

## aws_bedrock.invoke_model

```sql
SELECT aws_bedrock.invoke_model (
    'amazon.titan-text-express-v1',
    'application/json',
    'application/json',
    '{"inputText": "what is orange"}'
    );
```

```text:結果
{
  "inputTextTokenCount": 3,
  "results": [
    {
      "tokenCount": 8,
      "outputText": "\nAn orange is a citrus fruit.",
      "completionReason": "FINISH"
    }
  ]
}
```

`outputText`の部分だけほしい場合は、

```sql
SELECT aws_bedrock.invoke_model (
    'amazon.titan-text-express-v1',
    'application/json',
    'application/json',
    '{"inputText": "what is orange"}'
  )::json->'results'->0->>'outputText';
```

```text:結果
\nAn orange is a citrus fruit.
```


* `::json`

    JSON型に変換

    https://www.postgresql.jp/document/13/html/datatype-json.html


* `-> 'results'`、`->0`

    ->の後ろがtextの場合はJSONのキーを指定してバリューを取得、intの場合は指定した番号の要素を取得

    https://www.postgresql.jp/docs/9.4/functions-json.html

* ->>'outputText'

    指定した要素をtextとして取得

    https://www.postgresql.jp/docs/9.4/functions-json.html

## invoke_model_get_embeddings

埋め込みモデルの場合は、`invoke_model_get_embeddings`を使用します。

```sql
SELECT aws_bedrock.invoke_model_get_embeddings (
    'amazon.titan-embed-text-v1',
    'application/json',
    'embedding',
    '{"inputText": "what is orange"}'
  );
```

```text:結果
[0.8984375,0.10058594,0.37695312,0.0055236816,-1.1640625,...]
```

3つ目のパラメーターの`embedding`は、`invoke_model`で実行した際のレスポンスで、ベクトル配列が格納されているキーです。

```
{
  "embedding": [0.8984375,0.10058594,0.37695312, ...],        <--このキー
  "inputTextTokenCount": 3
}
```

なので、Cohere Embedの場合は`embeddings`になると思います。


## ベクトルデータベースを作成

ベクトル化した値をデータベースに格納します。

テーブルの形式はKnowledge Base for Amazon Bedrockのスキーマと合わせてみました。

https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/AuroraPostgreSQL.VectorDB.html

:::note
Knowledge Base for Amazon BedrockではData APIが必須ですが、今回の検証では不要です。
:::

上記サイトに従いテーブルを作成します。

```sql
CREATE EXTENSION IF NOT EXISTS vector;

CREATE SCHEMA bedrock_integration;

CREATE TABLE bedrock_integration.bedrock_kb (
  id uuid PRIMARY KEY,
  embedding vector(1536),
  chunks text,
  metadata json
);

CREATE INDEX on bedrock_integration.bedrock_kb USING hnsw (embedding vector_cosine_ops);
```

ベクトル化したデータを登録します。

```sql
INSERT INTO bedrock_integration.bedrock_kb (id, embedding, chunks, metadata)
VALUES (
    gen_random_uuid(),
    CAST(
      aws_bedrock.invoke_model_get_embeddings (
        'amazon.titan-embed-text-v1',
        'application/json',
        'embedding',
        '{"inputText": "what is orange"}'
      ) as vector
    ),
    'what is orange',
    '{"outputText": "An orange is a citrus fruit."}'
  );
```

:::note
metadataは好きなJSONが登録できます。質問に対する回答を入れてみました。
:::

登録できたので検索してみましょう。

```sql
SELECT *
FROM bedrock_integration.bedrock_kb
ORDER BY embedding <->CAST(
    aws_bedrock.invoke_model_get_embeddings (
      'amazon.titan-embed-text-v1',
      'application/json',
      'embedding',
      '{"inputText": "what is orange"}'
    ) AS vector
  )
LIMIT 5;
```

`<->`でベクトル同士の近さが取得できます。他に`<#>`や`<=>`もサポートされているようです。近い順に並び替えて上位5件を取得しています。

https://github.com/pgvector/pgvector

## Aurora単体でRAG

ここまで来たらRAG化したくなりました。SQLが長くなってきました。

サンプルとして、以下2件のデータを登録しました。

| Question | Answer | 
| --- | --- |
| Q: What is Amazon Elastic Compute Cloud (Amazon EC2)? | Amazon EC2 is a web service that provides resizable compute capacity in the cloud. It is designed to make web-scale computing easier for developers. |
| What is Amazon S3? | Amazon S3 is object storage built to store and retrieve any amount of data from anywhere. S3 is a simple storage service that offers industry leading durability, availability, performance, security, and virtually unlimited scalability at very low costs. |

<details><summary>INSERT文はこちら</summary>

```sql
INSERT INTO bedrock_integration.bedrock_kb (id, embedding, chunks, metadata)
VALUES (
    gen_random_uuid(),
    CAST(
      aws_bedrock.invoke_model_get_embeddings (
        'amazon.titan-embed-text-v1',
        'application/json',
        'embedding',
        '{"inputText": "Q: What is Amazon Elastic Compute Cloud (Amazon EC2)?"}'
      ) as vector
    ),
    'Q: What is Amazon Elastic Compute Cloud (Amazon EC2)?',
    '{"outputText": "Amazon EC2 is a web service that provides resizable compute capacity in the cloud. It is designed to make web-scale computing easier for developers."}'
  ),
  (
    gen_random_uuid(),
    CAST(
      aws_bedrock.invoke_model_get_embeddings (
        'amazon.titan-embed-text-v1',
        'application/json',
        'embedding',
        '{"inputText": "What is Amazon S3?"}'
      ) as vector
    ),
    'What is Amazon S3?',
    '{"outputText": "Amazon S3 is object storage built to store and retrieve any amount of data from anywhere. S3 is a simple storage service that offers industry leading durability, availability, performance, security, and virtually unlimited scalability at very low costs."}'
  );
```

</details>

プロンプトは[こちら](https://platform.openai.com/docs/guides/prompt-engineering/strategy-provide-reference-text)を参考に以下の文としました。

> Human:
> 
> Use the provided articles to answer questions. If the answer cannot be found in the articles, write  \"I could not find an answer.\"
> 
> <knowledge>${KNOWLEDGE}</knowledge>
> 
> Question: ${QUESTION} 
> 
> Assistant:


`${KNOWLEDGE}`と`${QUESTION}`はSQLの中で置換します。

Claude Instantを使用して、上位2件の結果を元にサマライズします。

```sql
SELECT aws_bedrock.invoke_model (
    'anthropic.claude-instant-v1',
    'application/json',
    'application/json',
    T4.PROMPT
  )::json->'completion' as A,
  Q,
  PROMPT
FROM (
    SELECT REPLACE(
        REPLACE(
          '{"prompt": "Human:\n\nUse the provided articles to answer questions. If the answer cannot be found in the articles, write  \"I could not find an answer.\"\n\n<knowledge>${KNOWLEDGE}</knowledge>\n\nQuestion: ${QUESTION} \n\nAssistant:", "max_tokens_to_sample": 300}',
          '${KNOWLEDGE}',
          T3.outputText
        ),
        '${QUESTION}',
        Q
      ) as PROMPT,
      Q
    FROM (
        SELECT STRING_AGG(T2.outputText, '\n\n') as outputText,
          Q
        FROM (
            SELECT metadata->>'outputText' as outputText,
              Q
            FROM bedrock_integration.bedrock_kb,
              (
                SELECT 'What is EC2?' as Q
              ) T1
            ORDER BY embedding <->CAST(
                aws_bedrock.invoke_model_get_embeddings (
                  'amazon.titan-embed-text-v1',
                  'application/json',
                  'embedding',
                  REPLACE(
                    '{"inputText": "${QUESTION}"}',
                    '${QUESTION}',
                    Q
                  )
                ) AS vector
              )
            LIMIT 2
          ) T2
        GROUP BY Q
      ) T3
  ) T4
```

```text:結果（prompt）
{""prompt"": ""Human:\n\nUse the provided articles to answer questions. If the answer cannot be found in the articles, write  \""I could not find an answer.\""\n\n<knowledge>Amazon EC2 is a web service that provides resizable compute capacity in the cloud. It is designed to make web-scale computing easier for developers.\n\nAmazon S3 is object storage built to store and retrieve any amount of data from anywhere. S3 is a simple storage service that offers industry leading durability, availability, performance, security, and virtually unlimited scalability at very low costs.</knowledge>\n\nQuestion: What is EC2? \n\nAssistant:"", ""max_tokens_to_sample"": 300}
```

```text:結果（q）
What is EC2?
```

```text:結果（a）
 Amazon EC2 is a web service that provides resizable compute capacity in the cloud. It is designed to make web-scale computing easier for developers.
```


うまくいきました！！

## まとめ

Aurora Postgresのベクトルデータベース化とBedrock呼び出しを試しました。単体でのRAG化もできました。ベクトル化や検索をデータベース側にオフロードできますね。