Skip to content

Commit 05eb0a7

Browse files
Will-ShaoHuaWill-ShaoHua
committed
enhance inline completion supplemental context fetching
1 parent effbebd commit 05eb0a7

File tree

6 files changed

+98
-39
lines changed

6 files changed

+98
-39
lines changed

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt‎

Lines changed: 75 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ import software.aws.toolkits.core.utils.debug
2525
import software.aws.toolkits.core.utils.getLogger
2626
import software.aws.toolkits.core.utils.info
2727
import software.aws.toolkits.core.utils.warn
28-
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
2928
import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextController
3029
import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil
3130
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
@@ -147,12 +146,18 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
147146
val latency = System.currentTimeMillis() - startFetchingTimestamp
148147
if (it.contents.isNotEmpty()) {
149148
val logStr = buildString {
150-
append("Successfully fetched supplemental context with strategy ${it.strategy} with $latency ms")
149+
append("""Q inline completion supplemental context:
150+
| Strategy: ${it.strategy},
151+
| Latency: $latency ms,
152+
| Contents: ${it.contents.size} chunks,
153+
| ContentLength: ${it.contentLength} chars,
154+
| TargetFile: ${it.targetFileName},
155+
""".trimMargin())
151156
it.contents.forEachIndexed { index, chunk ->
152157
append(
153158
"""
154159
|
155-
| Chunk ${index + 1}:
160+
| Chunk ${index}:
156161
| path = ${chunk.path},
157162
| score = ${chunk.score},
158163
| contentLength = ${chunk.content.length}
@@ -219,23 +224,19 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
219224
val query = generateQuery(targetContext)
220225

221226
val contexts = withContext(coroutineContext) {
222-
val projectContextDeferred1 = if (CodeWhispererFeatureConfigService.getInstance().getInlineCompletion()) {
223-
async {
224-
val t0 = System.currentTimeMillis()
225-
val r = fetchProjectContext(query, psiFile, targetContext)
226-
val t1 = System.currentTimeMillis()
227-
LOG.debug {
228-
buildString {
229-
append("time elapse for fetching project context=${t1 - t0}ms; ")
230-
append("numberOfChunks=${r.contents.size}; ")
231-
append("totalLength=${r.contentLength}")
232-
}
227+
val projectContextDeferred1 = async {
228+
val t0 = System.currentTimeMillis()
229+
val r = fetchProjectContext(query, psiFile, targetContext)
230+
val t1 = System.currentTimeMillis()
231+
LOG.debug {
232+
buildString {
233+
append("time elapse for fetching project context=${t1 - t0}ms; ")
234+
append("numberOfChunks=${r.contents.size}; ")
235+
append("totalLength=${r.contentLength}")
233236
}
234-
235-
r
236237
}
237-
} else {
238-
null
238+
239+
r
239240
}
240241

241242
val openTabsContextDeferred1 = async {
@@ -253,20 +254,65 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
253254
r
254255
}
255256

256-
if (projectContextDeferred1 != null) {
257-
awaitAll(projectContextDeferred1, openTabsContextDeferred1)
258-
} else {
259-
awaitAll(openTabsContextDeferred1)
260-
}
257+
awaitAll(projectContextDeferred1, openTabsContextDeferred1)
261258
}
262259

263-
val projectContext = contexts.find { it.strategy == CrossFileStrategy.ProjectContext }
260+
val projectContext = contexts.find { it.strategy == CrossFileStrategy.Codemap }
264261
val openTabsContext = contexts.find { it.strategy == CrossFileStrategy.OpenTabsBM25 }
265262

266-
return if (projectContext != null && projectContext.contents.isNotEmpty()) {
267-
projectContext
268-
} else {
269-
openTabsContext ?: SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
263+
return when {
264+
projectContext == null && openTabsContext == null -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
265+
266+
projectContext != null && openTabsContext != null -> {
267+
val context1 = projectContext.contents
268+
val context2 = openTabsContext.contents
269+
val mergedContext = (context1 + context2).filter { it.content.isNotEmpty() }
270+
271+
val strategy = if (projectContext.contentLength != 0 && openTabsContext.contentLength != 0) {
272+
CrossFileStrategy.Codemap
273+
} else if (projectContext.contentLength != 0) {
274+
CrossFileStrategy.Codemap
275+
} else if (openTabsContext.contentLength != 0) {
276+
CrossFileStrategy.OpenTabsBM25
277+
} else {
278+
CrossFileStrategy.Empty
279+
}
280+
281+
SupplementalContextInfo(
282+
isUtg = false,
283+
contents = mergedContext,
284+
targetFileName = targetContext.filename,
285+
strategy = strategy
286+
)
287+
}
288+
289+
projectContext != null -> {
290+
return if (projectContext.contentLength == 0) {
291+
SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
292+
} else {
293+
SupplementalContextInfo(
294+
isUtg = false,
295+
contents = projectContext.contents,
296+
targetFileName = targetContext.filename,
297+
strategy = CrossFileStrategy.Codemap
298+
)
299+
}
300+
}
301+
302+
openTabsContext != null -> {
303+
return if (openTabsContext.contentLength == 0) {
304+
SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
305+
} else {
306+
SupplementalContextInfo(
307+
isUtg = false,
308+
contents = openTabsContext.contents,
309+
targetFileName = targetContext.filename,
310+
strategy = CrossFileStrategy.OpenTabsBM25
311+
)
312+
}
313+
}
314+
315+
else -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
270316
}
271317
}
272318

@@ -285,7 +331,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
285331
)
286332
},
287333
targetFileName = targetContext.filename,
288-
strategy = CrossFileStrategy.ProjectContext
334+
strategy = CrossFileStrategy.Codemap
289335
)
290336
}
291337

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/SupplementalContextStrategy.kt‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,21 +12,23 @@ enum class UtgStrategy : SupplementalContextStrategy {
1212
;
1313

1414
override fun toString() = when (this) {
15-
ByName -> "ByName"
16-
ByContent -> "ByContent"
17-
Empty -> "Empty"
15+
ByName -> "byname"
16+
ByContent -> "bycontent"
17+
Empty -> "empty"
1818
}
1919
}
2020

2121
enum class CrossFileStrategy : SupplementalContextStrategy {
2222
OpenTabsBM25,
2323
Empty,
2424
ProjectContext,
25+
Codemap
2526
;
2627

2728
override fun toString() = when (this) {
28-
OpenTabsBM25 -> "OpenTabs_BM25"
29-
Empty -> "Empty"
30-
ProjectContext -> "ProjectContext"
29+
OpenTabsBM25 -> "opentabs"
30+
Empty -> "empty"
31+
ProjectContext -> "projectcontext"
32+
Codemap -> "codemap"
3133
}
3234
}

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,7 @@ class CodeWhispererFileContextProviderTest {
208208
assertThat(providerContext.constructed()).hasSize(1)
209209
assertThat(serverContext.constructed()).hasSize(1)
210210

211-
whenever(providerContext.constructed()[0].queryInline(any(), any())).thenThrow(RuntimeException("mock exception"))
211+
whenever(providerContext.constructed()[0].queryInline(any(), any(), any())).thenThrow(RuntimeException("mock exception"))
212212

213213
val result = controller.queryInline("query", "filePath")
214214
assertThat(result).isEmpty()

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/LspMessage.kt‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ data class QueryChatRequest(
5252
data class QueryInlineCompletionRequest(
5353
val query: String,
5454
val filePath: String,
55+
val target: String
5556
) : LspRequest
5657

5758
data class LspResponse(
@@ -70,6 +71,16 @@ enum class IndexOption(val command: String) {
7071
DEFAULT("default"),
7172
}
7273

74+
enum class InlineContextTarget(private val v: String) {
75+
DEFAULT("default"),
76+
CODEMAP("codemap"),
77+
BM25("bm25"), ;
78+
79+
override fun toString(): String {
80+
return this.v
81+
}
82+
}
83+
7384
// TODO: unify with [software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk]
7485
data class InlineBm25Chunk(
7586
val content: String,

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextController.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ class ProjectContextController(private val project: Project, private val cs: Cor
5858

5959
suspend fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> =
6060
try {
61-
projectContextProvider.queryInline(query, filePath)
61+
projectContextProvider.queryInline(query, filePath, InlineContextTarget.CODEMAP)
6262
} catch (e: Exception) {
6363
var logStr = "error while querying inline for project context $e.message"
6464
if (e is TimeoutCancellationException || e is TimeoutException) {

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,9 +171,9 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
171171
}.await()
172172
}
173173

174-
suspend fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> = withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
174+
suspend fun queryInline(query: String, filePath: String, target: InlineContextTarget): List<InlineBm25Chunk> = withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
175175
cs.async {
176-
val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath))
176+
val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath, target.toString()))
177177
val r = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted)
178178
return@async mapper.readValue<List<InlineBm25Chunk>>(r.responseBody)
179179
}.await()

0 commit comments

Comments
 (0)