💡

TiDB (Serverless) × Amazon Bedrockで始めるRAGアプリケーション入門！

2024/05/22に公開

AWS

はじめに

PingCAPの小板橋です。はじめまして！
TiDBの入門記事から上級者編まで幅広く取り扱う本アカウント第8回目は「TiDB (Serverless) × Amazon Bedrockで始めるRAGアプリケーション入門！」です！

クレジット登録無しで始められるTiDB Serverlessを利用し、NewSQLとAmazon Bedrockを組み合わせてRAGのアプリケーションを構築してみましょう！

TiDB × Amazon Bedrockで始めるRAGアプリケーション

最終成果物は、下記のRepositoryにて公開しています。

実際にローカルで起動する画面

TiDBにおけるベクトル検索

TiDB (Serverless)に今年(2024/2/1)、ベクトル検索の機能が追加されました。
この機能というのは、MySQL互換のNewSQLであるTiDBにおいて、現時点ではTiDB Serverlrssで作成したクラスターの中で、ベクトル列を追加でき、ベクトル検索することができるといった機能になります。
また、このベクトルカラム、HNSWインデックスとして作成することができるのが特徴になっています。

よくあるベクトル検索のユースケースの中でも今回は、生成AIに絡め、RAG(Retrieval-Augmented Generation）アプリケーションをBedrockより提供されているモデルを組み合わせ作成するものを今回は記事として取り扱いたいと思います。

詳細は、以下の記事をご参照ください。

それではやっていきましょう！！

1: TiDB Serverlessの準備

まずは、TiDBコンソールへサインアップ/サインインしていきましょう。

1-1: TiDBコンソールへのサインアップ/サインイン

https://pingcap.co.jp/

右上の、「今すぐ始める」をクリックしてください。

そうすると、下記の画面に遷移します。
今回利用するTiDBクラスターは、クレジット登録無しでも利用できるTiDB Serverlessクラスターを作成していくので一番左のプランを選択します。

選択すると、下記の画面が表示されます。
ここで、既にユーザがある方はそのアカウント情報でサインインをします。
まだユーザ作成ができていない方は、名前、メールアドレス、パスワードを入力し、サインアップしていきます。

1-2: TiDB Serverlessクラスター(ベクトル検索機能をオンにした)を作成する

サインイン/サインアップすると下記の画面に遷移します。
私のコンソール画面ではいくつかのTiDBクラスターがありますが、ここは人によってクラスターの状況は違うので、流し見していただければと思います。
この画面では、右上にある「Create Cluster」をクリックし、TiDB Serverlessクラスターを作成していきます。

「Create Cluster」をクリックすると下記の画面に遷移します。
ここで入力する内容としては、「Choose a Tier」を「Serverless」を選択し、「Basic Setting」の「Region」を「Frankfurt(en-central-1)」を選択します。
現在、ベクトル検索の機能は、特定のリージョンのみで利用できるプライベートβ扱いのためです。(2024/5/22時点)

なお、ここで、「Region」を「Frankfurt(en-central-1)」を選択し、ベクトル検索の機能が表示されない場合は、下記のWaitlistへ登録いただけると順次ベクトル検索機能が使えるTiDB Serverlessクラスターの作成ができるようになります。

https://tidb.cloud/ai

ベクトル検索機能を有効化するトグル画面

そのほか、「Cluster Name」などを入力していきます。
なお、クレジットカードの登録をしているかしていないかで、下記の「Cluster Plan」の表示が異なります。
私のアカウントでは既にクレジットカードの登録がされており、Freeプランが選択できないため、「Scalable Cluster」を選択していきます。

また、「Monthly Spending Limit」の設定を念の為、$1設定し、TiDB Serverlessクラスターを作成していきます。

1-3: TiDB ServerlessクラスターへローカルPCから接続する

下記の状態「StatusがAvailable」になるとTiDB Serverlessへの接続が可能になります。
ここでは、ローカルのターミナルからMySQL CLIを利用し、実際にTiDB Serverlessへ接続できるか検証していきましょう。

右上にある「Connect」をクリックすると下記の画面が表示されます。
ここで、「Connect With」を「MySQL CLI」を選択すると、MySQl CLIを使用し、作成したTiDB Serverlessクラスターへの接続するための情報が表示されます。
コピペする前に「Generate password」をクリックするとランダムに生成されたパスワードが貼り付けられた状態で、MySQL CLIから接続するための情報がコピペできるようになるので、そちらを利用します。

実際にローカルのターミナルからMySQL CLIを利用し接続すると下記のように接続することが確認できます。

一旦ここまでできていれば大丈夫です。コピペしたTiDB Serverlessへの接続するための情報は忘れずにどこかに書き留めておきましょう。ユーザ名、パスワード、ホスト名、データベース名はこの後利用します。

2: Amazon Bedrockの準備

次にRAGのアプリケーションで利用していくBedrockで提供されているモデルを有効化していきます。
AWSコンソールへサインインしていただき、Amazon Bedrockのサービス画面へ遷移します。

今回は、AWSのリージョンをus-east-1を使用していますが、好きなモデルが提供されているリージョンであればどこでも大丈夫です。
その後、Bedrockの画面より「モデルアクセス」を選択し、利用したいモデルの有効化をしていきます。

こんな感じで、各モデルを有効化すると「アクセスが付与されました」と表示されます。
私は今回、テキストモデルとして「Claude 3 Sonnet」とembeddingモデルとして「Titan Embeddings G1 - Text」を利用しています。

その後、ローカルPCに設定されている「~/.aws/credentials」の情報をアクセスができるように設定指定おきます。この辺りはたくさん記事が世の中公開されているので、他を参照してみてください。

3: ローカルへソースコードをClone

次に今回利用するRAGのサンプルコードは下記になります。READMEにも詳細は書いているのですが、本記事中でも使い方については記述していきます。

git clone https://github.com/Yoshiitaka/TiDB-Bedrock.git

VSCodeなどのIDEでcloneしたソースコードを表示してください。

4: 事前にTiDBのベクトルカラムにパブリックに公開されているTiDBのブログをEmbeddingしたベクトルデータを流し込む

さて、ここからはコードの世界になります。全体と作りとしては、TiDB Serverlssクラスターのベクトルデータに事前にテキストデータをEmbeddingしベクトルデータを流し込むスクリプトと、streamlitを利用しUI画面から実際にチャットをし、回答を確認するコードの２つが存在します。

まず行うのは、前段の事前にTiDBのベクトルカラムにパブリックに公開されているTiDBのブログをEmbeddingしたベクトルデータを流し込む作業になります。

4-1: venvの環境を用意し、Active化します

今回、コードはPythonで書いているのですが、ライブラリやモジュール周りの依存関係をローカル環境の問題から切り離したいので、venvを利用します。

python -m venv .venv
source .venv/bin/activate

4-2: 依存関係のあるライブラリ、パッケージをインストール

その後、依存関係のあるライブラリ、パッケージをインストールしていきます。

pip install -r requirements.txt

4-3: 環境変数にTiDBへの接続情報を記載

最後に、環境変数にTiDBへ先ほどコピペした情報を元に接続情報を記載していきます。
コードの中に、ファイルとして、サンプルのファイル「sample-env」を用意しています。
このファイルを参考に「.env」ファイルを作成します。下記の接続情報をベクトル検索を有効化したTiDB Serverlessクラスターの接続情報で更新します。

TIDB_HOSTNAME='xxxxxxx'
TIDB_USERNAME='xxxxxx.root'
TIDB_PASSWORD='xxxxxx'
TIDB_DATABASE_NAME='test'

4-4: TiDBのベクトルカラムにパブリックに公開されているTiDBのブログをEmbeddingしたベクトルデータを流し込みます。

下記のコマンドを実行すると、47行目〜54行目に記載されているインターネット公開されているTiDBに関連する私のブログを読み取り、データをベクトル化し、TiDBへ格納する処理を行ってくれます。

python prepare.py

47行目〜54行のコード

    documents = SimpleWebPageReader(html_to_text=True).load_data(
        [
            "https://docs.pingcap.com/ja/tidb/stable/overview",
            "https://docs.pingcap.com/ja/tidb/stable/release-7.5.0",
            "https://docs.pingcap.com/ja/tidb/stable/mysql-compatibility",
            "https://docs.pingcap.com/ja/tidbcloud/tidb-cloud-intro#architecture"
        ]
    )

5: ローカルでUI画面を起動し、実際のRAGを体験する

ここまでできれば、あとはローカルでUI画面を起動し、実際のRAGを体験していきましょう！
streamlitを利用し書かれたフロントサーバをローカルで立ち上げ、実際にチャットを体験してみてください。

streamlit run main.py

立ち上がるローカルのポート番号は8501です。なので、下記へブラウザよりアクセスしてみてください。

http://localhost:8501/

下記の画面のように4の手順で読み込ませたインターネット公開されている私のTiDBに関連するブログをベクトル化し、TiDBに格納し利用したRAGを構成しているので、TiDBに関連する情報を回答してくれるチャットボットが完成しました。

あとは自由に改造して楽しんでください！

まとめ

いかがだったでしょうか？
TiDBの世界は奥が深いです。引き続き様々な機能についてを深掘りブログ化していきたいと思います。