Amazon Bedrock Knowledge Base に API 経由でドキュメントを登録する

背景

Amazon Bedrock Knowledge Bases とは

Amazon Bedrock には Knowledge Bases という機能がある。S3 などのデータソースを登録すると、そこに入っているドキュメントを勝手に読み込んで、それを元に RAG を構築したり出来る。

Foundation Models for RAG – Amazon Bedrock Knowledge Bases – AWS

WP RAG のバックエンド

話は変わって、以前、WP RAG というサービスをリリースした。WordPress のコンテンツを元に RAG を使ったチャットボットを作れるというもの。詳細は以下のページを参照の事。

WP RAG – Build RAG Systems with WordPress Content

これのバックエンドだが、弊社の自前のインフラでベクトルデータベースを持っていて、OpenAI の Text Embeddings API を使って埋め込み表現を作成して、それを保存している。

バックエンドのインフラを弊社で管理しているので、WordPress のサイト管理者に難しい知識がなくても RAG を使ったチャットボットが簡単に作れるというメリットがある。その一方、

• カスタマイズの自由度が低い
• データを外部のサーバー（弊社の API サーバー）に置かなければいけない

というデメリットもある。そこで、今回、WP RAG のバックエンドに Amazon Bedrock を使ったバージョンも作ってみたので、それの説明をする。

ソースコードは以下を参照。

mobalab/wp-rag-for-amazon-bedrock: Integrates WordPress content with Amazon Bedrock Knowledge Bases to create an intelligent RAG chatbot system with automatic content

なお、最初はこのプラグインの名前を「WP RAG for Amazon Bedrock」としていたが、理由があり「RAG Chat for Amazon Bedrock」に変更してプラグインディレクトリに申請した。レポジトリ名も変更予定。

WP から始まる WordPress プラグイン名 – もばらぶエンジニアブログ

やる事

事前準備

管理コンソールから Knowledge Base と Data source を作成しておく。Data source の種類は「Custom」にしておく。

あとは、IAM ユーザーを作成し、適切な権限を付与しておく。

また、RAG で使用する Foundation Model は、使用可能にしておく。（Amazon Bedrock の初期状態では、全てのモデルが無効化されている事に注意。）

API でデータ登録

後は、以下の API を呼び出せば良い。

IngestKnowledgeBaseDocuments – Amazon Bedrock

それだけだとあんまりなので、パラメータ等についていくつか補足しておく。

今回登録するのはテキストデータなので、 inlineContent > type は TEXT を指定し、textContent にテキストデータを入れる。

上述のページ（API ドキュメント）にある「Request Syntax」には optional なパラメータも全て記載されているが、不要なものも多い。例えば、今回は Custom の data source なので、S3 関連のパラメータは渡す必要はない。

API 経由で質問をしてみる

ドキュメントの登録が終わったら、以下の API を呼ぶ事で、RAG の動作を確認できる。

RetrieveAndGenerate – Amazon Bedrock

レスポンスに含まれるのは以下のもの。

• 回答文
• citations （回答を作る際に引用した引用元のドキュメントに関する情報）
• sessionId

次に API を呼び出す際に sessionId を含めると、今回の回答文を記憶した状態で回答をしてくれる。逆に sessionId を渡さないと、新規の質問という扱いになる。

注意点、補足、その他

API エンドポイント

Amazon Bedrock の API には以下のような複数の種類があり、エンドポイントも異なるので注意。

• Amazon Bedrock
• Agents for Amazon Bedrock
• Agents for Amazon Bedrock Runtime
• Data Automation for Amazon Bedrock
• Runtime for Amazon Bedrock Data Automation

エンドポイントの一覧は以下のドキュメントを参照。

Amazon Bedrock endpoints and quotas – AWS General Reference

Amazon Bedrock API keys が使える API は限られている

上の方で、「IAM ユーザーを作成し」と書いたが、Amazon Bedrock には独自の API key がある。ただし、この API key が使える API は Amazon Bedrock と Amazon Bedrock Runtime の actions に限られる。詳しくは以下のドキュメントを参照。

Supported Regions and AWS software development kits (SDKs) for Amazon Bedrock API keys – Amazon Bedrock

ログを有効化しておく

Knowledge Base にドキュメントを入れようとしてもエラーになる事もある。ログを有効化しておくとデバッグがしやすくて良い。詳しくは以下のドキュメントを参照。

Monitor knowledge bases using CloudWatch Logs – Amazon Bedrock

まとめ

Amazon Bedrock には Knowledge Bases という機能があり、これを使うと RAG が簡単に構築できる。データソースとしては S3 を使う事を主に想定しているようだが、Custom タイプのデータソースにすると、API 経由で外部からドキュメントを追加する事が可能。

埋め込み作成、ベクトルデータベース、回答生成、全て Amazon のインフラで完結するのでとても楽だった。独自インフラ＋OpenAI API を使った WP RAG の場合、ある程度動くまでに2ヶ月くらいはかかったけど、Amazon Bedrock を使ったバージョンは2〜3日くらいで基本機能は実装出来た。もちろん、WordPress プラグイン部分は WP RAG の流用だが、その分を考慮したとしてもサーバーサイドのインフラ構築・プログラムが完全に不要なので、労力としては1/3くらいのイメージ。