Reading · how.wtf · RSS Amplifier

invoke Amazon bedrock models directly from Step Function statemachines

AWS released an optimized integration for Amazon Bedrock. It allows StepFunction statemachines to directly call the Amazon Bedrock API without needing to write a Lambda function.

How to deploy a StepFunction with Amazon Bedrock integration using AWS CDK

Follow the instructions below to deploy a StepFunction with Amazon Bedrock integration using the AWS CDK. The following tutorial uses Python, but the same principles can be used with any language.

Create a requirements.txt

For this tutorial, version 2.115.0 of the AWS CDK was used.

1aws-cdk-lib==2.115.0
2constructs>=10.0.0,<11.0.0

Install the dependencies

1pip3 install -r requirements.txt

Create stack.py with the step function and bedrock task

For this tutorial, I opted to use the Claude V1 instant model. If you’re unsure where the sfn.TaskInput.from_object() parameters originate from, please review my article describing how to invoke Bedrock models directly.

 1from aws_cdk import (
 2    App,
 3    CfnOutput,
 4    Stack,
 5    aws_stepfunctions as sfn,
 6    aws_stepfunctions_tasks as tasks,
 7    aws_bedrock as bedrock,
 8)
 9
10
11class StepFunctionBedrockStack(Stack):
12    def __init__(self, scope: App, construct_id: str, **kwargs) -> None:
13        super().__init__(scope, construct_id, **kwargs)
14
15        model = bedrock.FoundationModel.from_foundation_model_id(
16            self, "Model", bedrock.FoundationModelIdentifier.ANTHROPIC_CLAUDE_INSTANT_V1
17        )
18        invoke_model_task = tasks.BedrockInvokeModel(
19            self,
20            "InvokeModel",
21            model=model,
22            body=sfn.TaskInput.from_object(
23                {
24                    "prompt": "\n\nHuman:Give me a cowboy joke\n\nAssistant:",
25                    "max_tokens_to_sample": 256,
26                }
27            ),
28            result_selector={"joke": sfn.JsonPath.string_at("$.Body.completion")},
29        )
30        state_machine = sfn.StateMachine(
31            self,
32            "StateMachine",
33            definition_body=sfn.DefinitionBody.from_chainable(invoke_model_task),
34        )
35
36        CfnOutput(self, "StateMachineArn", value=state_machine.state_machine_arn)

The code above creates a BedrockInvokeModel task with a body that’s specific for Anthropic Claude’s input and a result_selector that extracts the completion.

Create app.py

1import aws_cdk as cdk
2
3from stack import StepFunctionBedrockStack
4
5app = cdk.App()
6StepFunctionBedrockStack(app, "StepFunctionBedrockStack")
7
8app.synth()

Create cdk.json

1{
2  "app": "python3 app.py"
3}

The directory structure should look like this:

1project/
2├── app.py
3├── cdk.json
4├── requirements.txt
5└── stack.py

Deploy the stack

1cdk deploy

Because of the CfnOutput, the statemachine’s arn is in an output:

1Outputs:
2StepFunctionBedrockStack.StateMachineArn = arn:aws:states:us-east-1:0123456789101:stateMachine:StateMachine2E01A3A5-GeJycj800pUN

Start a statemachine execution

With the statemachine name, we can start an execution using the AWS CLI:

1aws stepfunctions start-execution  \
2    --state-machine-arn <statemachine arn here>

Output:

1{
2    "executionArn": "arn:aws:states:us-east-1:0123456789101:execution:StateMachine2E01A3A5-GeJycj800pUN:387b0976-feb6-4c8c-9abe-af427546df3c",
3    "startDate": "2023-12-20T01:48:56.396000-05:00"
4}

then, we can describe the execution using the execution arn:

1aws stepfunctions describe-execution \
2    --execution-arn <execution arn here>

Output:

 1{
 2    "executionArn": "arn:aws:states:us-east-1:0123456789101:execution:StateMachine2E01A3A5-GeJycj800pUN:387b0976-feb6-4c8c-9abe-af427546df3c",
 3    "stateMachineArn": "arn:aws:states:us-east-1:0123456789101:stateMachine:StateMachine2E01A3A5-GeJycj800pUN",
 4    "name": "387b0976-feb6-4c8c-9abe-af427546df3c",
 5    "status": "SUCCEEDED",
 6    "startDate": "2023-12-20T01:48:56.396000-05:00",
 7    "stopDate": "2023-12-20T01:48:57.814000-05:00",
 8    "input": "{}",
 9    "inputDetails": {
10        "included": true
11    },
12    "output": "{\"joke\":\" Here's one: Why don't seagulls fly over the bay? Because then they'd be bagels!\"}",
13    "outputDetails": {
14        "included": true
15    },
16    "redriveCount": 0,
17    "redriveStatus": "NOT_REDRIVABLE",
18    "redriveStatusReason": "Execution is SUCCEEDED and cannot be redriven"
19}

Pass information from the statemachine input to Bedrock

If you want to pass the input to the foundational model, you can do this using a combination of:

sfn.JsonPatch.string_at() and States.Format intrinsic function

 1from aws_cdk import (
 2    App,
 3    CfnOutput,
 4    Stack,
 5    aws_stepfunctions as sfn,
 6    aws_stepfunctions_tasks as tasks,
 7    aws_bedrock as bedrock,
 8)
 9
10
11class StepFunctionBedrockStack(Stack):
12    def __init__(self, scope: App, construct_id: str, **kwargs) -> None:
13        super().__init__(scope, construct_id, **kwargs)
14
15        model = bedrock.FoundationModel.from_foundation_model_id(
16            self, "Model", bedrock.FoundationModelIdentifier.ANTHROPIC_CLAUDE_INSTANT_V1
17        )
18        invoke_model_task = tasks.BedrockInvokeModel(
19            self,
20            "InvokeModel",
21            model=model,
22            body=sfn.TaskInput.from_object(
23                {
24                    "prompt": sfn.JsonPath.string_at(
25                        "States.Format('\n\nHuman:{}\n\nAssistant:', $$.Execution.Input.prompt)"
26                    ),
27                    "max_tokens_to_sample": 256,
28                }
29            ),
30            result_selector={"joke": sfn.JsonPath.string_at("$.Body.completion")},
31        )
32        state_machine = sfn.StateMachine(
33            self,
34            "StateMachine",
35            definition_body=sfn.DefinitionBody.from_chainable(invoke_model_task),
36        )
37
38        CfnOutput(self, "StateMachineArn", value=state_machine.state_machine_arn)

Now we can start the statemachine execution with a prompt:

1aws stepfunctions start-execution  \
2    --state-machine-arn <statemachine arn here> \
3    --input '{"prompt": "Write me a joke"}'

Invoke Amazon Bedrock models from Step Functions with AWS CDK