Skip to content

Commit 5dd4c0c

Browse files
evanliu048evanliu048
committed
delete prefetch method
1 parent 7e2e46a commit 5dd4c0c

File tree

1 file changed

+111
-118
lines changed
  • plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service

1 file changed

+111
-118
lines changed

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt‎

Lines changed: 111 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import com.intellij.util.messages.Topic
2525
import kotlinx.coroutines.CoroutineScope
2626
import kotlinx.coroutines.Deferred
2727
import kotlinx.coroutines.Job
28+
import kotlinx.coroutines.SupervisorJob
2829
import kotlinx.coroutines.async
2930
import kotlinx.coroutines.delay
3031
import kotlinx.coroutines.isActive
@@ -209,7 +210,110 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
209210
invokeCodeWhispererInBackground(requestContext)
210211
}
211212

212-
internal suspend fun invokeCodeWhispererInBackground(requestContext: RequestContext): Job {
213+
internal suspend fun invokeCodeWhispererInBackground(
214+
requestContext: RequestContext,
215+
isPrefetch: Boolean = false,
216+
currStates: InvocationContext? = null,
217+
): Job {
218+
if (isPrefetch && currStates?.recommendationContext?.details?.isNotEmpty() == true) {
219+
val firstValidRecommendation = currStates.recommendationContext.details
220+
.firstOrNull {
221+
!it.isDiscarded && it.recommendation.content().isNotEmpty()
222+
} ?: return SupervisorJob().apply { complete() }
223+
val job = cs.launch(getCoroutineBgContext()) {
224+
val latencyContext = LatencyContext().apply {
225+
codewhispererPreprocessingStart = System.nanoTime()
226+
codewhispererEndToEndStart = System.nanoTime()
227+
}
228+
229+
val nextCaretPosition = calculateNextCaretPosition(requestContext, firstValidRecommendation)
230+
val nextFileContextInfo = createNextFileContextInfo(requestContext, firstValidRecommendation)
231+
232+
val nextRequestContext = requestContext.copy(
233+
caretPosition = nextCaretPosition,
234+
fileContextInfo = nextFileContextInfo,
235+
latencyContext = latencyContext
236+
)
237+
val newVisualPosition = withContext(EDT) {
238+
runReadAction {
239+
nextRequestContext.editor.offsetToVisualPosition(nextRequestContext.caretPosition.offset)
240+
}
241+
}
242+
try {
243+
val nextResponse = CodeWhispererClientAdaptor
244+
.getInstance(nextRequestContext.project)
245+
.generateCompletions(
246+
buildCodeWhispererRequest(
247+
nextRequestContext.fileContextInfo,
248+
nextRequestContext.awaitSupplementalContext(),
249+
nextRequestContext.customizationArn
250+
)
251+
)
252+
val startTime = System.nanoTime()
253+
nextRequestContext.latencyContext.codewhispererPreprocessingEnd = System.nanoTime()
254+
nextRequestContext.latencyContext.paginationAllCompletionsStart = System.nanoTime()
255+
CodeWhispererInvocationStatus.getInstance().setInvocationStart()
256+
nextResponse.let {
257+
val endTime = System.nanoTime()
258+
val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble()
259+
val requestId = nextResponse.responseMetadata().requestId()
260+
val sessionId = nextResponse.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0]
261+
262+
nextRequestContext.latencyContext.apply {
263+
codewhispererPostprocessingStart = System.nanoTime()
264+
paginationFirstCompletionTime = (endTime - codewhispererEndToEndStart).toDouble()
265+
firstRequestId = requestId
266+
}
267+
268+
CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId)
269+
270+
val nextResponseContext = ResponseContext(sessionId)
271+
CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent(
272+
nextResponse.responseMetadata().requestId(),
273+
nextRequestContext,
274+
nextResponseContext,
275+
nextResponse.completions().size,
276+
true,
277+
latency,
278+
null
279+
)
280+
val validatedResponse = validateResponse(it)
281+
val detailContexts = withContext(EDT) {
282+
runReadAction {
283+
CodeWhispererRecommendationManager.getInstance().buildDetailContext(
284+
nextRequestContext,
285+
"",
286+
validatedResponse.completions(),
287+
validatedResponse.responseMetadata().requestId()
288+
)
289+
}
290+
}
291+
val nextRecommendationContext = RecommendationContext(detailContexts, "", "", newVisualPosition)
292+
val newPopup = withContext(EDT) {
293+
JBPopupFactory.getInstance().createMessage("Dummy popup")
294+
}
295+
296+
// send userDecision and trigger decision when next recommendation haven't been seen
297+
if (currStates.popup.isDisposed) {
298+
CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll(
299+
nextRequestContext,
300+
nextResponseContext,
301+
nextRecommendationContext,
302+
SessionContext(),
303+
false
304+
)
305+
} else {
306+
nextInvocationContext = InvocationContext(nextRequestContext, nextResponseContext, nextRecommendationContext, newPopup)
307+
}
308+
LOG.debug { "Prefetched next invocation stored in nextInvocationContext" }
309+
}
310+
} catch (ex: Exception) {
311+
LOG.warn { "Failed to prefetch next codewhisperer invocation: ${ex.message}" }
312+
}
313+
}
314+
return job
315+
}
316+
213317
val popup = withContext(EDT) {
214318
CodeWhispererPopupManager.getInstance().initPopup().also {
215319
Disposer.register(it) { CodeWhispererInvocationStatus.getInstance().finishInvocation() }
@@ -481,7 +585,6 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
481585
if (currStates == null) {
482586
// first response
483587
nextStates = initStates(requestContext, responseContext, response, caretMovement, popup)
484-
nextStates?.let { prefetchNextInvocationAsync(it) }
485588
isPopupShowing = false
486589

487590
// receiving a null state means caret has moved backward or there's a conflict with
@@ -491,6 +594,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
491594
CodeWhispererPopupManager.getInstance().cancelPopup(popup)
492595
return null
493596
}
597+
cs.launch(getCoroutineBgContext()) {
598+
invokeCodeWhispererInBackground(requestContext, true, nextStates)
599+
}
494600
} else {
495601
// subsequent responses
496602
nextStates = updateStates(currStates, response)
@@ -621,121 +727,6 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
621727
CodeWhispererPopupManager.getInstance().changeStates(states, 0, "", true, recommendationAdded)
622728
}
623729

624-
private fun prefetchNextInvocationAsync(currStates: InvocationContext) {
625-
val firstValidRecommendation = currStates.recommendationContext.details
626-
.firstOrNull { !it.isDiscarded && it.recommendation.content().isNotEmpty() }
627-
?: return
628-
629-
val latencyContext = LatencyContext().apply {
630-
codewhispererPreprocessingStart = System.nanoTime()
631-
codewhispererEndToEndStart = System.nanoTime()
632-
}
633-
cs.launch(getCoroutineBgContext()) {
634-
val psiFile =
635-
runReadAction {
636-
PsiDocumentManager.getInstance(currStates.requestContext.project).getPsiFile(currStates.requestContext.editor.document)
637-
} ?: run {
638-
LOG.debug { "No PSI file for the current document" }
639-
return@launch
640-
}
641-
val nextCaretPosition = calculateNextCaretPosition(currStates.requestContext, firstValidRecommendation)
642-
val nextFileContextInfo = createNextFileContextInfo(currStates.requestContext, firstValidRecommendation)
643-
644-
val nextRequestContext = try {
645-
getRequestContext(
646-
currStates.requestContext.triggerTypeInfo,
647-
currStates.requestContext.editor,
648-
currStates.requestContext.project,
649-
psiFile,
650-
latencyContext
651-
).copy(
652-
caretPosition = nextCaretPosition,
653-
fileContextInfo = nextFileContextInfo
654-
)
655-
} catch (e: Exception) {
656-
LOG.debug { e.message.toString() }
657-
CodeWhispererTelemetryService.getInstance().sendFailedServiceInvocationEvent(currStates.requestContext.project, e::class.simpleName)
658-
return@launch
659-
}
660-
661-
val newVisualPosition = withContext(EDT) {
662-
runReadAction {
663-
nextRequestContext.editor.offsetToVisualPosition(nextRequestContext.caretPosition.offset)
664-
}
665-
}
666-
try {
667-
val nextResponse = CodeWhispererClientAdaptor.getInstance(nextRequestContext.project)
668-
.generateCompletions(
669-
buildCodeWhispererRequest(
670-
nextRequestContext.fileContextInfo,
671-
nextRequestContext.awaitSupplementalContext(),
672-
nextRequestContext.customizationArn
673-
)
674-
)
675-
val startTime = System.nanoTime()
676-
nextRequestContext.latencyContext.codewhispererPreprocessingEnd = System.nanoTime()
677-
nextRequestContext.latencyContext.paginationAllCompletionsStart = System.nanoTime()
678-
CodeWhispererInvocationStatus.getInstance().setInvocationStart()
679-
nextResponse.let {
680-
val endTime = System.nanoTime()
681-
val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble()
682-
val requestId = nextResponse.responseMetadata().requestId()
683-
val sessionId = nextResponse.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0]
684-
685-
nextRequestContext.latencyContext.apply {
686-
codewhispererPostprocessingStart = System.nanoTime()
687-
paginationFirstCompletionTime = (endTime - codewhispererEndToEndStart).toDouble()
688-
firstRequestId = requestId
689-
}
690-
691-
CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId)
692-
693-
val nextResponseContext = ResponseContext(sessionId)
694-
CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent(
695-
nextResponse.responseMetadata().requestId(),
696-
nextRequestContext,
697-
nextResponseContext,
698-
nextResponse.completions().size,
699-
true,
700-
latency,
701-
null
702-
)
703-
val validatedResponse = validateResponse(it)
704-
val detailContexts = withContext(EDT) {
705-
runReadAction {
706-
CodeWhispererRecommendationManager.getInstance().buildDetailContext(
707-
nextRequestContext,
708-
"",
709-
validatedResponse.completions(),
710-
validatedResponse.responseMetadata().requestId()
711-
)
712-
}
713-
}
714-
val nextRecommendationContext = RecommendationContext(detailContexts, "", "", newVisualPosition)
715-
val popup = withContext(EDT) {
716-
JBPopupFactory.getInstance().createMessage("dummy popup")
717-
}
718-
719-
// send userDecision and trigger decision when next recommendation haven't been seen
720-
if (currStates.popup.isDisposed) {
721-
CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll(
722-
nextRequestContext,
723-
nextResponseContext,
724-
nextRecommendationContext,
725-
SessionContext(),
726-
false
727-
)
728-
} else {
729-
nextInvocationContext = InvocationContext(nextRequestContext, nextResponseContext, nextRecommendationContext, popup)
730-
}
731-
LOG.debug { "Prefetched next invocation stored in nextInvocationContext" }
732-
}
733-
} catch (ex: Exception) {
734-
LOG.warn { "Failed to prefetch next codewhisperer invocation: ${ex.message}" }
735-
}
736-
}
737-
}
738-
739730
fun promoteNextInvocationIfAvailable() {
740731
val nextStates = nextInvocationContext ?: run {
741732
LOG.debug { "No nextInvocationContext found, nothing to promote." }
@@ -758,7 +749,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
758749
typeaheadAdded = true,
759750
recommendationAdded = false
760751
)
761-
prefetchNextInvocationAsync(updatedNextStates)
752+
cs.launch(getCoroutineBgContext()) {
753+
invokeCodeWhispererInBackground(updatedNextStates.requestContext, true, updatedNextStates)
754+
}
762755
}
763756

764757
LOG.debug { "Promoted nextInvocationContext to current session and displayed next recommendation." }

0 commit comments

Comments
 (0)