diff --git a/lib/deploy/stepFunctions/compileIamRole.js b/lib/deploy/stepFunctions/compileIamRole.js index 5b3ed809..47b1e2e7 100644 --- a/lib/deploy/stepFunctions/compileIamRole.js +++ b/lib/deploy/stepFunctions/compileIamRole.js @@ -524,6 +524,23 @@ function getSageMakerPermissions(state) { ]; } +function getBedrockPermissions(state) { + const modelId = state.Parameters.ModelId; + const modelArn = modelId.startsWith('arn:') ? modelId : { + 'Fn::Sub': [ + `arn:\${AWS::Partition}:bedrock:$\{AWS::Region}::foundation-model/${modelId}`, + {}, + ], + }; + + return [ + { + action: 'bedrock:InvokeModel', + resource: modelArn, + }, + ]; +} + function getEventBridgePermissions(state) { const eventBuses = new Set(); @@ -683,6 +700,9 @@ function getIamPermissions(taskStates) { case 'arn:aws:states:::sagemaker:createTransformJob.sync': return getSageMakerPermissions(state); + case 'arn:aws:states:::bedrock:invokeModel': + return getBedrockPermissions(state); + case 'arn:aws:states:::events:putEvents': case 'arn:aws:states:::events:putEvents.waitForTaskToken': return getEventBridgePermissions(state); diff --git a/lib/deploy/stepFunctions/compileIamRole.test.js b/lib/deploy/stepFunctions/compileIamRole.test.js index f72480e4..c79d9da5 100644 --- a/lib/deploy/stepFunctions/compileIamRole.test.js +++ b/lib/deploy/stepFunctions/compileIamRole.test.js @@ -3559,6 +3559,69 @@ describe('#compileIamRole', () => { }]); }); + it('should give bedrock invoke permissions for foundation models', () => { + serverless.service.stepFunctions = { + stateMachines: { + myStateMachine1: { + id: 'StateMachine1', + definition: { + StartAt: 'A', + States: { + A: { + Type: 'Task', + Resource: 'arn:aws:states:::bedrock:invokeModel', + Parameters: { + ModelId: 'anthropic.claude-v2:1', + Body: { + prompt: 'your-prompt', + max_tokens_to_sample: 500, + temperature: 0.1, + }, + ContentType: 'application/json', + Accept: 'application/json', + }, + Next: 'B', + }, + B: { + Type: 'Task', + Resource: 'arn:aws:states:::bedrock:invokeModel', + Parameters: { + // modelId can be specified as an arn + ModelId: 'arn:aws:bedrock:us-east-1::foundation-model/meta.llama2-70b-chat-v1', + Body: { + prompt: 'your-prompt', + max_tokens_to_sample: 500, + temperature: 0.1, + }, + ContentType: 'application/json', + Accept: 'application/json', + }, + End: true, + }, + }, + }, + }, + }, + }; + + serverlessStepFunctions.compileIamRole(); + const statements = serverlessStepFunctions.serverless.service + .provider.compiledCloudFormationTemplate.Resources.StateMachine1Role + .Properties.Policies[0].PolicyDocument.Statement; + const bedrockPermissions = statements.filter(s => _.isEqual(s.Action, ['bedrock:InvokeModel'])); + expect(bedrockPermissions).to.have.lengthOf(1); + expect(bedrockPermissions[0].Resource).to.have.lengthOf(2); + expect(bedrockPermissions[0].Resource).to.deep.eq([ + { + 'Fn::Sub': [ + 'arn:${AWS::Partition}:bedrock:${AWS::Region}::foundation-model/anthropic.claude-v2:1', + {}, + ], + }, + 'arn:aws:bedrock:us-east-1::foundation-model/meta.llama2-70b-chat-v1', + ]); + }); + it('should give event bridge putEvents permissions', () => { const genStateMachine = id => ({ id,