CodeWhisperer: Correct completion type for every completion in user decision events (#3811)

1. Previously we just use the completion type of the first suggestion,
    as the completion type for all suggestions. Now since this is no longer true,
    plugin side also needs to make the change to compute the completion type
    for every suggestion. Since this is per-suggestion level this will
    be recorded in the userDecision event.

    2. Note that we still have 'completionType' field in other events(
    serviceInvocation, clientComponentLatency, etc) these values no longer
    make sense for analysis.

    3. For actually computing the completionType for each suggestion, the rule
    is:
    If a suggestion has multi-lines and they are non-blank lines, it's a
    block suggestion, otherwise it's a line suggestion.

Author: andrewyuq

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorManager.kt

@@ -59,7 +59,8 @@ class CodeWhispererEditorManager {
                     PsiDocumentManager.getInstance(project).getPsiFile(document)?.virtualFile,
                     rangeMarker,
                     remainingRecommendation,
-                    selectedIndex
+                    selectedIndex,
+                    detail.completionType
                 )
 
                 ApplicationManager.getApplication().messageBus.syncPublisher(

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt

@@ -16,6 +16,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispe
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
+import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import software.aws.toolkits.telemetry.CodewhispererTriggerType
 import software.aws.toolkits.telemetry.Result
 import java.util.concurrent.TimeUnit
@@ -63,7 +64,8 @@ data class DetailContext(
     val reformatted: Completion,
     val isDiscarded: Boolean,
     val isTruncatedOnRight: Boolean,
-    val rightOverlap: String = ""
+    val rightOverlap: String = "",
+    val completionType: CodewhispererCompletionType,
 )
 
 data class SessionContext(

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt

@@ -16,6 +16,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.Reference
 import software.amazon.awssdk.services.codewhispererruntime.model.Span
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationChunk
+import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getCompletionType
 import kotlin.math.max
 
 class CodeWhispererRecommendationManager {
@@ -121,7 +122,15 @@ class CodeWhispererRecommendationManager {
         return recommendations.map {
             val isDiscardedByUserInput = !it.content().startsWith(userInput) || it.content() == userInput
             if (isDiscardedByUserInput) {
-                return@map DetailContext(requestId, it, it, isDiscarded = true, isTruncatedOnRight = false, rightOverlap = "")
+                return@map DetailContext(
+                    requestId,
+                    it,
+                    it,
+                    isDiscarded = true,
+                    isTruncatedOnRight = false,
+                    rightOverlap = "",
+                    getCompletionType(it)
+                )
             }
 
             val overlap = findRightContextOverlap(requestContext, it)
@@ -137,7 +146,15 @@ class CodeWhispererRecommendationManager {
                 .build()
             val isDiscardedByUserInputForTruncated = !truncated.content().startsWith(userInput) || truncated.content() == userInput
             if (isDiscardedByUserInputForTruncated) {
-                return@map DetailContext(requestId, it, truncated, isDiscarded = true, isTruncatedOnRight = true, rightOverlap = overlap)
+                return@map DetailContext(
+                    requestId,
+                    it,
+                    truncated,
+                    isDiscarded = true,
+                    isTruncatedOnRight = true,
+                    rightOverlap = overlap,
+                    getCompletionType(it)
+                )
             }
 
             val reformatted = reformat(requestContext, truncated)
@@ -148,7 +165,8 @@ class CodeWhispererRecommendationManager {
                 reformatted,
                 isDiscardedByRightContextTruncationDedupe,
                 truncated.content().length != it.content().length,
-                overlap
+                overlap,
+                getCompletionType(it)
             )
         }
     }

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt

@@ -65,8 +65,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhis
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CaretMovement
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.SUPPLEMENTAL_CONTEXT_TIMEOUT
-import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.checkCompletionType
-import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.checkEmptyRecommendations
+import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getCompletionType
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyErrorCodeWhispererUsageLimit
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider
@@ -180,9 +179,7 @@ class CodeWhispererService {
                     if (response.nextToken().isEmpty()) {
                         requestContext.latencyContext.paginationAllCompletionsEnd = System.nanoTime()
                     }
-                    val emptyRecommendations = checkEmptyRecommendations(response.completions())
-                    val completionType = checkCompletionType(response.completions(), emptyRecommendations)
-                    val responseContext = ResponseContext(sessionId, completionType)
+                    val responseContext = ResponseContext(sessionId)
                     logServiceInvocation(requestId, requestContext, responseContext, response.completions(), latency, null)
                     lastRecommendationIndex += response.completions().size
                     ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_CODE_COMPLETION_PERFORMED)
@@ -269,7 +266,7 @@ class CodeWhispererService {
                     }
                 }
                 val exceptionType = e::class.simpleName
-                val responseContext = ResponseContext(sessionId, CodewhispererCompletionType.Unknown)
+                val responseContext = ResponseContext(sessionId)
                 CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId)
                 logServiceInvocation(requestId, requestContext, responseContext, emptyList(), null, exceptionType)
                 CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent(
@@ -373,10 +370,17 @@ class CodeWhispererService {
         if (nextStates.recommendationContext.details.isEmpty() && response.nextToken().isEmpty()) {
             LOG.debug { "Received just an empty list from this session, requestId: $requestId" }
             CodeWhispererTelemetryService.getInstance().sendUserDecisionEvent(
-                requestId,
                 requestContext,
                 responseContext,
-                Completion.builder().build(),
+                DetailContext(
+                    requestId,
+                    Completion.builder().build(),
+                    Completion.builder().build(),
+                    false,
+                    false,
+                    "",
+                    CodewhispererCompletionType.Unknown
+                ),
                 -1,
                 CodewhispererSuggestionState.Empty,
                 nextStates.recommendationContext.details.size
@@ -488,7 +492,9 @@ class CodeWhispererService {
         responseContext: ResponseContext,
         recommendations: List<Completion>
     ) {
-        val detailContexts = recommendations.map { DetailContext("", it, it, true, false) }
+        val detailContexts = recommendations.map {
+            DetailContext("", it, it, true, false, "", getCompletionType(it))
+        }
         val recommendationContext = RecommendationContext(detailContexts, "", "", VisualPosition(0, 0))
 
         CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll(
@@ -719,7 +725,6 @@ data class RequestContext(
 
 data class ResponseContext(
     val sessionId: String,
-    val completionType: CodewhispererCompletionType
 )
 
 interface CodeWhispererCodeCompletionServiceListener {

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt

@@ -11,10 +11,10 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import org.apache.commons.collections4.queue.CircularFifoQueue
 import org.jetbrains.annotations.TestOnly
-import software.amazon.awssdk.services.codewhispererruntime.model.Completion
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.CodeScanTelemetryEvent
+import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
@@ -28,6 +28,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseCo
 import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getConnectionStartUrl
 import software.aws.toolkits.jetbrains.settings.AwsSettings
+import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import software.aws.toolkits.telemetry.CodewhispererLanguage
 import software.aws.toolkits.telemetry.CodewhispererPreviousSuggestionState
 import software.aws.toolkits.telemetry.CodewhispererSuggestionState
@@ -104,7 +105,7 @@ class CodeWhispererTelemetryService {
         CodewhispererTelemetry.serviceInvocation(
             project = project,
             codewhispererAutomatedTriggerType = automatedTriggerType.telemetryType,
-            codewhispererCompletionType = responseContext.completionType,
+            codewhispererCompletionType = CodewhispererCompletionType.Line,
             codewhispererCursorOffset = offset,
             codewhispererLanguage = codewhispererLanguage,
             codewhispererLastSuggestionIndex = lastRecommendationIndex,
@@ -126,14 +127,15 @@ class CodeWhispererTelemetryService {
     }
 
     fun sendUserDecisionEvent(
-        requestId: String,
         requestContext: RequestContext,
         responseContext: ResponseContext,
-        detail: Completion,
+        detailContext: DetailContext,
         index: Int,
         suggestionState: CodewhispererSuggestionState,
         numOfRecommendations: Int
     ) {
+        val requestId = detailContext.requestId
+        val recommendation = detailContext.recommendation
         val (project, _, triggerTypeInfo) = requestContext
         val codewhispererLanguage = requestContext.fileContextInfo.programmingLanguage.toTelemetryType()
         val supplementalContext = requestContext.supplementalContext
@@ -143,21 +145,21 @@ class CodeWhispererTelemetryService {
                 "Index: $index, " +
                 "State: $suggestionState, " +
                 "Request ID: $requestId, " +
-                "Recommendation: ${detail.content()}"
+                "Recommendation: ${recommendation.content()}"
         }
         val startUrl = getConnectionStartUrl(requestContext.connection)
         val importEnabled = CodeWhispererSettings.getInstance().isImportAdderEnabled()
         CodewhispererTelemetry.userDecision(
             project = project,
-            codewhispererCompletionType = responseContext.completionType,
+            codewhispererCompletionType = detailContext.completionType,
             codewhispererLanguage = codewhispererLanguage,
             codewhispererPaginationProgress = numOfRecommendations,
             codewhispererRequestId = requestId,
             codewhispererSessionId = responseContext.sessionId,
             codewhispererSuggestionIndex = index,
-            codewhispererSuggestionReferenceCount = detail.references().size,
-            codewhispererSuggestionReferences = jacksonObjectMapper().writeValueAsString(detail.references().map { it.licenseName() }.toSet().toList()),
-            codewhispererSuggestionImportCount = if (importEnabled) detail.mostRelevantMissingImports().size else null,
+            codewhispererSuggestionReferenceCount = recommendation.references().size,
+            codewhispererSuggestionReferences = jacksonObjectMapper().writeValueAsString(recommendation.references().map { it.licenseName() }.toSet().toList()),
+            codewhispererSuggestionImportCount = if (importEnabled) recommendation.mostRelevantMissingImports().size else null,
             codewhispererSuggestionState = suggestionState,
             codewhispererTriggerType = triggerTypeInfo.triggerType,
             credentialStartUrl = startUrl,
@@ -204,6 +206,12 @@ class CodeWhispererTelemetryService {
         }
 
         val supplementalContext = requestContext.supplementalContext
+        val completionType = if (recommendationContext.details.any {
+                it.completionType == CodewhispererCompletionType.Block
+            }
+        ) {
+            CodewhispererCompletionType.Block
+        } else CodewhispererCompletionType.Line
 
         CodewhispererTelemetry.userTriggerDecision(
             project = requestContext.project,
@@ -214,7 +222,7 @@ class CodeWhispererTelemetryService {
             codewhispererPartialAcceptanceCount = null,
             codewhispererCharactersAccepted = null,
             codewhispererCharactersRecommended = null,
-            codewhispererCompletionType = responseContext.completionType,
+            codewhispererCompletionType = completionType,
             codewhispererLanguage = language.toTelemetryType(),
             codewhispererTriggerType = requestContext.triggerTypeInfo.triggerType,
             codewhispererAutomatedTriggerType = automatedTriggerType.telemetryType,
@@ -286,13 +294,13 @@ class CodeWhispererTelemetryService {
         vFile: VirtualFile?,
         range: RangeMarker,
         suggestion: String,
-        selectedIndex: Int
+        selectedIndex: Int,
+        completionType: CodewhispererCompletionType
     ) {
-        val (sessionId, completionType) = responseContext
         val codewhispererLanguage = requestContext.fileContextInfo.programmingLanguage.toTelemetryType()
         CodeWhispererUserModificationTracker.getInstance(requestContext.project).enqueue(
             AcceptedSuggestionEntry(
-                time, vFile, range, suggestion, sessionId, requestId, selectedIndex,
+                time, vFile, range, suggestion, responseContext.sessionId, requestId, selectedIndex,
                 requestContext.triggerTypeInfo.triggerType, completionType,
                 codewhispererLanguage, null, null,
                 requestContext.connection
@@ -312,16 +320,15 @@ class CodeWhispererTelemetryService {
         val decisions = mutableListOf<CodewhispererSuggestionState>()
 
         detailContexts.forEachIndexed { index, detailContext ->
-            val (requestId, detail, _, isDiscarded) = detailContext
             val suggestionState = recordSuggestionState(
                 index,
                 sessionContext.selectedIndex,
                 sessionContext.seen.contains(index),
                 hasUserAccepted,
-                isDiscarded,
-                detail.content().isEmpty()
+                detailContext.isDiscarded,
+                detailContext.recommendation.content().isEmpty()
             )
-            sendUserDecisionEvent(requestId, requestContext, responseContext, detail, index, suggestionState, detailContexts.size)
+            sendUserDecisionEvent(requestContext, responseContext, detailContext, index, suggestionState, detailContexts.size)
 
             decisions.add(suggestionState)
         }
@@ -381,7 +388,7 @@ class CodeWhispererTelemetryService {
         val startUrl = getConnectionStartUrl(requestContext.connection)
         CodewhispererTelemetry.perceivedLatency(
             project = project,
-            codewhispererCompletionType = responseContext.completionType,
+            codewhispererCompletionType = CodewhispererCompletionType.Line,
             codewhispererLanguage = codewhispererLanguage,
             codewhispererRequestId = requestId,
             codewhispererSessionId = responseContext.sessionId,
@@ -409,7 +416,7 @@ class CodeWhispererTelemetryService {
             codewhispererPostprocessingLatency = requestContext.latencyContext.getCodeWhispererPostprocessingLatency(),
             codewhispererCredentialFetchingLatency = requestContext.latencyContext.getCodeWhispererCredentialFetchingLatency(),
             codewhispererTriggerType = requestContext.triggerTypeInfo.triggerType,
-            codewhispererCompletionType = responseContext.completionType,
+            codewhispererCompletionType = CodewhispererCompletionType.Line,
             codewhispererLanguage = codewhispererLanguage,
             credentialStartUrl = startUrl,
             codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt

@@ -109,24 +109,17 @@ fun VirtualFile.toCodeChunk(path: String): Sequence<Chunk> = sequence {
 }
 
 object CodeWhispererUtil {
-    fun checkCompletionType(
-        results: List<Completion>,
-        noRecommendation: Boolean
-    ): CodewhispererCompletionType {
-        if (noRecommendation) {
-            return CodewhispererCompletionType.Unknown
-        }
-        return if (results[0].content().contains("\n")) {
-            CodewhispererCompletionType.Block
-        } else {
-            CodewhispererCompletionType.Line
+    fun getCompletionType(completion: Completion): CodewhispererCompletionType {
+        val content = completion.content()
+        val nonBlankLines = content.split("\n").count { it.isNotBlank() }
+
+        return when {
+            content.isEmpty() -> CodewhispererCompletionType.Unknown
+            nonBlankLines > 1 -> CodewhispererCompletionType.Block
+            else -> CodewhispererCompletionType.Line
         }
     }
 
-    // return true if every recommendation is empty
-    fun checkEmptyRecommendations(recommendations: List<Completion>): Boolean =
-        recommendations.all { it.content().isEmpty() }
-
     fun notifyWarnCodeWhispererUsageLimit(project: Project? = null) {
         notifyWarn(
             message("codewhisperer.notification.usage_limit.warn.title"),

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt

@@ -156,9 +156,19 @@ internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCov
             null,
             mock()
         )
-        val responseContext = ResponseContext("sessionId", CodewhispererCompletionType.Block)
+        val responseContext = ResponseContext("sessionId")
         val recommendationContext = RecommendationContext(
-            listOf(DetailContext("requestId", pythonResponse.completions()[0], pythonResponse.completions()[0], false, false)),
+            listOf(
+                DetailContext(
+                    "requestId",
+                    pythonResponse.completions()[0],
+                    pythonResponse.completions()[0],
+                    false,
+                    false,
+                    "",
+                    CodewhispererCompletionType.Block
+                )
+            ),
             "x, y",
             "x, y",
             mock()

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererStateTest.kt

@@ -10,7 +10,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestU
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
-import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import software.aws.toolkits.telemetry.CodewhispererLanguage
 import software.aws.toolkits.telemetry.CodewhispererTriggerType
 
@@ -52,7 +51,6 @@ class CodeWhispererStateTest : CodeWhispererTestBase() {
             assertThat(listOf(actualResponseContext.sessionId)).isEqualTo(
                 pythonResponse.sdkHttpResponse().headers()[CodeWhispererService.KET_SESSION_ID]
             )
-            assertThat(actualResponseContext.completionType).isEqualTo(CodewhispererCompletionType.Block)
         }
     }