refactor: consolidate code completion request building; remove old code

Author: carlrobertoh

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt

@@ -71,20 +71,19 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
                 CompletionProgressNotifier.update(project, true)
 
                 var eventListener = CodeCompletionEventListener(request.editor, this)
+                val infillRequest = InfillRequestUtil.buildInfillRequest(request)
 
                 if (service<ModelSelectionService>().getServiceForFeature(FeatureType.CODE_COMPLETION) == ServiceType.PROXYAI) {
-                    project.service<GrpcClientService>().cancelCodeCompletion()
-                    project.service<GrpcClientService>()
-                        .getCodeCompletionAsync(eventListener, request, this)
+                    val grpcClient = project.service<GrpcClientService>()
+                    grpcClient.cancelCodeCompletion()
+                    grpcClient.getCodeCompletionAsync(infillRequest, eventListener, this)
                     return@channelFlow
                 }
 
-                val infillRequest = InfillRequestUtil.buildInfillRequest(request)
                 val call = service<CodeCompletionService>().getCodeCompletionAsync(
                     infillRequest,
                     CodeCompletionEventListener(request.editor, this)
                 )
-
                 currentCallRef.set(call)
             } finally {
                 awaitClose { currentCallRef.getAndSet(null)?.cancel() }
@@ -131,13 +130,8 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
             return false
         }
 
-        if (event is LookupInlineCompletionEvent || event is InlineCompletionEvent.DirectCall) {
-            return true
-        }
-
-        val hasActiveCompletion =
-            REMAINING_CODE_COMPLETION.get(event.toRequest()?.editor)?.partialCompletion?.isNotEmpty() == true
-
-        return event is InlineCompletionEvent.DocumentChange || hasActiveCompletion
+        return event is LookupInlineCompletionEvent
+                || event is InlineCompletionEvent.DirectCall
+                || event is InlineCompletionEvent.DocumentChange
     }
 }

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt

@@ -1,7 +1,8 @@
 package ee.carlrobert.codegpt.codecompletions
 
+import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.service
-import com.intellij.openapi.editor.Document
+import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import ee.carlrobert.codegpt.EncodingManager
@@ -15,7 +16,9 @@ class InfillRequest private constructor(
     val prefix: String,
     val suffix: String,
     val caretOffset: Int,
+    val editor: Editor?,
     val fileDetails: FileDetails?,
+    val gitDiff: String?,
     val repositoryName: String?,
     val dependenciesStructure: Set<ClassStructure>?,
     val context: InfillContext?,
@@ -29,7 +32,8 @@ class InfillRequest private constructor(
         private val suffix: String
         private val caretOffset: Int
         private var fileDetails: FileDetails? = null
-        private var additionalContext: String? = null
+        private var editor: Editor? = null
+        private var gitDiff: String? = null
         private var repositoryName: String? = null
         private var dependenciesStructure: Set<ClassStructure>? = null
         private var context: InfillContext? = null
@@ -46,10 +50,9 @@ class InfillRequest private constructor(
             this.stopTokens = getStopTokens()
         }
 
-        constructor(
-            document: Document,
-            caretOffset: Int,
-        ) {
+        constructor(editor: Editor) {
+            val document = editor.document
+            val caretOffset = runReadAction { editor.caretModel.offset }
             prefix =
                 document.getText(TextRange(0, caretOffset))
                     .truncateText(MAX_PROMPT_TOKENS, false)
@@ -58,11 +61,12 @@ class InfillRequest private constructor(
                     .truncateText(MAX_PROMPT_TOKENS)
             this.caretOffset = caretOffset
             this.stopTokens = getStopTokens()
+            this.editor = editor
+            this.fileDetails = FileDetails(editor.document.text, editor.virtualFile.path)
         }
 
-        fun fileDetails(fileDetails: FileDetails) = apply { this.fileDetails = fileDetails }
-        fun additionalContext(additionalContext: String) =
-            apply { this.additionalContext = additionalContext }
+        fun gitDiff(gitDiff: String) =
+            apply { this.gitDiff = gitDiff }
 
         fun addRepositoryName(repositoryName: String) =
             apply { this.repositoryName = repositoryName }
@@ -90,16 +94,18 @@ class InfillRequest private constructor(
         }
 
         fun build(): InfillRequest {
-            val modifiedPrefix = if (!additionalContext.isNullOrEmpty()) {
-                "/*\n${additionalContext}\n*/\n\n$prefix"
+            val modifiedPrefix = if (!gitDiff.isNullOrEmpty()) {
+                "/*\n${gitDiff}\n*/\n\n$prefix"
             } else {
                 prefix
             }
             return InfillRequest(
                 modifiedPrefix,
                 suffix,
                 caretOffset,
+                editor,
                 fileDetails,
+                gitDiff,
                 repositoryName,
                 dependenciesStructure,
                 context,

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt

@@ -1,80 +1,34 @@
 package ee.carlrobert.codegpt.codecompletions
 
 import com.intellij.codeInsight.inline.completion.InlineCompletionRequest
-import com.intellij.openapi.application.readAction
 import com.intellij.openapi.components.service
-import com.intellij.refactoring.suggested.startOffset
-import ee.carlrobert.codegpt.EncodingManager
-import ee.carlrobert.codegpt.codecompletions.psi.CompletionContextService
-import ee.carlrobert.codegpt.codecompletions.psi.readText
 import ee.carlrobert.codegpt.psistructure.PsiStructureProvider
 import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings
 import ee.carlrobert.codegpt.util.GitUtil
 
-
 object InfillRequestUtil {
 
     suspend fun buildInfillRequest(request: InlineCompletionRequest): InfillRequest {
-        val caretOffset = readAction { request.editor.caretModel.offset }
-        val infillRequestBuilder = InfillRequest.Builder(request.document, caretOffset)
-            .fileDetails(
-                InfillRequest.FileDetails(
-                    request.document.text,
-                    request.file.virtualFile.path
-                )
-            )
+        val infillRequestBuilder = InfillRequest.Builder(request.editor)
 
         val project = request.editor.project ?: return infillRequestBuilder.build()
         if (service<ConfigurationSettings>().state.codeCompletionSettings.gitDiffEnabled) {
-            val additionalContext = GitUtil.getCurrentChanges(project)
-            if (!additionalContext.isNullOrEmpty()) {
-                infillRequestBuilder.additionalContext(additionalContext)
-            }
-        }
-
-        if (service<ConfigurationSettings>().state.codeCompletionSettings.contextAwareEnabled) {
-            getInfillContext(request, caretOffset)?.let {
-                infillRequestBuilder.context(it)
-                infillRequestBuilder.addRepositoryName(it.getRepoName())
+            GitUtil.getCurrentChanges(project)?.let { diff ->
+                if (diff.isNotEmpty()) {
+                    infillRequestBuilder.gitDiff(diff)
+                }
             }
         }
 
         if (service<ConfigurationSettings>().state.codeCompletionSettings.collectDependencyStructure) {
-            val depth = service<ConfigurationSettings>().state.codeCompletionSettings.psiStructureAnalyzeDepth
+            val depth =
+                service<ConfigurationSettings>().state.codeCompletionSettings.psiStructureAnalyzeDepth
             val psiStructure = PsiStructureProvider().get(listOf(request.file), depth)
             if (psiStructure.isNotEmpty()) {
                 infillRequestBuilder.addDependenciesStructure(psiStructure)
-                infillRequestBuilder.addRepositoryName(psiStructure.first().repositoryName)
             }
         }
 
         return infillRequestBuilder.build()
     }
-
-    private fun getInfillContext(
-        request: InlineCompletionRequest,
-        caretOffset: Int
-    ): InfillContext? {
-        val infillContext =
-            service<CompletionContextService>().findContext(request.editor, caretOffset)
-                ?: return null
-        val caretInEnclosingElement =
-            caretOffset - infillContext.enclosingElement.psiElement.startOffset
-        val entireText = infillContext.enclosingElement.psiElement.readText()
-        val prefix = entireText.take(caretInEnclosingElement)
-        val suffix =
-            if (entireText.length < caretInEnclosingElement) "" else entireText.takeLast(
-                entireText.length - caretInEnclosingElement
-            )
-        return truncateContext(prefix + suffix, infillContext)
-    }
-
-    private fun truncateContext(prompt: String, infillContext: InfillContext): InfillContext {
-        var promptTokens = EncodingManager.getInstance().countTokens(prompt)
-        val truncatedContextElements = infillContext.contextElements.takeWhile {
-            promptTokens += it.tokens
-            promptTokens <= MAX_PROMPT_TOKENS
-        }.toSet()
-        return InfillContext(infillContext.enclosingElement, truncatedContextElements)
-    }
 }

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt

@@ -1,16 +1,15 @@
 package ee.carlrobert.codegpt.codecompletions.edit
 
-import com.intellij.codeInsight.inline.completion.InlineCompletionRequest
 import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
 import com.intellij.openapi.Disposable
-import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.util.net.ssl.CertificateManager
 import ee.carlrobert.codegpt.CodeGPTPlugin
 import ee.carlrobert.codegpt.codecompletions.CodeCompletionEventListener
+import ee.carlrobert.codegpt.codecompletions.InfillRequest
 import ee.carlrobert.codegpt.credentials.CredentialsStore
 import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CodeGptApiKey
 import ee.carlrobert.codegpt.settings.service.FeatureType
@@ -48,15 +47,16 @@ class GrpcClientService(private val project: Project) : Disposable {
     }
 
     fun getCodeCompletionAsync(
+        request: InfillRequest,
         eventListener: CodeCompletionEventListener,
-        request: InlineCompletionRequest,
         channel: ProducerScope<InlineCompletionElement>
     ) {
         ensureCodeCompletionConnection()
 
+        val editor = request.editor ?: return
         val grpcRequest = createCodeCompletionGrpcRequest(request)
         codeCompletionObserver =
-            CodeCompletionStreamObserver(request.editor, channel, eventListener)
+            CodeCompletionStreamObserver(editor, channel, eventListener)
         codeCompletionContext?.cancel(null)
         val ctx = Context.current().withCancellation()
         codeCompletionContext = ctx
@@ -186,16 +186,17 @@ class GrpcClientService(private val project: Project) : Disposable {
         }
     }
 
-    private fun createCodeCompletionGrpcRequest(request: InlineCompletionRequest): GrpcCodeCompletionRequest {
-        val editor = request.editor
+    private fun createCodeCompletionGrpcRequest(request: InfillRequest): GrpcCodeCompletionRequest {
+        val fileDetails = request.fileDetails ?: throw IllegalArgumentException("File details are required")
+        val gitDiff = request.gitDiff ?: ""
         return GrpcCodeCompletionRequest.newBuilder()
             .setModel(
                 ModelSelectionService.getInstance().getModelForFeature(FeatureType.CODE_COMPLETION)
             )
-            .setFilePath(editor.virtualFile.path)
-            .setFileContent(editor.document.text)
-            .setGitDiff(GitUtil.getCurrentChanges(project) ?: "")
-            .setCursorPosition(runReadAction { editor.caretModel.offset })
+            .setFilePath(fileDetails.filePath)
+            .setFileContent(fileDetails.fileContent)
+            .setGitDiff(gitDiff)
+            .setCursorPosition(request.caretOffset)
             .setPluginVersion(CodeGPTPlugin.getVersion())
             .build()
     }

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/NextEditStreamObserver.kt

@@ -11,6 +11,7 @@ import ee.carlrobert.service.NextEditResponse
 import io.grpc.Status
 import io.grpc.StatusRuntimeException
 import io.grpc.stub.StreamObserver
+import ee.carlrobert.codegpt.codecompletions.CodeCompletionFormatter
 
 class NextEditStreamObserver(
     private val editor: Editor,

src/main/kotlin/ee/carlrobert/codegpt/nextedit/NextEditDiffViewer.kt

@@ -83,7 +83,7 @@ class NextEditDiffViewer(
 
         val size = computeCompactSize(change)
         myEditor.component.preferredSize = size
-        adjustPopupSize(popup, myEditor)
+        adjustPopupSize(popup, myEditor, mainEditor)
 
         val changeOffset = change.lineFragment.startOffset1
         val adjustedLocation = getAdjustedPopupLocation(mainEditor, changeOffset, size)
@@ -221,7 +221,7 @@ class NextEditDiffViewer(
 
             coroutineScope.launch {
                 withContext(Dispatchers.EDT) {
-                    adjustPopupSize(popup, myEditor)
+                    adjustPopupSize(popup, myEditor, mainEditor)
                     val adjustedLocation = getAdjustedPopupLocation(
                         mainEditor,
                         change.lineFragment.startOffset1,
@@ -342,31 +342,60 @@ fun getAdjustedPopupLocation(editor: Editor, changeOffset: Int, popupSize: Dimen
         SwingUtilities.convertPointToScreen(point, editor.component)
         return point
     }
+
     val visibleArea = editor.scrollingModel.visibleArea
     val editorLocationOnScreen = editor.component.locationOnScreen
-    val vsOffset = editor.scrollingModel.verticalScrollOffset
-    val marginX = +(editor as EditorEx).gutterComponentEx.size.width
+    val verticalOffset = editor.scrollingModel.verticalScrollOffset
+    val gutterWidth = (editor as? EditorEx)?.gutterComponentEx?.size?.width ?: 0
+    val margin = JBUI.scale(8)
+
     val yInEditor = pointInEditor.y
+    val visibleTopY = visibleArea.y
+    val visibleBottomY = visibleArea.y + visibleArea.height
+    val changeIsVisible = yInEditor in visibleTopY until visibleBottomY
+
+    if (changeIsVisible) {
+        val spaceBelow = visibleBottomY - yInEditor - editor.lineHeight
+        val fitsBelow = spaceBelow >= popupSize.height
+        if (fitsBelow) {
+            val rawLeft = editorLocationOnScreen.x + pointInEditor.x + gutterWidth
+            val minLeft = editorLocationOnScreen.x + visibleArea.x + margin
+            val maxLeft = editorLocationOnScreen.x + visibleArea.x + visibleArea.width - popupSize.width - margin
+            val left = rawLeft.coerceIn(minLeft, maxLeft)
+            val top = editorLocationOnScreen.y + yInEditor - verticalOffset
+            return Point(left, top)
+        }
 
-    val spaceBelow = visibleArea.y + visibleArea.height - yInEditor - editor.lineHeight
-    val fitsBelow = spaceBelow >= popupSize.height
-    if (fitsBelow) {
-        val top = editorLocationOnScreen.y + yInEditor - vsOffset
-        val left = editorLocationOnScreen.x + pointInEditor.x + marginX
+        val screenChangeY = editorLocationOnScreen.y + yInEditor - verticalOffset
+        val spaceAbove = screenChangeY - editorLocationOnScreen.y
+        val canPlaceAbove = spaceAbove >= popupSize.height + margin
+        val top = if (canPlaceAbove) {
+            screenChangeY - popupSize.height - margin
+        } else {
+            val centered = screenChangeY - popupSize.height / 2
+            val minTop = editorLocationOnScreen.y + margin
+            val maxTop = editorLocationOnScreen.y + visibleArea.height - popupSize.height - margin
+            centered.coerceIn(minTop, maxTop)
+        }
+        val left = editorLocationOnScreen.x + visibleArea.x + visibleArea.width - popupSize.width - gutterWidth
         return Point(left, top)
     }
 
-    val left =
-        editorLocationOnScreen.x + visibleArea.x + visibleArea.width - popupSize.width - marginX
-    val visibleBottom = editorLocationOnScreen.y + visibleArea.height
-    var top = editorLocationOnScreen.y + yInEditor - vsOffset - popupSize.height / 2
-    top = top.coerceIn(editorLocationOnScreen.y, visibleBottom - popupSize.height)
+    val topAligned = editorLocationOnScreen.y + margin
+    val bottomAligned = editorLocationOnScreen.y + visibleArea.height - popupSize.height - margin
+    val top = if (yInEditor < visibleTopY) topAligned else bottomAligned
+    val left = editorLocationOnScreen.x + visibleArea.x + visibleArea.width - popupSize.width - gutterWidth
     return Point(left, top)
 }
 
-fun adjustPopupSize(popup: JBPopup, editor: Editor) {
-    val newWidth = editor.component.preferredSize.width
-    val newHeight = editor.component.preferredSize.height
+fun adjustPopupSize(popup: JBPopup, contentEditor: Editor, anchorEditor: Editor) {
+    val preferred = contentEditor.component.preferredSize
+    val visibleArea = anchorEditor.scrollingModel.visibleArea
+    val margin = JBUI.scale(16)
+    val maxWidth = (visibleArea.width - margin).coerceAtLeast(JBUI.scale(64))
+    val maxHeight = (visibleArea.height - margin).coerceAtLeast(JBUI.scale(48))
+    val newWidth = preferred.width.coerceAtMost(maxWidth)
+    val newHeight = preferred.height.coerceAtMost(maxHeight)
     popup.size = Dimension(newWidth, newHeight)
     popup.content.revalidate()
     popup.content.repaint()

src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt

@@ -153,17 +153,4 @@ object GitUtil {
                     !it.startsWith("commit ")
         }
     }
-
-    private fun String.cleanDiff(showContext: Boolean = false): String =
-        lineSequence()
-            .filterNot { line ->
-                line.startsWith("index ") ||
-                        line.startsWith("diff --git") ||
-                        line.startsWith("---") ||
-                        line.startsWith("+++") ||
-                        line.startsWith("===") ||
-                        line.contains("\\ No newline at end of file")
-                (!showContext && line.startsWith(" "))
-            }
-            .joinToString("\n")
 }