addressing comments and last feedback

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

@@ -103,7 +103,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}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
+       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names: {}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.Body.results[0].outputText\",\"Resource\":\"arn:",
        {
         "Ref": "AWS::Partition"
        },
@@ -115,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}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
+       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Echo list of first names: {}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
        {
         "Ref": "AWS::Partition"
        },
@@ -127,7 +127,7 @@
        {
         "Ref": "AWS::Region"
        },
-       "::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri.$\":\"$.names\"},\"Output\":{\"S3Uri.$\":\"$.names\"}}}},\"TimeoutSeconds\":30}"
+       "::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri\":\"$.names\"},\"Output\":{\"S3Uri\":\"$.names\"}}}},\"TimeoutSeconds\":30}"
       ]
      ]
     },

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts

@@ -41,7 +41,7 @@ const prompt2 = new BedrockInvokeModel(stack, 'Prompt2', {
   body: sfn.TaskInput.fromObject(
     {
       inputText: sfn.JsonPath.format(
-        'Alphabetize this list of first names:\n{}',
+        'Alphabetize this list of first names: {}',
         sfn.JsonPath.stringAt('$.names'),
       ),
       textGenerationConfig: {
@@ -62,7 +62,7 @@ const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
   body: sfn.TaskInput.fromObject(
     {
       inputText: sfn.JsonPath.format(
-        'Alphabetize this list of first names:\n{}',
+        'Echo list of first names: {}',
         sfn.JsonPath.stringAt('$.names'),
       ),
       textGenerationConfig: {
@@ -71,35 +71,18 @@ const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
       },
     },
   ),
-  outputPath: sfn.JsonPath.stringAt('$.names'),
-});
-
-/**Test for Bedrock Input Path */
-const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
-  model,
-  body: sfn.TaskInput.fromObject(
-    {
-      inputText: sfn.JsonPath.format(
-        'Alphabetize this list of first names:\n{}',
-        sfn.JsonPath.stringAt('$.names'),
-      ),
-      textGenerationConfig: {
-        maxTokenCount: 100,
-        temperature: 1,
-      },
-    },
-  ),
-  inputPath: sfn.JsonPath.stringAt('$.names'),
+  outputPath: '$.Body.results[0].outputText',
 });
 
 /** Test for Bedrock s3 URI Path */
-const prompt5 = new BedrockInvokeModel(stack, 'Prompt5', {
+//Execution will fail for the following input as it expects a valid s3 URI from previous prompt
+const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
   model,
-  input: { s3InputUri: sfn.JsonPath.stringAt('$.names') },
-  output: { s3OutputUri: sfn.JsonPath.stringAt('$.names') },
+  input: { s3InputUri: '$.names' },
+  output: { s3OutputUri: '$.names' },
 });
 
-const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4).next(prompt5);
+const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4);
 
 new sfn.StateMachine(stack, 'StateMachine', {
   definitionBody: sfn.DefinitionBody.fromChainable(chain),

packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts

@@ -164,12 +164,12 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
 
     validatePatternSupported(this.integrationPattern, BedrockInvokeModel.SUPPORTED_INTEGRATION_PATTERNS);
 
-    const isFeatureFlagEnabled = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
+    const useNewS3UriParamsForTask = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
 
     const isBodySpecified = props.body !== undefined;
 
     let isInputSpecified: boolean;
-    if (!isFeatureFlagEnabled) {
+    if (!useNewS3UriParamsForTask) {
       isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.inputPath !== undefined);
     } else {
       //Either specific props.input with bucket name and object key or input s3 path
@@ -188,17 +188,24 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     if (props.output?.s3Location?.objectVersion !== undefined) {
       throw new Error('Output S3 object version is not supported.');
     }
+    if (props.input?.s3InputUri && props.input.s3Location) {
+      throw new Error('Cannot specify both S3 InputUri and S3 location');
+    }
+    if (props.input?.s3InputUri === '') {
+      throw new Error('S3 InputUri cannot be an empty string');
+    }
 
     //Warning to let users know about the newly introduced props
-    if (props.inputPath || props.outputPath && !isFeatureFlagEnabled) {
+    if (props.inputPath || props.outputPath && !useNewS3UriParamsForTask) {
       Annotations.of(scope).addWarningV2('aws-cdk-lib/aws-stepfunctions-taks',
-        'These props will set the value of inputPath/outputPath as s3 URI under input/output field in state machine JSON definition. To modify the behaviour set feature flag `@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true` and use props s3InputUri/s3OutputUri');
+        'These props will set the value of inputPath/outputPath as s3 URI under input/output field in state machine JSON definition. To modify the behaviour set feature flag `@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true` and use props input.s3InputUri/output.s3OutputUri');
     }
 
-    this.taskPolicies = this.renderPolicyStatements(isFeatureFlagEnabled);
+    this.taskPolicies = this.renderPolicyStatements();
   }
 
-  private renderPolicyStatements(isFeatureFlagEnabled?: boolean): iam.PolicyStatement[] {
+  private renderPolicyStatements(): iam.PolicyStatement[] {
+    const useNewS3UriParamsForTask = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
     const policyStatements = [
       new iam.PolicyStatement({
         actions: ['bedrock:InvokeModel'],
@@ -207,7 +214,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     ];
 
     //For Compatibility with existing behaviour of input path
-    if (this.props.input?.s3InputUri !== undefined || (!isFeatureFlagEnabled && this.props.inputPath !== undefined)) {
+    if (this.props.input?.s3InputUri !== undefined || (!useNewS3UriParamsForTask && this.props.inputPath !== undefined)) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:GetObject'],
@@ -239,7 +246,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     }
 
     //For Compatibility with existing behaviour of output path
-    if (this.props.output?.s3OutputUri !== undefined || (!isFeatureFlagEnabled && this.props.outputPath !== undefined)) {
+    if (this.props.output?.s3OutputUri !== undefined || (!useNewS3UriParamsForTask && this.props.outputPath !== undefined)) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:PutObject'],
@@ -298,18 +305,18 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
    */
   protected _renderTask(): any {
 
-    const isFeatureFlagEnabled = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
-    const inputSource = this.getInputSource(this.props.input, this.props.inputPath, isFeatureFlagEnabled);
-    const outputSource = this.getOutputSource(this.props.output, this.props.outputPath, isFeatureFlagEnabled);
+    const useNewS3UriParamsForTask = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
+    const inputSource = this.getInputSource(this.props.input, this.props.inputPath, useNewS3UriParamsForTask);
+    const outputSource = this.getOutputSource(this.props.output, this.props.outputPath, useNewS3UriParamsForTask);
     return {
       Resource: integrationResourceArn('bedrock', 'invokeModel'),
       Parameters: sfn.FieldUtils.renderObject({
         ModelId: this.props.model.modelArn,
         Accept: this.props.accept,
         ContentType: this.props.contentType,
         Body: this.props.body?.value,
-        Input: this.props.input || (this.props.inputPath && !isFeatureFlagEnabled) ? { S3Uri: inputSource } : undefined,
-        Output: this.props.output || ( this.props.outputPath && !isFeatureFlagEnabled) ? { S3Uri: outputSource } : undefined,
+        Input: inputSource ? { S3Uri: inputSource } : undefined,
+        Output: outputSource ? { S3Uri: outputSource } : undefined,
         GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
         GuardrailVersion: this.props.guardrail?.guardrailVersion,
         Trace: this.props.traceEnabled === undefined
@@ -321,23 +328,23 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     };
   };
 
-  private getInputSource(props?: BedrockInvokeModelInputProps, inputPath?: string, isFeatureFlagEnabled?: boolean): string | undefined {
+  private getInputSource(props?: BedrockInvokeModelInputProps, inputPath?: string, useNewS3UriParamsForTask?: boolean): string | undefined {
     if (props?.s3Location) {
       return `s3://${props.s3Location.bucketName}/${props.s3Location.objectKey}`;
-    } else if (isFeatureFlagEnabled && props?.s3InputUri) {
+    } else if (useNewS3UriParamsForTask && props?.s3InputUri) {
       return props.s3InputUri;
-    } else if (!isFeatureFlagEnabled && inputPath) {
+    } else if (!useNewS3UriParamsForTask && inputPath) {
       return inputPath;
     }
     return undefined;
   }
 
-  private getOutputSource(props?: BedrockInvokeModelOutputProps, outputPath?: string, isFeatureFlagEnabled?: boolean): string | undefined {
+  private getOutputSource(props?: BedrockInvokeModelOutputProps, outputPath?: string, useNewS3UriParamsForTask?: boolean): string | undefined {
     if (props?.s3Location) {
       return `s3://${props.s3Location.bucketName}/${props.s3Location.objectKey}`;
-    } else if (isFeatureFlagEnabled && props?.s3OutputUri) {
+    } else if (useNewS3UriParamsForTask && props?.s3OutputUri) {
       return props.s3OutputUri;
-    } else if (!isFeatureFlagEnabled && outputPath) {
+    } else if (!useNewS3UriParamsForTask && outputPath) {
       return outputPath;
     }
     return undefined;

packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md

@@ -72,7 +72,7 @@ Flags come in three types:
 | [@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm](#aws-cdkaws-ecsremovedefaultdeploymentalarm) | When enabled, remove default deployment alarm settings | 2.143.0 | (default) |
 | [@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault](#aws-cdkcustom-resourceslogapiresponsedatapropertytruedefault) | When enabled, the custom resource used for `AwsCustomResource` will configure the `logApiResponseData` property as true by default | 2.145.0 | (fix) |
 | [@aws-cdk/aws-s3:keepNotificationInImportedBucket](#aws-cdkaws-s3keepnotificationinimportedbucket) | When enabled, Adding notifications to a bucket in the current stack will not remove notification from imported stack. | 2.155.0 | (fix) |
-| [@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | 2.155.1 | (fix) |
+| [@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | 2.156.0 | (fix) |
 
 <!-- END table -->
 
@@ -1372,7 +1372,7 @@ When this feature flag is enabled, specify newly introduced props 's3InputUri' a
 | Since | Default | Recommended |
 | ----- | ----- | ----- |
 | (not in v1) |  |  |
-| 2.155.1 | `true` | `true` |
+| 2.156.0 | `true` | `true` |
 
 
 <!-- END details -->