From 08be4669945c82a72ea2c8b253e6309f6a18f78e Mon Sep 17 00:00:00 2001
From: Abhishek Patil <abhirpat@amazon.com>
Date: Tue, 12 Nov 2024 10:28:03 -0500
Subject: [PATCH] feat: streaming support for qnabot II

---
 config/utils/merge-config.js    |  2 ++
 templates/codebuild-deploy.yaml | 18 +++++-------------
 templates/master.yaml           | 10 +---------
 3 files changed, 8 insertions(+), 22 deletions(-)

diff --git a/config/utils/merge-config.js b/config/utils/merge-config.js
index fab35fed..2026d857 100644
--- a/config/utils/merge-config.js
+++ b/config/utils/merge-config.js
@@ -33,6 +33,8 @@ module.exports = function mergeConfig(baseConfig, srcConfig) {
             key=== 'initialText' ||
             key=== 'avatarImageUrl' ||
             key=== 'toolbarLogo' ||
+            key=== 'streamingWebSocketEndpoint' ||
+            key=== 'streamingDynamoDbTable' ||
             !isEmpty(srcConfig[key]) ) {
             value = (typeof (baseConfig[key]) === 'object') ?
                 // recursively merge sub-objects in both directions
diff --git a/templates/codebuild-deploy.yaml b/templates/codebuild-deploy.yaml
index 1e82a7a2..e5822145 100644
--- a/templates/codebuild-deploy.yaml
+++ b/templates/codebuild-deploy.yaml
@@ -546,13 +546,7 @@ Parameters:
         Type: String
         Default: ''
         Description: >
-            If not using the Web socket endpoint created via AllowStreamingResponses, you can optionally specify an alternative streaming endpoint here
-
-    StreamingDynamoDbTable:
-        Type: String
-        Default: ''
-        Description: >
-            If not using the DynamoDB table through Lex Web UI, you can optionally specify an alternative table for streaming purposes here
+            If you have an existing WebSocket API Gateway endpoint, you can specify it using this parameter. This requires parameter AllowStreamingResponses set to True.
 
     ShouldEnableUpload:
         Type: String
@@ -587,8 +581,7 @@ Conditions:
     UseDefaultCloudfrontUrl: !Or [ !Equals [!Ref WebAppConfCname, ''], !Equals [!Ref WebAppAcmCertificateArn, ''] ]
     ShouldNotSpecifyWafAcl: !Equals [!Ref WebAppWafAclArn, '']
     EnableStreaming: !Equals [!Ref AllowStreamingResponses, true]
-    HasStreamingWebSocketEndpoint: !Not [!Equals [!Ref StreamingWebSocketEndpoint, ""]]
-    HasStreamingDynamoDbTable: !Not [!Equals [!Ref StreamingDynamoDbTable, ""]]
+    DeployStreamingStack: !And [!Equals [!Ref AllowStreamingResponses, true], !Equals [!Ref StreamingWebSocketEndpoint, ""]]
     EnableUpload: !Equals [!Ref ShouldEnableUpload, true]
     NeedsVpc:  !And [ !Not [ !Equals [!Ref VpcSubnetId, ''] ], !Not [ !Equals [!Ref VpcSecurityGroupId, ''] ] ] 
 
@@ -816,7 +809,7 @@ Resources:
     
     StreamingSupport:
         Type: AWS::CloudFormation::Stack
-        Condition: EnableStreaming
+        Condition: DeployStreamingStack
         Properties:
             TimeoutInMinutes: 15
             TemplateURL: !Sub "https://${SourceBucket}.s3.${AWS::Region}.amazonaws.com/${SourcePrefix}/templates/streaming-support.yaml"
@@ -980,9 +973,9 @@ Resources:
                     - Name: ALLOW_STREAMING_RESPONSES
                       Value: !Ref AllowStreamingResponses
                     - Name: STREAMING_WEB_SOCKET_ENDPOINT
-                      Value: !If [EnableStreaming, !If [HasStreamingWebSocketEndpoint, !Ref StreamingWebSocketEndpoint, !Sub "wss://${StreamingSupport.Outputs.WebSocketId}.execute-api.${AWS::Region}.amazonaws.com/Prod"], ""]
+                      Value: !If [DeployStreamingStack, !Sub "wss://${StreamingSupport.Outputs.WebSocketId}.execute-api.${AWS::Region}.amazonaws.com/Prod", !If [EnableStreaming, !Ref StreamingWebSocketEndpoint, ""]]
                     - Name: STREAMING_DYNAMO_TABLE
-                      Value: !If [EnableStreaming, !If [HasStreamingDynamoDbTable, !Ref StreamingDynamoDbTable, !Sub "${StreamingSupport.Outputs.DynamoTableName}"], ""]
+                      Value: !If [DeployStreamingStack, !Sub "${StreamingSupport.Outputs.DynamoTableName}", ""]
                     - Name: ENABLE_UPLOAD
                       Value: !Ref ShouldEnableUpload
                     - Name: UPLOAD_BUCKET_NAME
@@ -1074,7 +1067,6 @@ Resources:
             BotAvatarImgUrl: !Ref BotAvatarImgUrl
             AllowStreamingResponses: !Ref AllowStreamingResponses
             StreamingWebSocketEndpoint: !Ref StreamingWebSocketEndpoint
-            StreamingDynamoDbTable: !Ref StreamingDynamoDbTable
             ShouldEnableUpload: !Ref ShouldEnableUpload
             UploadBucket: !Ref UploadBucket
             Timestamp: !Ref Timestamp
diff --git a/templates/master.yaml b/templates/master.yaml
index dd0fcd83..9551aa54 100644
--- a/templates/master.yaml
+++ b/templates/master.yaml
@@ -574,13 +574,7 @@ Parameters:
         Type: String
         Default: ''
         Description: >
-            If not using the Web socket endpoint created via AllowStreamingResponses, you can optionally specify an alternative streaming endpoint here
-            
-    StreamingDynamoDbTable:
-        Type: String
-        Default: ''
-        Description: >
-            If not using the DynamoDB table through Lex Web UI, you can optionally specify an alternative table for streaming purposes here
+            If you have an existing WebSocket API Gateway endpoint, you can specify it using this parameter. This requires parameter AllowStreamingResponses set to True.
     
     ShouldEnableUpload:
         Type: String
@@ -675,7 +669,6 @@ Metadata:
                   - retryCountPostTextTimeout
                   - AllowStreamingResponses
                   - StreamingWebSocketEndpoint
-                  - StreamingDynamoDbTable
                   - ShouldEnableUpload
                   - UploadBucket
             - Label:
@@ -812,7 +805,6 @@ Resources:
                 SourceBucket: !Ref BootstrapBucket
                 SourcePrefix: !Ref BootstrapPrefix
                 StreamingWebSocketEndpoint: !Ref StreamingWebSocketEndpoint
-                StreamingDynamoDbTable: !Ref StreamingDynamoDbTable
                 SourceObject: !Sub "${BootstrapPrefix}/src-v0.21.6.zip"
                 CustomResourceCodeObject: !Sub "${BootstrapPrefix}/custom-resources-v0.21.6.zip"
                 InitiateChatLambdaCodeObject: !Sub "${BootstrapPrefix}/initiate-chat-lambda-v0.21.6.zip"