Step FunctionsがBedrockに対応したのでやってみた

[Amazon Bedrock Advent Calendar 2023](https://qiita.com/advent-calendar/2023/amazon-bedrock-generative-ai-aws) 5日目の投稿です。

---

re:Inventの直前にStep FunctionsがBedrockに対応したとのアナウンスがありました。

https://aws.amazon.com/jp/about-aws/whats-new/2023/11/aws-step-functions-optimized-integration-bedrock/

個人的にとても期待していたアップデートなので、re:Inventの帰りの飛行機を待ちながら検証しました。

## テンプレートからデプロイしてみる

1. マネジメントコンソールにアクセスし`今すぐ始める`をクリックします。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/c31b6306-492e-4f88-e842-42178e21aabd.png)

1. 真ん中の`テンプレートを選択`をクリックします。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/c9fac5e1-993e-2131-51e9-c2714844f134.png)

1. １つ目に`Bedrockでプロンプトチェイニングを実行する`を選択して`次へ`をクリックします。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/823161dc-b6e8-c984-8f59-7cab0d7f608a.png)

1. テンプレートの内容が表示されます。`テンプレートの使用`をクリックします。

    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/cdd69979-4c3f-bb52-9961-e9862e3f6fc2.png)

1. CloudFormationでデプロイするようです。右上の`デプロイと実行`をクリックします。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/4e2f99f8-67c9-fa2f-772a-b091519270e4.png)

これでStep Functionsのステートマシンが作成されました。

## 使ってみる

ステートマシンの作成後に、実行前の状態になりますので、このまま実行します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/aef41666-660c-8bd0-8952-e6a4d27e86a9.png)

インプットのJSONを整形するとこのような形です。3回Bedrockを呼び出すのでそれぞれのプロンプトを指定します。

```json
{
    "prompt_one": "Name a random city. You must provide a non-fictitious city. You must only provide the city's name, followed by a comma and the city's country of origin.",
    "prompt_two": "Write two to three sentences describing the city mentioned above",
    "prompt_three": "Write three to four sentences describing the city mentioned above"
}
```

```text:１つ目のプロンプト（日本語訳）
任意の都市名を挙げてください。架空の都市ではありません。都市名の後にカンマ、そしてその都市の出身国のみを記入すること。
```

```text:２つ目のプロンプト（日本語訳）
上記の都市について、2～3つの文章で説明してください。
```

```text:３つ目のプロンプト（日本語訳）
上記の都市について、3～4つの文章で説明してください。
```

実際に呼び出してみましょう。

* Invoke model with first promptステップ

    １回目のBedrock呼び出しです。

    ```json:入力
    {
      "prompt_one": "Name a random city. You must provide a non-fictitious city. You must only provide the city's name, followed by a comma and the city's country of origin.",
      "prompt_two": "Write two to three sentences describing the city mentioned above",
      "prompt_three": "Write three to four sentences describing the city mentioned above"
    }
    ```

    ```json: 出力
    {
      "prompt_one": "Name a random city. You must provide a non-fictitious city. You must only provide the city's name, followed by a comma and the city's country of origin.",
      "prompt_two": "Write two to three sentences describing the city mentioned above",
      "prompt_three": "Write three to four sentences describing the city mentioned above",
      "result_one": {
        "result_one": " Well, the following are some random cities from around the world:\n\n1. Kyoto, Japan"
      }
    }
    ```

    Bedrockで生成した文書が`result_one`に出力されました。

* Add first result to conversation historyステップ

    チャット履歴を保持します。

    ```json:入力
    {
      "prompt_one": "Name a random city. You must provide a non-fictitious city. You must only provide the city's name, followed by a comma and the city's country of origin.",
      "prompt_two": "Write two to three sentences describing the city mentioned above",
      "prompt_three": "Write three to four sentences describing the city mentioned above",
      "result_one": {
        "result_one": " Well, the following are some random cities from around the world:\n\n1. Kyoto, Japan"
      }
    }
    ```
    
    ```json:出力
    {
      "prompt_one": "Name a random city. You must provide a non-fictitious city. You must only provide the city's name, followed by a comma and the city's country of origin.",
      "prompt_two": "Write two to three sentences describing the city mentioned above",
      "prompt_three": "Write three to four sentences describing the city mentioned above",
      "result_one": {
        "result_one": " Well, the following are some random cities from around the world:\n\n1. Kyoto, Japan"
      },
      "convo_one": {
        "convo_one": "Name a random city. You must provide a non-fictitious city. You must only provide the city's name, followed by a comma and the city's country of origin.\n Well, the following are some random cities from around the world:\n\n1. Kyoto, Japan"
      }
    }
    ```

    チャット履歴が`convo_one`に追加されています。内容は１度目のBedrock呼び出しの入力プロンプトと出力を連結して保持しています。

同様に２回目、３回目とBedrockを呼び出します。

３回目の呼び出し結果はこのようになります。

```json
{
  "result_three": " Kyoto is a city on the most populous island of Japan, Honshu. It is a place of cultural importance in Japan, having been the country's capital for more than one thousand years. Kyoto is home to hundreds of temples, gardens, and Shinto shrines, most of which are preserved thanks to the city's late 19th-century survival of American bombing campaigns during World War II.  This city is a favorite destination for those seeking historical and cultural experiences, especially in the spring when the cherry blossoms bloom.  The cherry blossom season, beginning in March, is the most popular time to visit Kyoto and attracts visitors worldwide.  The city also has a vibrant cultural scene and offers traditional Japanese arts, performances, and festivals.  In addition, Kyoto is known for its savory cuisine, including its fresh, local dishes and its vegetarian temple food.  With so much to offer, it's no wonder this beautiful city is a popular tourist destination. \n\nWould you like to know more about"
}
```

>  京都は日本で最も人口の多い島、本州にある都市である。千年以上もの間、日本の首都であった京都は、日本文化にとって重要な場所である。京都には何百もの寺院、庭園、神社があり、そのほとんどは19世紀後半に第二次世界大戦中のアメリカの爆撃を免れたおかげで保存されている。 この街は、歴史的、文化的な体験を求める人々にとって、特に桜が咲く春の旅行先として人気がある。 3月から始まる桜の季節は、京都を訪れるのに最も人気のある時期で、世界中から観光客が訪れる。 また、京都には活気ある文化シーンがあり、日本の伝統芸能やパフォーマンス、お祭りを楽しむことができる。 さらに、京都は新鮮な郷土料理や精進料理の寺院料理など、味わい深い料理でも知られている。 これだけの魅力があれば、この美しい街が人気の観光地であるのもうなずける。についてもっと知りたいですか？

## 仕組みにDive Deep

一体どのようにBedrockを呼び出しているのでしょうか？

１つ目のBedrock呼び出しの`設定`タブはこの様になっていました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/4cf31a81-b0f9-df83-5761-6ef2be74f2e3.png)

設定タブでは主に以下の設定を行います。

* API パラメータ

    通常かプロビジョンドかを選択します

* Bedrock model identifier

    使用する基盤モデルのModel IDを指定します。直接入力だけでなく選択肢も出してくれます。実行時に前ステップの変数を指定することもできます。

    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/64a1c59c-1d33-d2e0-7b20-ec4a8e597d20.png)

    ちなみに、画像生成のStable Diffusionも使えそうです。
    
* Bedrock モデルパラメータ

    モデル呼び出し時のパラメーターを指定します。サンプルテンプレートでは以下のないようです。

    ```json
    {
        "prompt.$": "$.prompt_one",
        "max_tokens": 20
    }
    ```

    また、プリセットから選ぶこともできます。

    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/80a0fae0-80e6-d46c-0d1c-a962eaf24d51.png)

    Claudeを選んだ場合はこのようにひな形がセットされるので、あとはプロンプト部分を修正するだけです。

    ```json
    {
      "prompt": "\n\nHuman:<prompt>\n\nAssistant:",
      "max_tokens_to_sample": 200
    }
    ```

`出力`タブも確認しましょう。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/5cc8ee36-d7a3-aea3-543e-b84615db1097.png)

* ResultSelector を使用して結果を変換

    Bedrockのinvoke model APIのレスポンスから、生成した文言だけを抽出しています。Bedrockを使ったことがあれば違和感はないと思います。（PythonのBoto3の場合はStreamingBodyからreadしたりjsonに変換したり面倒なコードが必要でしたが、Step Functionsでは不要でいいですね！）
    
    ```json
    {
      "result_one.$": "$.Body.generations[0].text"
    }
    ```

* ResultPath を使用して元の入力を出力に追加

今回のデモは3つのプロんプロトを入力パラメーターで指定しているので、後続のステップに入力パラメーターを引き継ぐ必要があります。そのため、ここで設定し、Bedrockで生成した文書は`result_one`にセットしています。

## まとめ

Bedrock用に最適化されてインテグレーションされていることが確認できました。生成系AIの呼び出しは複数ステップになることが多いので、Bedrockを使う際はStep Functionsの利用も検討してみてはいかがでしょうか？