Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(stepfunctions-tasks): add guardrailConfiguration and trace property to the BedrockInvokeModel #30426

Merged
merged 17 commits into from
Jul 5, 2024
Merged

feat(stepfunctions-tasks): add guardrailConfiguration and trace property to the BedrockInvokeModel #30426

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
747827a
feat: add guardrailConfiguration an trace to the BedrockInvokeModel
mazyu36 Jun 2, 2024
a54228f
fix: update README
mazyu36 Jun 2, 2024
094a115
a
mazyu36 Jun 5, 2024
9234818
incorporated review comments
mazyu36 Jun 7, 2024
aae5122
Update packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invok…
mazyu36 Jun 8, 2024
d7c850a
Update packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invok…
mazyu36 Jun 8, 2024
0a63f7f
error fixed
mazyu36 Jun 8, 2024
5f5072e
add validation
mazyu36 Jun 20, 2024
5781569
rename guardrailConfiguration to guardrail
mazyu36 Jun 21, 2024
ea54fba
fix: fix policy statement when guardrailIdentifier is set to arn
mazyu36 Jun 22, 2024
d8d5721
Merge branch 'main' into sfn-bedrock-guardrail-trace
mazyu36 Jun 22, 2024
be97141
Merge branch 'main' into sfn-bedrock-guardrail-trace
paulhcsun Jun 24, 2024
ad01e3b
incorporated review comments
mazyu36 Jun 29, 2024
4603188
Merge branch 'main' into sfn-bedrock-guardrail-trace
mazyu36 Jun 29, 2024
17506d6
update README
mazyu36 Jun 30, 2024
7f307c9
fix: validate arn
mazyu36 Jul 4, 2024
80fc07b
Merge branch 'main' into sfn-bedrock-guardrail-trace
mergify[bot] Jul 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,15 +221,18 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
}

if (this.props.guardrail) {
const isArn = this.props.guardrail.guardrailIdentifier.startsWith('arn:');
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The policy document was incorrect for the case when guardrailIdentifier is an ARN, so I added a condition to address this.

policyStatements.push(
new iam.PolicyStatement({
actions: ['bedrock:ApplyGuardrail'],
resources: [
Stack.of(this).formatArn({
service: 'bedrock',
resource: 'guardrail',
resourceName: this.props.guardrail.guardrailIdentifier,
}),
isArn
? this.props.guardrail.guardrailIdentifier
: Stack.of(this).formatArn({
service: 'bedrock',
resource: 'guardrail',
resourceName: this.props.guardrail.guardrailIdentifier,
}),
],
}),
);
Expand Down
111 changes: 110 additions & 1 deletion packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -361,7 +361,7 @@ describe('Invoke Model', () => {
}).toThrow(/Output S3 object version is not supported./);
});

test('guardrail', () => {
test('guardrail when gurdarilIdentifier is set to arn', () => {
// GIVEN
const stack = new cdk.Stack();
const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
Expand All @@ -381,6 +381,10 @@ describe('Invoke Model', () => {
},
});

new sfn.StateMachine(stack, 'StateMachine', {
definitionBody: sfn.DefinitionBody.fromChainable(task),
});

// THEN
expect(stack.resolve(task.toStateJson())).toEqual({
Type: 'Task',
Expand All @@ -407,6 +411,111 @@ describe('Invoke Model', () => {
GuardrailVersion: 'DRAFT',
},
});

Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', {
PolicyDocument: {
Statement: [
{
Action: 'bedrock:InvokeModel',
Effect: 'Allow',
Resource: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
},
{
Action: 'bedrock:ApplyGuardrail',
Effect: 'Allow',
Resource: 'arn:aws:bedrock:us-turbo-2:123456789012:guardrail/testid',
},
],
},
});
});

test('guardrail when gurdarilIdentifier is set to id', () => {
// 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,
contentType: 'application/json',
body: sfn.TaskInput.fromObject(
{
prompt: 'Hello world',
},
),
guardrail: {
guardrailIdentifier: 'testid',
guardrailVersion: 'DRAFT',
},
});

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,
Parameters: {
ModelId: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
Body: {
prompt: 'Hello world',
},
ContentType: 'application/json',
GuardrailIdentifier: 'testid',
GuardrailVersion: 'DRAFT',
},
});

Template.fromStack(stack).hasResourceProperties('AWS::IAM::Policy', {
PolicyDocument: {
Statement: [
{
Action: 'bedrock:InvokeModel',
Effect: 'Allow',
Resource: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
},
{
Action: 'bedrock:ApplyGuardrail',
Effect: 'Allow',
Resource: {
'Fn::Join': [
'',
[
'arn:',
{
Ref: 'AWS::Partition',
},
':bedrock:',
{
Ref: 'AWS::Region',
},
':',
{
Ref: 'AWS::AccountId',
},
':guardrail/testid',
],
],
},
},
],
},
});
});

test('guardrail fails when invalid guardrailIdentifier is set', () => {
Expand Down