Chat customization support (#4642)

* q streaming service sdk udpate * telemetry addMessage will send customizationArn * gamma endpoint * fix sending empty string instead null to the service if no customization * internal scope * revert dev config change * patch * move selectCustomization action under Q:Other Feature instead of inline suggestion * typo resourceArn -> customizationArn * Revert "typo resourceArn -> customizationArn" This reverts commit eaf014b. * patch * compile error due to stale code changes * detekt * detekt * remove stale comment
aws · Jul 5, 2024 · 937e9b4 · 937e9b4
1 parent 5862224
commit 937e9b4
Show file tree

Hide file tree

Showing 8 changed files with 36 additions and 17 deletions.
diff --git a/...community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/model/Requests.kt b/...community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/model/Requests.kt
@@ -5,6 +5,7 @@ package software.aws.toolkits.jetbrains.services.cwc.clients.chat.model
 
 import com.fasterxml.jackson.annotation.JsonProperty
 import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent
+import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.cwc.editor.context.ActiveFileContext
 
 enum class TriggerType {
@@ -19,7 +20,8 @@ data class ChatRequestData(
     val message: String,
     val activeFileContext: ActiveFileContext,
     val userIntent: UserIntent?,
-    val triggerType: TriggerType
+    val triggerType: TriggerType,
+    val customization: CodeWhispererCustomization?
 )
 
 interface CodeNames {

diff --git a/...mmunity/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt b/...mmunity/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt
@@ -211,6 +211,7 @@ class ChatSessionV1(
             .conversationId(conversationId)
             .currentMessage(ChatMessage.fromUserInputMessage(userInput))
             .chatTriggerType(ChatTriggerType.MANUAL)
+            .customizationArn(customization?.arn)
             .build()
         return GenerateAssistantResponseRequest.builder()
             .conversationState(conversationState)

diff --git a/...s-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt b/...s-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt
@@ -28,6 +28,7 @@ import kotlinx.coroutines.flow.mapNotNull
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.job
 import kotlinx.coroutines.withContext
+import migration.software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
@@ -356,6 +357,7 @@ class ChatController private constructor(
             activeFileContext = activeFileContext,
             userIntent = userIntent,
             triggerType = triggerType,
+            customization = CodeWhispererModelConfigurator.getInstance().activeCustomization(context.project)
         )
 
         val sessionInfo = getSessionInfo(tabId)

diff --git a/...software/aws/toolkits/jetbrains/services/cwc/controller/chat/telemetry/TelemetryHelper.kt b/...software/aws/toolkits/jetbrains/services/cwc/controller/chat/telemetry/TelemetryHelper.kt
@@ -104,7 +104,8 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
             cwsprChatRequestLength = data.message.length,
             cwsprChatResponseLength = responseLength,
             cwsprChatConversationType = CwsprChatConversationType.Chat,
-            credentialStartUrl = getStartUrl(context.project)
+            credentialStartUrl = getStartUrl(context.project),
+            codewhispererCustomizationArn = data.customization?.arn
         )
 
         val programmingLanguage = data.activeFileContext.fileContext?.fileLanguage

diff --git a/...e/aws/toolkits/jetbrains/services/codewhisperer/explorer/QStatusBarLoggedInActionGroup.kt b/...e/aws/toolkits/jetbrains/services/codewhisperer/explorer/QStatusBarLoggedInActionGroup.kt
@@ -59,7 +59,7 @@ class QStatusBarLoggedInActionGroup : DefaultActionGroup() {
 
             add(Separator.create())
             add(Separator.create(message("codewhisperer.statusbar.sub_menu.other_features.title")))
-            addAll(buildActionListForOtherFeatures(actionProvider))
+            addAll(buildActionListForOtherFeatures(it, actionProvider))
 
             add(Separator.create())
             add(Separator.create(message("codewhisperer.statusbar.sub_menu.connect_help.title")))

diff --git a/.../software/aws/toolkits/jetbrains/services/codewhisperer/explorer/actions/ActionFactory.kt b/.../software/aws/toolkits/jetbrains/services/codewhisperer/explorer/actions/ActionFactory.kt
@@ -4,8 +4,10 @@
 package software.aws.toolkits.jetbrains.services.codewhisperer.explorer.actions
 
 import com.intellij.openapi.project.Project
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
+import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
+import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
-import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isUserBuilderId
@@ -27,9 +29,9 @@ interface ActionProvider<T> {
     val documentation: T
 }
 
-fun<T> buildActionListForInlineSuggestions(project: Project, actionProvider: ActionProvider<T>): List<T> {
+@SuppressWarnings("UnusedParameter")
+fun <T> buildActionListForInlineSuggestions(project: Project, actionProvider: ActionProvider<T>): List<T> {
     val manager = CodeWhispererExplorerActionManager.getInstance()
-    val activeConnectionType = manager.checkActiveCodeWhispererConnectionType(project)
 
     return buildList {
         if (manager.isAutoEnabled()) {
@@ -40,18 +42,11 @@ fun<T> buildActionListForInlineSuggestions(project: Project, actionProvider: Act
 
         add(actionProvider.openCodeReference)
 
-        // We only show this customization node to SSO users who are in CodeWhisperer Gated Preview list
-        if (activeConnectionType == CodeWhispererLoginType.SSO &&
-            CodeWhispererModelConfigurator.getInstance().shouldDisplayCustomNode(project)
-        ) {
-            add(actionProvider.customize)
-        }
-
         add(actionProvider.learn)
     }
 }
 
-fun<T> buildActionListForCodeScan(project: Project, actionProvider: ActionProvider<T>): List<T> =
+fun <T> buildActionListForCodeScan(project: Project, actionProvider: ActionProvider<T>): List<T> =
     buildList {
         val codeScanManager = CodeWhispererCodeScanManager.getInstance(project)
         val manager = CodeWhispererExplorerActionManager.getInstance()
@@ -69,14 +64,21 @@ fun<T> buildActionListForCodeScan(project: Project, actionProvider: ActionProvid
         }
     }
 
-fun<T> buildActionListForOtherFeatures(actionProvider: ActionProvider<T>): List<T> =
+fun <T> buildActionListForOtherFeatures(project: Project, actionProvider: ActionProvider<T>): List<T> =
     buildList {
         if (!isRunningOnRemoteBackend()) {
             add(actionProvider.openChatPanel)
         }
+
+        val isIdC = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()).let { conn ->
+            conn != null && !conn.isSono()
+        }
+        if (isIdC && CodeWhispererModelConfigurator.getInstance().shouldDisplayCustomNode(project)) {
+            add(actionProvider.customize)
+        }
     }
 
-fun<T> buildActionListForConnectHelp(actionProvider: ActionProvider<T>): List<T> =
+fun <T> buildActionListForConnectHelp(actionProvider: ActionProvider<T>): List<T> =
     buildList {
         add(actionProvider.sendFeedback)
         add(actionProvider.connectOnGithub)

diff --git a/plugins/core/jetbrains-community/resources/telemetryOverride.json b/plugins/core/jetbrains-community/resources/telemetryOverride.json
@@ -461,6 +461,10 @@
                 {
                     "type": "credentialStartUrl",
                     "required": false
+                },
+                {
+                    "type": "codewhispererCustomizationArn",
+                    "required": false
                 }
             ]
         },

diff --git a/plugins/core/sdk-codegen/codegen-resources/codewhispererstreaming/service-2.json b/plugins/core/sdk-codegen/codegen-resources/codewhispererstreaming/service-2.json
@@ -214,7 +214,8 @@
                 "conversationId":{"shape":"ConversationId"},
                 "history":{"shape":"ChatHistory"},
                 "currentMessage":{"shape":"ChatMessage"},
-                "chatTriggerType":{"shape":"ChatTriggerType"}
+                "chatTriggerType":{"shape":"ChatTriggerType"},
+                "customizationArn":{"shape": "ResourceArn"}
             }
         },
         "CursorState":{
@@ -479,6 +480,12 @@
             "max":10,
             "min":0
         },
+        "ResourceArn":{
+            "type":"string",
+            "max":1224,
+            "min":0,
+            "pattern":"arn:([-.a-z0-9]{1,63}:){2}([-.a-z0-9]{0,63}:){2}([a-zA-Z0-9-_:/]){1,1023}"
+        },
         "ResourceNotFoundException":{
             "type":"structure",
             "required":["message"],