Knowledge Base for Amazon Bedrock で EventBridge と Step Functions を使ったスケジュール Sync 実行をやってみた

はじめに

Knowledge Base for Amazon Bedrock は、S3 に格納したデータを使って簡単に RAG のデータソースにする機能があります。S3 に格納したうえで Sync を行うことで、初めて RAG として活用できるようになります。そのため、定期的に、もしくは準リアルタイムに Sync を行いたいたです。

この記事では、EventBridge と StepFunctions を使って、日次で Sync を行う方法を紹介します。なお、Sync のタイミングについては、ユースケースに応じてタイミングは自由にコントロールが可能です。

Sync を行う API

Knowledge Base でデータの Sync を行う API は、StartIngestionJob です。

ちょっと名称が直感的ではないので迷いそうになりますが、これを使って定期的に Sync を行う方法を紹介します。

Step Functions のステートマシンを作成

StartIngestionJob を行うための Step Functions ステートマシンを作成します。今回は、Knowledge Base はバージニア北部リージョンで構成しているため、Step Functions も同じリージョンのバージニア北部を利用します。

Create state machine を選択します。

テンプレートから選択できるようになっていますが、Blank のまま Select をします。

今回は、Code で編集を行います。

Code に以下の JSON 文字列を全部上書き貼り付けをします。

{
  "Comment": "A description of my state machine",
  "StartAt": "StartIngestionJob",
  "States": {
    "StartIngestionJob": {
      "Type": "Task",
      "Parameters": {
        "DataSourceId.$": "$.DataSourceId",
        "KnowledgeBaseId.$": "$.KnowledgeBaseId"
      },
      "Resource": "arn:aws:states:::aws-sdk:bedrockagent:startIngestionJob",
      "Next": "Wait"
    },
    "Wait": {
      "Type": "Wait",
      "Seconds": 5,
      "Next": "GetIngestionJob"
    },
    "GetIngestionJob": {
      "Type": "Task",
      "Parameters": {
        "DataSourceId.$": "$.IngestionJob.DataSourceId",
        "IngestionJobId.$": "$.IngestionJob.IngestionJobId",
        "KnowledgeBaseId.$": "$.IngestionJob.KnowledgeBaseId"
      },
      "Resource": "arn:aws:states:::aws-sdk:bedrockagent:getIngestionJob",
      "Next": "Choice"
    },
    "Choice": {
      "Type": "Choice",
      "Choices": [
        {
          "Variable": "$.IngestionJob.Status",
          "StringEquals": "COMPLETE",
          "Next": "Success"
        },
        {
          "Or": [
            {
              "Variable": "$.IngestionJob.Status",
              "StringEquals": "STARTING"
            },
            {
              "Variable": "$.IngestionJob.Status",
              "StringEquals": "IN_PROGRESS"
            }
          ],
          "Next": "Wait"
        }
      ],
      "Default": "Fail"
    },
    "Fail": {
      "Type": "Fail"
    },
    "Success": {
      "Type": "Succeed"
    }
  }
}

Code を張り付けた結果、以下のフローが出来上がります。ポイントは以下の通りです。

• StartIngestionJob : データの Sync を実行
• Wait : 5 秒待機
• GetIngestionJob : データを Sync した Job の Status を取得
• Choice : STARTING や IN_PROGRESS のステータスの場合は、WAIT に戻って再度 5 秒待機する。COMPLETE の場合は、正常終了。それ以外のステータスは FAIL で異常終了

Config を開き、State machine の名前を適当にいれて、Create を押します。

Bedrock に関する IAM Role は自動生成できないため、手動で設定するように促されます。Confirm を押します。

IAM Role の詳細画面を開きます。

Attach Policy を押します。

今回は手順をシンプルにするために、AmazonBedrockFullAccess を選択します。

作成した State machine をテスト実行します。

以下の JSON を、各環境に合わせて ID を変更したうえで、張り付けます。

• DataSourceId : Knowledge Base の DataSource ID を指定
• KnowledgeBaseId : Knowledge Base の ID を指定

{
  "DataSourceId": "D8UUCPWZSA",
  "KnowledgeBaseId": "CVXUPL53HG"
}

Success することを確認します。

Knowledge Base の Data source の画面を開いて、Status が Completed であることも確認します。

EventBridge Scheduler を作成

上記で作成した State machine を日次実行するために、EventBridge Scheduler を作成します。

Schedule の名前や、cron 式で日次実行を指定します。各環境に合わせて好きに設定してください。

実行する State Machine を選択します。

JSON の Input は、環境に合わせて ID を書き換えてください。

{
  "DataSourceId": "D8UUCPWZSA",
  "KnowledgeBaseId": "CVXUPL53HG"
}

このあたりはデフォルトのままで OK です。リトライに関する設定もあるので、内容を確認しながら指定しても良いと思います。

Create

動作確認

指定した Schedule を手動で実行することは出来なさそうだったので、時間になったあと再度確認します。

時間経過後、State Machine がきちんと実行されています。EventBridge 側に指定した Input を正しく受け取れています。

Knowledge Base 側もきちんと実行されています。

参考 URL