From 7e5ff0090975a767a2ed1ca6a2bda6ac1941fe16 Mon Sep 17 00:00:00 2001 From: shikha372 Date: Tue, 23 Jul 2024 17:11:49 -0700 Subject: [PATCH] feat(stepfunctions-tasks): allow BedrockInvokeModel to use JsonPath (#30298) ### Issue # (if applicable) Closes #29229. ### Reason for this change When trying to use JsonPath to specify the S3 URIs that BedrockInvokeModel will read from and write from, you get an error. Example of the Error message: `jsii.errors.JavaScriptError: Error: Field references must be the entire string, cannot concatenate them (found 's3://${Token[prompt_bucket.348]}/${Token[prompt_key.349]}')` ### Description of changes Extended the inputPath property to be allowed as an input value for the task state. Instead of adding a new S3Uri props in current `BedrockInvokeModelProps` as proposed in the original issue, leveraged the `inputPath` property that is already defined in `sfn.TaskStateBaseProps` and being extended by `BedrockInvokeModelInputProps` and `BedrockInvokeModelOutputProps` **Limitation:** We cannot limit the resource policy to specific input token for which the value might be coming from the prompt, so had to keep it as [*] here. ### Description of how you validated changes Added unit tests. Successful deployment of integration tests in the account. ### Checklist - [x] Unit Tests - [x] Integration Tests - [x] Updated ReadMe - [x] My code adheres to the [CONTRIBUTING GUIDE](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md) and [DESIGN GUIDELINES](https://github.com/aws/aws-cdk/blob/main/docs/DESIGN_GUIDELINES.md) ---- *By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license* --- ...efaultTestDeployAssert9C0D2DFC.assets.json | 2 +- ...sks-bedrock-invoke-model-integ.assets.json | 6 +- ...s-bedrock-invoke-model-integ.template.json | 35 +++++- .../integ.invoke-model.js.snapshot/cdk.out | 2 +- .../integ.invoke-model.js.snapshot/integ.json | 2 +- .../manifest.json | 4 +- .../integ.invoke-model.js.snapshot/tree.json | 43 ++++++- .../test/bedrock/integ.invoke-model.ts | 8 +- .../aws-stepfunctions-tasks/README.md | 21 ++++ .../lib/bedrock/invoke-model.ts | 43 ++++++- .../test/bedrock/invoke-model.test.ts | 108 ++++++++++++++++++ 11 files changed, 255 insertions(+), 19 deletions(-) diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json index ec22c34fcf613..62c1eb1585be8 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json @@ -1,5 +1,5 @@ { - "version": "35.0.0", + "version": "36.0.0", "files": { "21fbb51d7b23f6a6c262b46a9caee79d744a3ac019fd45422d988b96d44b2a22": { "source": { diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json index 2d7004e16caf5..ea5387a06b82c 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json @@ -1,7 +1,7 @@ { - "version": "35.0.0", + "version": "36.0.0", "files": { - "19db222d8d51351d1127c4b099aa6545a4c1ddd9425a2e0f78c328f39ff74edf": { + "e9d946bbac52fb88d3fc7c9ea4f26da0e6a6965417c8ae9fb5464e5342269444": { "source": { "path": "aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json", "packaging": "file" @@ -9,7 +9,7 @@ "destinations": { "current_account-current_region": { "bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}", - "objectKey": "19db222d8d51351d1127c4b099aa6545a4c1ddd9425a2e0f78c328f39ff74edf.json", + "objectKey": "e9d946bbac52fb88d3fc7c9ea4f26da0e6a6965417c8ae9fb5464e5342269444.json", "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}" } } diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json index ce41a5e96cc20..a67f72f4ea158 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json @@ -41,6 +41,25 @@ ] ] } + }, + { + "Action": [ + "s3:GetObject", + "s3:PutObject" + ], + "Effect": "Allow", + "Resource": { + "Fn::Join": [ + "", + [ + "arn:", + { + "Ref": "AWS::Partition" + }, + ":s3:::*" + ] + ] + } } ], "Version": "2012-10-17" @@ -72,7 +91,19 @@ { "Ref": "AWS::Region" }, - "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText\":\"Generate a list of five first names.\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt2\":{\"End\":true,\"Type\":\"Task\",\"ResultPath\":\"$\",\"ResultSelector\":{\"names.$\":\"$.Body.results[0].outputText\"},\"Resource\":\"arn:", + "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText\":\"Generate a list of five first names.\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt2\":{\"Next\":\"Prompt3\",\"Type\":\"Task\",\"ResultPath\":\"$\",\"ResultSelector\":{\"names.$\":\"$.Body.results[0].outputText\"},\"Resource\":\"arn:", + { + "Ref": "AWS::Partition" + }, + ":states:::bedrock:invokeModel\",\"Parameters\":{\"ModelId\":\"arn:", + { + "Ref": "AWS::Partition" + }, + ":bedrock:", + { + "Ref": "AWS::Region" + }, + "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"End\":true,\"Type\":\"Task\",\"InputPath\":\"$.names\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:", { "Ref": "AWS::Partition" }, @@ -84,7 +115,7 @@ { "Ref": "AWS::Region" }, - "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}}},\"TimeoutSeconds\":30}" + "::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri.$\":\"$.names\"},\"Output\":{\"S3Uri.$\":\"$.names\"}}}},\"TimeoutSeconds\":30}" ] ] }, diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out index c5cb2e5de6344..1f0068d32659a 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out @@ -1 +1 @@ -{"version":"35.0.0"} \ No newline at end of file +{"version":"36.0.0"} \ No newline at end of file diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json index 5eb622d1d7b82..36143be21ab57 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json @@ -1,5 +1,5 @@ { - "version": "35.0.0", + "version": "36.0.0", "testCases": { "InvokeModel/DefaultTest": { "stacks": [ diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json index 60de1b8de8ab7..677092f63b132 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json @@ -1,5 +1,5 @@ { - "version": "35.0.0", + "version": "36.0.0", "artifacts": { "aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets": { "type": "cdk:asset-manifest", @@ -18,7 +18,7 @@ "validateOnSynth": false, "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-deploy-role-${AWS::AccountId}-${AWS::Region}", "cloudFormationExecutionRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-cfn-exec-role-${AWS::AccountId}-${AWS::Region}", - "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/19db222d8d51351d1127c4b099aa6545a4c1ddd9425a2e0f78c328f39ff74edf.json", + "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/e9d946bbac52fb88d3fc7c9ea4f26da0e6a6965417c8ae9fb5464e5342269444.json", "requiresBootstrapStackVersion": 6, "bootstrapStackVersionSsmParameter": "/cdk-bootstrap/hnb659fds/version", "additionalDependencies": [ diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json index b3a2882dbd6fb..17b23007671bd 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json @@ -24,6 +24,14 @@ "version": "0.0.0" } }, + "Prompt3": { + "id": "Prompt3", + "path": "aws-stepfunctions-tasks-bedrock-invoke-model-integ/Prompt3", + "constructInfo": { + "fqn": "aws-cdk-lib.aws_stepfunctions_tasks.BedrockInvokeModel", + "version": "0.0.0" + } + }, "StateMachine": { "id": "StateMachine", "path": "aws-stepfunctions-tasks-bedrock-invoke-model-integ/StateMachine", @@ -96,6 +104,25 @@ ] ] } + }, + { + "Action": [ + "s3:GetObject", + "s3:PutObject" + ], + "Effect": "Allow", + "Resource": { + "Fn::Join": [ + "", + [ + "arn:", + { + "Ref": "AWS::Partition" + }, + ":s3:::*" + ] + ] + } } ], "Version": "2012-10-17" @@ -147,7 +174,19 @@ { "Ref": "AWS::Region" }, - "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText\":\"Generate a list of five first names.\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt2\":{\"End\":true,\"Type\":\"Task\",\"ResultPath\":\"$\",\"ResultSelector\":{\"names.$\":\"$.Body.results[0].outputText\"},\"Resource\":\"arn:", + "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText\":\"Generate a list of five first names.\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt2\":{\"Next\":\"Prompt3\",\"Type\":\"Task\",\"ResultPath\":\"$\",\"ResultSelector\":{\"names.$\":\"$.Body.results[0].outputText\"},\"Resource\":\"arn:", + { + "Ref": "AWS::Partition" + }, + ":states:::bedrock:invokeModel\",\"Parameters\":{\"ModelId\":\"arn:", + { + "Ref": "AWS::Partition" + }, + ":bedrock:", + { + "Ref": "AWS::Region" + }, + "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"End\":true,\"Type\":\"Task\",\"InputPath\":\"$.names\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:", { "Ref": "AWS::Partition" }, @@ -159,7 +198,7 @@ { "Ref": "AWS::Region" }, - "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}}},\"TimeoutSeconds\":30}" + "::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri.$\":\"$.names\"},\"Output\":{\"S3Uri.$\":\"$.names\"}}}},\"TimeoutSeconds\":30}" ] ] }, diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts index 95860f09e624e..2d98f804b4475 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts @@ -52,7 +52,13 @@ const prompt2 = new BedrockInvokeModel(stack, 'Prompt2', { resultPath: '$', }); -const chain = sfn.Chain.start(prompt1).next(prompt2); +const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', { + model, + inputPath: sfn.JsonPath.stringAt('$.names'), + outputPath: sfn.JsonPath.stringAt('$.names'), +}); + +const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3); new sfn.StateMachine(stack, 'StateMachine', { definitionBody: sfn.DefinitionBody.fromChainable(chain), diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md b/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md index 4b288aafd7f93..9f5943afaa16b 100644 --- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md +++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md @@ -398,6 +398,27 @@ const task = new tasks.BedrockInvokeModel(this, 'Prompt Model', { names: sfn.JsonPath.stringAt('$.Body.results[0].outputText'), }, }); +``` +### Using Input Path + +Provide S3 URI as an input or output path to invoke a model + +```ts + +import * as bedrock from 'aws-cdk-lib/aws-bedrock'; + +const model = bedrock.FoundationModel.fromFoundationModelId( + this, + 'Model', + bedrock.FoundationModelIdentifier.AMAZON_TITAN_TEXT_G1_EXPRESS_V1, +); + +const task = new tasks.BedrockInvokeModel(this, 'Prompt Model', { + model, + inputPath: sfn.JsonPath.stringAt('$.prompt'), + outputPath: sfn.JsonPath.stringAt('$.prompt'), +}); + ``` You can apply a guardrail to the invocation by setting `guardrail`. diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts index f831cf606c8c5..3abc7338cab8c 100644 --- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts +++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts @@ -140,12 +140,14 @@ export class BedrockInvokeModel extends sfn.TaskStateBase { constructor(scope: Construct, id: string, private readonly props: BedrockInvokeModelProps) { super(scope, id, props); + this.integrationPattern = props.integrationPattern ?? sfn.IntegrationPattern.REQUEST_RESPONSE; validatePatternSupported(this.integrationPattern, BedrockInvokeModel.SUPPORTED_INTEGRATION_PATTERNS); const isBodySpecified = props.body !== undefined; - const isInputSpecified = props.input !== undefined && props.input.s3Location !== undefined; + //Either specific props.input with bucket name and object key or input s3 path + const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.inputPath !== undefined); if (isBodySpecified && isInputSpecified) { throw new Error('Either `body` or `input` must be specified, but not both.'); @@ -171,7 +173,21 @@ export class BedrockInvokeModel extends sfn.TaskStateBase { }), ]; - if (this.props.input !== undefined && this.props.input.s3Location !== undefined) { + if (this.props.inputPath !== undefined) { + policyStatements.push( + new iam.PolicyStatement({ + actions: ['s3:GetObject'], + resources: [ + Stack.of(this).formatArn({ + region: '', + account: '', + service: 's3', + resource: '*', + }), + ], + }), + ); + } else if (this.props.input !== undefined && this.props.input.s3Location !== undefined) { policyStatements.push( new iam.PolicyStatement({ actions: ['s3:GetObject'], @@ -188,7 +204,21 @@ export class BedrockInvokeModel extends sfn.TaskStateBase { ); } - if (this.props.output !== undefined && this.props.output.s3Location !== undefined) { + if (this.props.outputPath !== undefined) { + policyStatements.push( + new iam.PolicyStatement({ + actions: ['s3:PutObject'], + resources: [ + Stack.of(this).formatArn({ + region: '', + account: '', + service: 's3', + resource: '*', + }), + ], + }), + ); + } else if (this.props.output !== undefined && this.props.output.s3Location !== undefined) { policyStatements.push( new iam.PolicyStatement({ actions: ['s3:PutObject'], @@ -241,10 +271,10 @@ export class BedrockInvokeModel extends sfn.TaskStateBase { Body: this.props.body?.value, Input: this.props.input?.s3Location ? { S3Uri: `s3://${this.props.input.s3Location.bucketName}/${this.props.input.s3Location.objectKey}`, - } : undefined, + } : this.props.inputPath ? { S3Uri: this.props.inputPath } : undefined, Output: this.props.output?.s3Location ? { S3Uri: `s3://${this.props.output.s3Location.bucketName}/${this.props.output.s3Location.objectKey}`, - } : undefined, + } : this.props.outputPath ? { S3Uri: this.props.outputPath }: undefined, GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier, GuardrailVersion: this.props.guardrail?.guardrailVersion, Trace: this.props.traceEnabled === undefined @@ -254,5 +284,6 @@ export class BedrockInvokeModel extends sfn.TaskStateBase { : 'DISABLED', }), }; - } + }; } + diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts index 3c480ea1638a3..02fa2737bb2b8 100644 --- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts +++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts @@ -204,6 +204,52 @@ describe('Invoke Model', () => { }); }); + test('invoke model allows input and output json path', () => { + const stack = new cdk.Stack(); + const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123'); + + const task = new BedrockInvokeModel(stack, 'Invoke', { + model, + inputPath: sfn.JsonPath.stringAt('$.prompt'), + outputPath: sfn.JsonPath.stringAt('$.prompt'), + }); + + new sfn.StateMachine(stack, 'StateMachine', { + definitionBody: sfn.DefinitionBody.fromChainable(task), + }); + + // THEN + expect(stack.resolve(task.toStateJson())).toEqual({ + Type: 'Task', + Resource: { + 'Fn::Join': [ + '', + [ + 'arn:', + { + Ref: 'AWS::Partition', + }, + ':states:::bedrock:invokeModel', + ], + ], + }, + End: true, + InputPath: '$.prompt', + OutputPath: '$.prompt', + Parameters: { + ModelId: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123', + Input: { + //Expected key modified from S3Uri to S3Uri.$ as per the State Machine context key field transformation + //Reference: https://docs.aws.amazon.com/step-functions/latest/dg/input-output-example.html + 'S3Uri.$': '$.prompt', + }, + Output: { + 'S3Uri.$': '$.prompt', + }, + }, + }); + }); + test('S3 permissions are created in generated policy when input and output locations are specified', () => { // GIVEN const stack = new cdk.Stack(); @@ -276,6 +322,68 @@ describe('Invoke Model', () => { }); }); + test('S3 permissions are created in generated policy when input and output path are specified', () => { + // GIVEN + const stack = new cdk.Stack(); + const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123'); + + // WHEN + const task = new BedrockInvokeModel(stack, 'Invoke', { + model, + inputPath: sfn.JsonPath.stringAt('$.prompt'), + outputPath: sfn.JsonPath.stringAt('$.prompt'), + }); + + new sfn.StateMachine(stack, 'StateMachine', { + definitionBody: sfn.DefinitionBody.fromChainable(task), + }); + + // THEN + Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', { + PolicyDocument: Match.objectLike({ + Statement: Match.arrayWith([ + { + Action: 'bedrock:InvokeModel', + Effect: 'Allow', + Resource: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123', + }, + { + Action: 's3:GetObject', + Effect: 'Allow', + Resource: { + 'Fn::Join': [ + '', + [ + 'arn:', + { + Ref: 'AWS::Partition', + }, + ':s3:::*', + ], + ], + }, + }, + { + Action: 's3:PutObject', + Effect: 'Allow', + Resource: { + 'Fn::Join': [ + '', + [ + 'arn:', + { + Ref: 'AWS::Partition', + }, + ':s3:::*', + ], + ], + }, + }, + ]), + }), + }); + }); + test('fails on neither input nor body set', () => { // GIVEN const stack = new cdk.Stack();