add validation

aws · Jun 21, 2024 · 57de8ad · 57de8ad
1 parent 0a63f7f
commit 57de8ad
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
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
@@ -276,11 +276,15 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     if (!Token.isUnresolved(guardrailIdentifier)) {
       const guardrailConfigurationPattern = /^(([a-z0-9]+)|(arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:[0-9]{12}:guardrail\/[a-z0-9]+))$/;
       if (!guardrailConfigurationPattern.test(guardrailIdentifier)) {
-        throw new Error(`guardrailIdentifier must match the ^(([a-z0-9]+)|(arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:[0-9]{12}:guardrail\/[a-z0-9]+))$ pattern, got ${guardrailIdentifier}`);
+        throw new Error(`You must set guardrailIdentifier to the id or the arn of Guardrail, got ${guardrailIdentifier}`);
       }
       if (props.contentType !== 'application/json') {
         throw new Error(`You must set contentType to \'application/json\' when using guardrailConfiguration, got '${props.contentType}'.`);
       }
+      if (guardrailIdentifier.length > 2048) {
+        throw new Error(`\`guardrailIdentifier\` length must be between 0 and 2048, got ${guardrailIdentifier.length}.`);
+      }
+
     }
 
     const guardrailVersionPattern = /^(([1-9][0-9]{0,7})|(DRAFT))$/;

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
@@ -416,7 +416,7 @@ describe('Invoke Model', () => {
 
     expect(() => {
       // WHEN
-      const task = new BedrockInvokeModel(stack, 'Invoke', {
+      new BedrockInvokeModel(stack, 'Invoke', {
         model,
         contentType: 'application/json',
         body: sfn.TaskInput.fromObject(
@@ -430,7 +430,32 @@ describe('Invoke Model', () => {
         },
       });
       // THEN
-    }).toThrow('guardrailIdentifier must match the ^(([a-z0-9]+)|(arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:[0-9]{12}:guardrail/[a-z0-9]+))$ pattern, got invalid-id');
+    }).toThrow('You must set guardrailIdentifier to the id or the arn of Guardrail, got invalid-id');
+  });
+
+  test('guardrail configuration fails when guardrailIdentifier length is invalid', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+    const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
+    const guardrailIdentifier = 'a'.repeat(2049);
+
+    expect(() => {
+      // WHEN
+      new BedrockInvokeModel(stack, 'Invoke', {
+        model,
+        contentType: 'application/json',
+        body: sfn.TaskInput.fromObject(
+          {
+            prompt: 'Hello world',
+          },
+        ),
+        guardrailConfiguration: {
+          guardrailIdentifier,
+          guardrailVersion: 'DRAFT',
+        },
+      });
+      // THEN
+    }).toThrow(`\`guardrailIdentifier\` length must be between 0 and 2048, got ${guardrailIdentifier.length}.`);
   });
 
   test('guardrail configuration fails when invalid guardrailVersion is set', () => {