feat: add IAM Role support for bedrock invokeModel with foundation mo…

…dels
serverless-operations · Feb 2, 2024 · 8282395 · 8282395
1 parent f2f379d
commit 8282395
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 0 deletions.
diff --git 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
@@ -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,