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` です。

https://docs.aws.amazon.com/ja_jp/bedrock/latest/APIReference/API_agent_StartIngestionJob.html



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



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

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

Create state machine を選択します。

![image-20240504120037343.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/892777c3-7613-0b5d-5347-3a63059aeda5.png)



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

![image-20240504120205837.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/15716e26-1dca-1690-5b0c-4115ab21810e.png)



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

![image-20240504120257392.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/33c4246c-d6a9-59f2-cb68-93d9afe40737.png)



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 で異常終了

![image-20240504123225908.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/6a4e6b91-8eab-4117-f9f3-677ee3b16629.png)



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

![image-20240504123516616.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/c08fac0f-e8b3-c9aa-5148-1ab348eb1aeb.png)



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

![image-20240504123623331.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/a7ad3434-2aa4-1e33-11df-29429de7c11b.png)



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

![image-20240504123655737.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/37f6ef39-4909-dbbe-0750-218e3b8284ad.png)



Attach Policy を押します。

![image-20240504123736222.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/13f8c5a1-e2d5-a85c-bafd-65ef268decbc.png)



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

![image-20240504123820534.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/ac064f03-7226-8a73-5526-996371e14d3f.png)



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

![image-20240504123927181.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/c2512392-9b5c-ac5e-de8e-c4f793dbe0e9.png)



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

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

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

![image-20240504124048248.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/2ce5c6b3-ac55-7dd7-54c7-798c40a0c356.png)



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

![image-20240504124123009.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/081750a6-a319-550e-087e-77aa94bd2f50.png)



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

![image-20240504124204946.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/76090a17-d8dd-41e9-0cf4-62de9fa2ad41.png)



# EventBridge Scheduler を作成

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

![image-20240504125121029.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/be07fc6d-f5d4-14ca-9336-4a72533f069b.png)



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

![image-20240504125357180.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/713a400f-732f-913f-3dc4-352db029aac0.png)



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

![image-20240504125717297.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/b6a42ba7-90ca-1fa4-0abc-c050d68347f8.png)



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

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



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

![image-20240504125946591.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/31105cfd-3d34-5c2f-8ff7-4b84bb158ff7.png)



Create

![image-20240504130026024.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/dd86b870-a93f-e92e-6069-b2ad7a529ad9.png)



# 動作確認

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

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

![image-20240504130930645.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/361315ca-15b9-c5e9-44c5-07416068f5c6.png)



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

![image-20240504131005016.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1002774/6ee9174b-0843-b28b-6d8e-63948ef47a3d.png)





# 参考 URL

https://xp-cloud.jp/blog/2024/01/12/110000