Skip to content

Commit 35da4fb

Browse files
Will-ShaoHuaWill-ShaoHua
authored
send supplemental context strategy telemetry field (#3852)
1 parent 7fc2b48 commit 35da4fb

File tree

11 files changed

+137
-38
lines changed

11 files changed

+137
-38
lines changed

‎gradle/libs.versions.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ kotlinCoroutines = "1.6.4"
2020
mockito = "4.6.1"
2121
mockitoKotlin = "4.0.0"
2222
mockk = "1.13.4"
23-
telemetryGenerator = "1.0.140"
23+
telemetryGenerator = "1.0.143"
2424
testLogger = "3.1.0"
2525
testRetry = "1.5.2"
2626
slf4j = "1.7.36"

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

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import com.intellij.openapi.Disposable
77
import com.intellij.openapi.editor.VisualPosition
88
import com.intellij.openapi.editor.markup.RangeHighlighter
99
import com.intellij.openapi.ui.popup.JBPopup
10+
import com.intellij.openapi.vfs.VirtualFile
1011
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
1112
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
1213
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
@@ -16,6 +17,9 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispe
1617
import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext
1718
import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
1819
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
20+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CrossFileStrategy
21+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.SupplementalContextStrategy
22+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.UtgStrategy
1923
import software.aws.toolkits.telemetry.CodewhispererCompletionType
2024
import software.aws.toolkits.telemetry.CodewhispererTriggerType
2125
import software.aws.toolkits.telemetry.Result
@@ -28,6 +32,11 @@ data class Chunk(
2832
val score: Double = 0.0
2933
)
3034

35+
data class ListUtgCandidateResult(
36+
val vfile: VirtualFile?,
37+
val strategy: UtgStrategy
38+
)
39+
3140
data class CaretContext(val leftFileContext: String, val rightFileContext: String, val leftContextOnCurrentLine: String = "")
3241

3342
data class FileContextInfo(
@@ -39,8 +48,9 @@ data class FileContextInfo(
3948
data class SupplementalContextInfo(
4049
val isUtg: Boolean,
4150
val contents: List<Chunk>,
42-
val latency: Long,
43-
val targetFileName: String
51+
val targetFileName: String,
52+
val strategy: SupplementalContextStrategy,
53+
val latency: Long = 0L,
4454
) {
4555
val contentLength: Int
4656
get() = contents.fold(0) { acc, chunk ->
@@ -49,6 +59,24 @@ data class SupplementalContextInfo(
4959

5060
val isProcessTimeout: Boolean
5161
get() = latency > CodeWhispererConstants.SUPPLEMENTAL_CONTEXT_TIMEOUT
62+
63+
companion object {
64+
fun emptyCrossFileContextInfo(targetFileName: String): SupplementalContextInfo = SupplementalContextInfo(
65+
isUtg = false,
66+
contents = emptyList(),
67+
targetFileName = targetFileName,
68+
strategy = CrossFileStrategy.Empty,
69+
latency = 0L
70+
)
71+
72+
fun emptyUtgFileContextInfo(targetFileName: String): SupplementalContextInfo = SupplementalContextInfo(
73+
isUtg = true,
74+
contents = emptyList(),
75+
targetFileName = targetFileName,
76+
strategy = UtgStrategy.Empty,
77+
latency = 0L
78+
)
79+
}
5280
}
5381

5482
data class RecommendationContext(

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,9 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhisperer
6868
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getCompletionType
6969
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyErrorCodeWhispererUsageLimit
7070
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
71+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CrossFileStrategy
7172
import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider
73+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.UtgStrategy
7274
import software.aws.toolkits.jetbrains.utils.isInjectedText
7375
import software.aws.toolkits.resources.message
7476
import software.aws.toolkits.telemetry.CodewhispererCompletionType
@@ -534,7 +536,8 @@ class CodeWhispererService {
534536
isUtg = isTstFile,
535537
contents = emptyList(),
536538
latency = System.currentTimeMillis() - startFetchingTimestamp,
537-
targetFileName = fileContext.filename
539+
targetFileName = fileContext.filename,
540+
strategy = if (isTstFile) UtgStrategy.Empty else CrossFileStrategy.Empty
538541
)
539542
} else {
540543
LOG.debug { "Run into unexpected error when fetching supplemental context, error: ${e.message}" }

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,7 @@ class CodeWhispererTelemetryService {
275275
codewhispererSupplementalContextIsUtg = supplementalContext?.isUtg,
276276
codewhispererSupplementalContextLength = supplementalContext?.contentLength,
277277
codewhispererSupplementalContextTimeout = supplementalContext?.isProcessTimeout,
278+
codewhispererSupplementalContextStrategyId = supplementalContext?.strategy.toString(),
278279
codewhispererUserGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup().name
279280
)
280281
}

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

Lines changed: 38 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -114,10 +114,10 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
114114
val language = targetContext.programmingLanguage
115115
val group = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
116116

117-
val chunks = if (isTst) {
117+
val supplementalContext = if (isTst) {
118118
when (shouldFetchUtgContext(language, group)) {
119119
true -> extractSupplementalFileContextForTst(psiFile, targetContext)
120-
false -> emptyList()
120+
false -> SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename)
121121
null -> {
122122
LOG.debug { "UTG is not supporting ${targetContext.programmingLanguage.languageId}" }
123123
null
@@ -126,18 +126,18 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
126126
} else {
127127
when (shouldFetchCrossfileContext(language, group)) {
128128
true -> extractSupplementalFileContextForSrc(psiFile, targetContext)
129-
false -> emptyList()
129+
false -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
130130
null -> {
131131
LOG.debug { "Crossfile is not supporting ${targetContext.programmingLanguage.languageId}" }
132132
null
133133
}
134134
}
135135
}
136136

137-
return chunks?.let {
138-
if (it.isNotEmpty()) {
137+
return supplementalContext?.let {
138+
if (it.contents.isNotEmpty()) {
139139
LOG.info { "Successfully fetched supplemental context." }
140-
it.forEachIndexed { index, chunk ->
140+
it.contents.forEachIndexed { index, chunk ->
141141
LOG.info {
142142
"""
143143
|---------------------------------------------------------------
@@ -153,12 +153,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
153153
LOG.warn { "Failed to fetch supplemental context, empty list." }
154154
}
155155

156-
SupplementalContextInfo(
157-
isUtg = isTst,
158-
contents = it,
159-
latency = System.currentTimeMillis() - startFetchingTimestamp,
160-
targetFileName = targetContext.filename
161-
)
156+
it.copy(latency = System.currentTimeMillis() - startFetchingTimestamp)
162157
}
163158
}
164159

@@ -192,8 +187,10 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
192187
override fun isTestFile(psiFile: PsiFile) = psiFile.programmingLanguage().fileCrawler.isTestFile(psiFile.virtualFile, psiFile.project)
193188

194189
@VisibleForTesting
195-
suspend fun extractSupplementalFileContextForSrc(psiFile: PsiFile, targetContext: FileContextInfo): List<Chunk> {
196-
if (!targetContext.programmingLanguage.isSupplementalContextSupported()) return emptyList()
190+
suspend fun extractSupplementalFileContextForSrc(psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo {
191+
if (!targetContext.programmingLanguage.isSupplementalContextSupported()) {
192+
return SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
193+
}
197194

198195
// takeLast(11) will extract 10 lines (exclusing current line) of left context as the query parameter
199196
val query = targetContext.caretContext.leftFileContext.split("\n").takeLast(11).joinToString("\n")
@@ -212,7 +209,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
212209
"0 chunks was found for supplemental context, fileName=${targetContext.filename}, " +
213210
"programmingLanaugage: ${targetContext.programmingLanguage}"
214211
}
215-
return emptyList()
212+
return SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
216213
}
217214

218215
// we need to keep the reference to Chunk object because we will need to get "nextChunk" later after calculation
@@ -227,7 +224,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
227224
yield()
228225

229226
// we use nextChunk as supplemental context
230-
return top3Chunks.mapNotNull { bm25Result ->
227+
val crossfileContext = top3Chunks.mapNotNull { bm25Result ->
231228
contentToChunk[bm25Result.docString]?.let {
232229
if (it.nextChunk.isNotBlank()) {
233230
Chunk(content = it.nextChunk, path = it.path, score = bm25Result.score)
@@ -236,20 +233,31 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
236233
}
237234
}
238235
}
236+
237+
return SupplementalContextInfo(
238+
isUtg = false,
239+
contents = crossfileContext,
240+
targetFileName = targetContext.filename,
241+
strategy = CrossFileStrategy.OpenTabsBM25
242+
)
239243
}
240244

241245
@VisibleForTesting
242-
fun extractSupplementalFileContextForTst(psiFile: PsiFile, targetContext: FileContextInfo): List<Chunk> {
243-
if (!targetContext.programmingLanguage.isUTGSupported()) return emptyList()
246+
fun extractSupplementalFileContextForTst(psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo {
247+
if (!targetContext.programmingLanguage.isUTGSupported()) {
248+
return SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename)
249+
}
244250

245-
val focalFile = targetContext.programmingLanguage.fileCrawler.listUtgCandidate(psiFile)
251+
val utgCandidateResult = targetContext.programmingLanguage.fileCrawler.listUtgCandidate(psiFile)
252+
val focalFile = utgCandidateResult.vfile
253+
val strategy = utgCandidateResult.strategy
246254

247255
return focalFile?.let { file ->
248256
runReadAction {
249257
val relativePath = contentRootPathProvider.getPathToElement(project, file, null) ?: file.path
250258
val content = file.content()
251259

252-
if (content.isBlank()) {
260+
val utgContext = if (content.isBlank()) {
253261
emptyList()
254262
} else {
255263
listOf(
@@ -264,8 +272,17 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
264272
)
265273
)
266274
}
275+
276+
SupplementalContextInfo(
277+
isUtg = true,
278+
contents = utgContext,
279+
targetFileName = targetContext.filename,
280+
strategy = strategy
281+
)
267282
}
268-
}.orEmpty()
283+
} ?: run {
284+
return SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename)
285+
}
269286
}
270287

271288
companion object {

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import com.intellij.openapi.vfs.VirtualFile
1313
import com.intellij.psi.PsiFile
1414
import com.intellij.psi.PsiManager
1515
import software.aws.toolkits.core.utils.tryOrNull
16+
import software.aws.toolkits.jetbrains.services.codewhisperer.model.ListUtgCandidateResult
1617

1718
/**
1819
* An interface define how do we parse and fetch files provided a psi file or project
@@ -39,7 +40,7 @@ interface FileCrawler {
3940
* @param target psi of the test file we are searching with, e.g. MainTest.java
4041
* @return its source file e.g. Main.java, main.py or most relevant file if any
4142
*/
42-
fun listUtgCandidate(target: PsiFile): VirtualFile?
43+
fun listUtgCandidate(target: PsiFile): ListUtgCandidateResult
4344

4445
/**
4546
* List files opened in the editors and sorted by file distance @see [CodeWhispererFileCrawler.getFileDistance]
@@ -61,7 +62,7 @@ class NoOpFileCrawler : FileCrawler {
6162
override suspend fun listFilesImported(psiFile: PsiFile): List<VirtualFile> = emptyList()
6263

6364
override fun listFilesUnderProjectRoot(project: Project): List<VirtualFile> = emptyList()
64-
override fun listUtgCandidate(target: PsiFile): VirtualFile? = null
65+
override fun listUtgCandidate(target: PsiFile) = ListUtgCandidateResult(null, UtgStrategy.Empty)
6566

6667
override fun listFilesWithinSamePackage(psiFile: PsiFile): List<VirtualFile> = emptyList()
6768

@@ -129,7 +130,19 @@ abstract class CodeWhispererFileCrawler : FileCrawler {
129130
return fileToFileDistanceList.sortedBy { it.second }.map { it.first }
130131
}
131132

132-
override fun listUtgCandidate(target: PsiFile): VirtualFile? = findSourceFileByName(target) ?: findSourceFileByContent(target)
133+
override fun listUtgCandidate(target: PsiFile): ListUtgCandidateResult {
134+
val byName = findSourceFileByName(target)
135+
if (byName != null) {
136+
return ListUtgCandidateResult(byName, UtgStrategy.ByName)
137+
}
138+
139+
val byContent = findSourceFileByContent(target)
140+
if (byContent != null) {
141+
return ListUtgCandidateResult(byContent, UtgStrategy.ByContent)
142+
}
143+
144+
return ListUtgCandidateResult(null, UtgStrategy.Empty)
145+
}
133146

134147
abstract fun findSourceFileByName(target: PsiFile): VirtualFile?
135148

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.codewhisperer.util
5+
6+
interface SupplementalContextStrategy
7+
8+
enum class UtgStrategy : SupplementalContextStrategy {
9+
ByName,
10+
ByContent,
11+
Empty;
12+
13+
override fun toString() = when (this) {
14+
ByName -> "ByName"
15+
ByContent -> "ByContent"
16+
Empty -> "Empty"
17+
}
18+
}
19+
20+
enum class CrossFileStrategy : SupplementalContextStrategy {
21+
OpenTabsBM25,
22+
Empty;
23+
24+
override fun toString() = when (this) {
25+
OpenTabsBM25 -> "OpenTabs_BM25"
26+
Empty -> "Empty"
27+
}
28+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeCove
6161
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererCodeCoverageTracker
6262
import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager
6363
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_SECONDS_IN_MINUTE
64+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CrossFileStrategy
6465
import software.aws.toolkits.jetbrains.services.telemetry.NoOpPublisher
6566
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
6667
import software.aws.toolkits.jetbrains.settings.AwsSettings
@@ -153,7 +154,7 @@ internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCov
153154
mock(),
154155
mock(),
155156
FileContextInfo(mock(), pythonFileName, CodeWhispererPython.INSTANCE),
156-
SupplementalContextInfo(false, emptyList(), 0, ""),
157+
SupplementalContextInfo(isUtg = false, contents = emptyList(), targetFileName = "", strategy = CrossFileStrategy.Empty, latency = 0L),
157158
null,
158159
mock()
159160
)

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -184,12 +184,12 @@ class CodeWhispererFileContextProviderTest {
184184
runBlocking {
185185
var context = aFileContextInfo(CodeWhispererCsharp.INSTANCE)
186186

187-
assertThat(sut.extractSupplementalFileContextForSrc(psi, context)).isEmpty()
188-
assertThat(sut.extractSupplementalFileContextForTst(psi, context)).isEmpty()
187+
assertThat(sut.extractSupplementalFileContextForSrc(psi, context).contents).isEmpty()
188+
assertThat(sut.extractSupplementalFileContextForTst(psi, context).contents).isEmpty()
189189

190190
context = aFileContextInfo(CodeWhispererKotlin.INSTANCE)
191-
assertThat(sut.extractSupplementalFileContextForSrc(psi, context)).isEmpty()
192-
assertThat(sut.extractSupplementalFileContextForTst(psi, context)).isEmpty()
191+
assertThat(sut.extractSupplementalFileContextForSrc(psi, context).contents).isEmpty()
192+
assertThat(sut.extractSupplementalFileContextForTst(psi, context).contents).isEmpty()
193193
}
194194
}
195195

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

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.util.JavaCodeWhisp
2222
import software.aws.toolkits.jetbrains.services.codewhisperer.util.JavascriptCodeWhispererFileCrawler
2323
import software.aws.toolkits.jetbrains.services.codewhisperer.util.PythonCodeWhispererFileCrawler
2424
import software.aws.toolkits.jetbrains.services.codewhisperer.util.TypescriptCodeWhispererFileCrawler
25+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.UtgStrategy
2526
import software.aws.toolkits.jetbrains.services.codewhisperer.util.content
2627
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
2728
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
@@ -341,7 +342,8 @@ class JavaCodeWhispererFileCrawlerTest {
341342

342343
val actual = sut.listUtgCandidate(tstPsi)
343344

344-
assertThat(actual).isNotNull.isEqualTo(mainPsi.virtualFile)
345+
assertThat(actual.vfile).isNotNull.isEqualTo(mainPsi.virtualFile)
346+
assertThat(actual.strategy).isNotNull.isEqualTo(UtgStrategy.ByName)
345347
}
346348
}
347349

@@ -392,7 +394,8 @@ class JavaCodeWhispererFileCrawlerTest {
392394
val actual = sut.listUtgCandidate(tstPsi)
393395

394396
assertThat(openedFiles).isEqualTo(5)
395-
assertThat(actual).isNotNull.isEqualTo(mainPsi.virtualFile)
397+
assertThat(actual.vfile).isNotNull.isEqualTo(mainPsi.virtualFile)
398+
assertThat(actual.strategy).isNotNull.isEqualTo(UtgStrategy.ByContent)
396399
}
397400
}
398401

@@ -491,7 +494,8 @@ class PythonCodeWhispererFileCrawlerTest {
491494
runInEdtAndWait {
492495
fixture.openFileInEditor(tstPsi.virtualFile)
493496
val actual = sut.listUtgCandidate(tstPsi)
494-
assertThat(actual).isNotNull.isEqualTo(mainPsi.virtualFile)
497+
assertThat(actual.vfile).isNotNull.isEqualTo(mainPsi.virtualFile)
498+
assertThat(actual.strategy).isNotNull.isEqualTo(UtgStrategy.ByName)
495499
}
496500
}
497501

@@ -534,7 +538,8 @@ class PythonCodeWhispererFileCrawlerTest {
534538
val actual = sut.listUtgCandidate(tstPsi)
535539

536540
assertThat(openedFiles).isEqualTo(5)
537-
assertThat(actual).isNotNull.isEqualTo(mainPsi.virtualFile)
541+
assertThat(actual.vfile).isNotNull.isEqualTo(mainPsi.virtualFile)
542+
assertThat(actual.strategy).isNotNull.isEqualTo(UtgStrategy.ByContent)
538543
}
539544
}
540545

0 commit comments

Comments
 (0)