Invoke Amazon Bedrock models from Step Functions with AWS CDK
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.0Install the dependencies
1pip3 install -r requirements.txtCreate 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.pyDeploy the stack
1cdk deployBecause of the CfnOutput, the statemachine’s arn is in an output:
1Outputs:
2StepFunctionBedrockStack.StateMachineArn = arn:aws:states:us-east-1:0123456789101:stateMachine:StateMachine2E01A3A5-GeJycj800pUNStart 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"}'