Generative AI: Amazon Bedrock with AWS Step Functions

December 3rd, 2023 369 Words

AWS Step Functions can easily be used to wrap existing AWS services and persist specific use cases. For example, AWS Step Functions can call the Amazon Bedrock API to generate text responses using the AI21 Labs Jurassic-2 models. You can create a Step Function using the AWS Management Console, or the AWS Cloud Development Kit.

AWS Management Console

First, head over to the AWS Step Functions console in your browser and create a new blank Step Function within the UI.

Create new AWS Step Function in AWS Management Console

Next, search for Bedrock on the left side of the available actions. In the list of results, drag InvokeModel into your Step Function’s definition:

Edit your AWS Step Function

To configure Amazon Bedrock, select a Foundation Model: ai21.j2-mid-v1

Use Amazon Bedrock action in AWS Step Functions

For the Bedrock Model Parameters, you need to configure the required API parameters. Additionally, you can retrieve the default input parameter Comment of the Step Function execution as the value for the prompt of the API request:

{
  "prompt.$": "$.Comment",
  "maxTokens": 300,
  "temperature": 0.5,
  "topP": 0.9
}

To retrieve the generate text, you need to edit the Output configuration and provide $.Body.completions[0].data.text as the filter for the OutputPath:

Configure output in AWS Step Functions

Finally, click on Create and confirm the creation of the needed IAM Role.

Confirm IAM Role permissions for AWS Step Functions and Amazon Bedrock

Now, to use generative AI with Amazon Bedrock with AWS Step Functions, select Start Execution on the overview page:

Overview of executions in AWS Step Functions

The value of Content is passed as the prompt to Amazon Bedrock.

Start new execution in AWS Step Functions for Amazon Bedrock

Wait for the successful execution of your AWS Step Function and select Execution input and output to see the generated text.

Details of AWS Step Functions execution Output response of Amazon Bedrock in AWS Step Functions

Infrastructure as Code

You can create AWS Step Functions with CloudFormation and, of course, the AWS Cloud Development Kit. In the Management Console, you can easily retrieve your Step Function’s defintion:

{
  "Comment": "A description of my state machine",
  "StartAt": "Bedrock InvokeModel",
  "States": {
    "Bedrock InvokeModel": {
      "Type": "Task",
      "Resource": "arn:aws:states:::bedrock:invokeModel",
      "Parameters": {
        "ModelId": "arn:aws:bedrock:us-east-1::foundation-model/ai21.j2-mid-v1",
        "Body": {
          "prompt.$": "$.Comment",
          "maxTokens": 300,
          "temperature": 0.5,
          "topP": 0.9
        }
      },
      "End": true,
      "OutputPath": "$.Body.completions[0].data.text"
    }
  }
}

Following soon, a detailed post will show examples for the AWS Cloud Development Kit and Step Functions. Basics about AWS Step Functions and CDK are already available:

Well done, have a nice day!