feat: code assistant (#810)

* feat: code assistant implementation * refactor: clean up
carlrobertoh · Dec 26, 2024 · 8d6ca73 · 8d6ca73
1 parent d816485
commit 8d6ca73
Show file tree

Hide file tree

Showing 32 changed files with 944 additions and 62 deletions.
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -12,7 +12,7 @@ jsoup = "1.17.2"
 jtokkit = "1.1.0"
 junit = "5.11.0"
 kotlin = "2.0.0"
-llm-client = "0.8.30"
+llm-client = "0.8.32"
 okio = "3.9.0"
 tree-sitter = "0.24.4"
 

diff --git a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java
@@ -1,10 +1,12 @@
 package ee.carlrobert.codegpt;
 
 import com.intellij.openapi.util.Key;
+import ee.carlrobert.codegpt.predictions.CodeSuggestionDiffViewer;
 import ee.carlrobert.codegpt.settings.prompts.PersonaDetails;
 import ee.carlrobert.codegpt.ui.DocumentationDetails;
 import ee.carlrobert.llm.client.codegpt.CodeGPTUserDetails;
 import java.util.List;
+import okhttp3.Call;
 
 public class CodeGPTKeys {
 
@@ -26,4 +28,8 @@ public class CodeGPTKeys {
       Key.create("codegpt.isFetchingCompletion");
   public static final Key<Boolean> IS_PROMPT_TEXT_FIELD_DOCUMENT =
       Key.create("codegpt.isPromptTextFieldDocument");
+  public static final Key<CodeSuggestionDiffViewer> EDITOR_PREDICTION_DIFF_VIEWER =
+      Key.create("codegpt.editorPredictionDiffViewer");
+  public static final Key<Call> PENDING_PREDICTION_CALL =
+      Key.create("codegpt.editorPendingPredictionCall");
 }
diff --git a/src/main/java/ee/carlrobert/codegpt/statusbar/CodeGPTStatusBarWidget.java b/src/main/java/ee/carlrobert/codegpt/statusbar/CodeGPTStatusBarWidget.java
@@ -1,6 +1,7 @@
 package ee.carlrobert.codegpt.statusbar;
 
 import static ee.carlrobert.codegpt.CodeGPTKeys.IS_FETCHING_COMPLETION;
+import static ee.carlrobert.codegpt.CodeGPTKeys.PENDING_PREDICTION_CALL;
 
 import com.intellij.openapi.actionSystem.ActionGroup;
 import com.intellij.openapi.actionSystem.ActionManager;
@@ -44,7 +45,9 @@ public CodeGPTStatusBarWidget(Project project) {
   protected @NotNull WidgetState getWidgetState(@Nullable VirtualFile file) {
     var state = new WidgetState(CodeGPTBundle.get("statusBar.widget.tooltip"), "", true);
     var fetchingCompletion = IS_FETCHING_COMPLETION.get(getEditor());
-    var loading = fetchingCompletion != null && fetchingCompletion;
+    var pendingPredicationCall = PENDING_PREDICTION_CALL.get(getEditor());
+    var loading =
+        (fetchingCompletion != null && fetchingCompletion) || pendingPredicationCall != null;
 
     state.setIcon(loading ? Icons.StatusBarCompletionInProgress : Icons.DefaultSmall);
     return state;

diff --git a/src/main/java/ee/carlrobert/codegpt/ui/UIUtil.java b/src/main/java/ee/carlrobert/codegpt/ui/UIUtil.java
@@ -139,8 +139,14 @@ public static <T extends JComponent> T withEmptyLeftBorder(T component) {
   }
 
   public static JLabel createComment(String messageKey) {
+    return createComment(messageKey, ComponentPanelBuilder.MAX_COMMENT_WIDTH);
+  }
+
+  public static JLabel createComment(String messageKey, int maxLineLength) {
     var comment = ComponentPanelBuilder.createCommentComponent(
-        CodeGPTBundle.get(messageKey), true);
+        CodeGPTBundle.get(messageKey),
+        true,
+        maxLineLength);
     comment.setBorder(JBUI.Borders.empty(0, 4));
     return comment;
   }

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/CodeGPTLookupListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/CodeGPTLookupListener.kt
@@ -0,0 +1,54 @@
+package ee.carlrobert.codegpt
+
+import com.intellij.codeInsight.lookup.Lookup
+import com.intellij.codeInsight.lookup.LookupEvent
+import com.intellij.codeInsight.lookup.LookupListener
+import com.intellij.codeInsight.lookup.LookupManagerListener
+import com.intellij.codeInsight.lookup.impl.LookupImpl
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.components.service
+import ee.carlrobert.codegpt.predictions.PredictionService
+import ee.carlrobert.codegpt.settings.GeneralSettings
+import ee.carlrobert.codegpt.settings.service.ServiceType
+import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
+
+class CodeGPTLookupListener : LookupManagerListener {
+    override fun activeLookupChanged(oldLookup: Lookup?, newLookup: Lookup?) {
+        if (newLookup is LookupImpl) {
+            newLookup.addLookupListener(object : LookupListener {
+
+                var beforeApply: String = ""
+                var cursorOffset: Int = 0
+
+                override fun beforeItemSelected(event: LookupEvent): Boolean {
+                    beforeApply = newLookup.editor.document.text
+                    cursorOffset = runReadAction {
+                        newLookup.editor.caretModel.offset
+                    }
+
+                    return true
+                }
+
+                override fun itemSelected(event: LookupEvent) {
+                    val editor = newLookup.editor
+
+                    if (GeneralSettings.getSelectedService() != ServiceType.CODEGPT
+                        || !service<CodeGPTServiceSettings>().state.codeAssistantEnabled
+                        || service<EncodingManager>().countTokens(editor.document.text) > 4098
+                    ) {
+                        return
+                    }
+
+                    ApplicationManager.getApplication().executeOnPooledThread {
+                        service<PredictionService>().displayLookupPrediction(
+                            editor,
+                            event,
+                            beforeApply
+                        )
+                    }
+                }
+            })
+        }
+    }
+}
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeAssistantFeatureToggleActions.kt.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeAssistantFeatureToggleActions.kt.kt
@@ -0,0 +1,35 @@
+package ee.carlrobert.codegpt.actions
+
+import com.intellij.openapi.actionSystem.ActionUpdateThread
+import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.components.service
+import com.intellij.openapi.project.DumbAwareAction
+import ee.carlrobert.codegpt.settings.GeneralSettings
+import ee.carlrobert.codegpt.settings.service.ServiceType.CODEGPT
+import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
+
+abstract class CodeAssistantFeatureToggleAction(
+    private val enableFeatureAction: Boolean
+) : DumbAwareAction() {
+
+    override fun actionPerformed(e: AnActionEvent) {
+        val settings = service<CodeGPTServiceSettings>().state
+        settings.codeAssistantEnabled = enableFeatureAction
+    }
+
+    override fun update(e: AnActionEvent) {
+        val codeAssistantEnabled = service<CodeGPTServiceSettings>().state.codeAssistantEnabled
+
+        e.presentation.isVisible = GeneralSettings.getSelectedService() == CODEGPT
+                    && codeAssistantEnabled != enableFeatureAction
+        e.presentation.isEnabled = GeneralSettings.getSelectedService() == CODEGPT
+    }
+
+    override fun getActionUpdateThread(): ActionUpdateThread {
+        return ActionUpdateThread.BGT
+    }
+}
+
+class EnableCodeAssistantAction : CodeAssistantFeatureToggleAction(true)
+
+class DisableCodeAssistantAction : CodeAssistantFeatureToggleAction(false)
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/InlayActionPromoter.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/InlayActionPromoter.kt
@@ -7,11 +7,14 @@ import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.actionSystem.DataContext
 import ee.carlrobert.codegpt.codecompletions.AcceptNextLineInlayAction
 import ee.carlrobert.codegpt.codecompletions.AcceptNextWordInlayAction
+import ee.carlrobert.codegpt.predictions.TriggerCustomPredictionAction
 
 class InlayActionPromoter : ActionPromoter {
     override fun promote(actions: List<AnAction>, context: DataContext): List<AnAction> {
         val editor = CommonDataKeys.EDITOR.getData(context) ?: return emptyList()
 
+        actions.filterIsInstance<TriggerCustomPredictionAction>().takeIf { it.isNotEmpty() }?.let { return it }
+
         if (InlineCompletionContext.getOrNull(editor) == null) {
             return emptyList()
         }

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt
@@ -8,13 +8,20 @@ import com.intellij.codeInsight.inline.completion.session.InlineCompletionContex
 import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession
 import com.intellij.codeInsight.lookup.LookupManager
 import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.components.service
 import com.intellij.openapi.editor.Caret
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.actionSystem.EditorAction
 import com.intellij.openapi.editor.actionSystem.EditorWriteActionHandler
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.util.concurrency.ThreadingAssertions
 import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
+import ee.carlrobert.codegpt.EncodingManager
+import ee.carlrobert.codegpt.predictions.PredictionService
+import ee.carlrobert.codegpt.settings.GeneralSettings
+import ee.carlrobert.codegpt.settings.service.ServiceType
+import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
 
 class CodeCompletionInsertAction :
     EditorAction(InsertInlineCompletionHandler()), HintManagerImpl.ActionToIgnore {
@@ -34,11 +41,27 @@ class CodeCompletionInsertAction :
                 val textToInsert = context.textToInsert()
                 val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) ?: ""
                 if (remainingCompletion.isNotEmpty()) {
-                    REMAINING_EDITOR_COMPLETION.set(editor, remainingCompletion.removePrefix(textToInsert))
+                    REMAINING_EDITOR_COMPLETION.set(
+                        editor,
+                        remainingCompletion.removePrefix(textToInsert)
+                    )
                 }
 
+                val beforeApply = editor.document.text
                 InlineCompletion.getHandlerOrNull(editor)?.insert()
-                return
+
+                if (GeneralSettings.getSelectedService() == ServiceType.CODEGPT
+                    && service<CodeGPTServiceSettings>().state.codeAssistantEnabled
+                    && service<EncodingManager>().countTokens(editor.document.text) <= 4098) {
+                    ApplicationManager.getApplication().executeOnPooledThread {
+                        service<PredictionService>().displayAutocompletePrediction(
+                            editor,
+                            textToInsert,
+                            beforeApply
+                        )
+                    }
+                    return
+                }
             }
 
             for (element in elements) {
@@ -92,7 +115,10 @@ class CodeCompletionInsertAction :
             processRemainingCompletion(remainingCompletionLine, editor, endOffset)
         }
 
-        private fun processPartialCompletionElement(element: CodeCompletionTextElement, editor: Editor) {
+        private fun processPartialCompletionElement(
+            element: CodeCompletionTextElement,
+            editor: Editor
+        ) {
             val lineNumber = editor.document.getLineNumber(editor.caretModel.offset)
             val lineEndOffset = editor.document.getLineEndOffset(lineNumber)
             editor.caretModel.moveToOffset(lineEndOffset)

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionProgressNotifier.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionProgressNotifier.kt
@@ -1,5 +1,6 @@
 package ee.carlrobert.codegpt.codecompletions
 
+import com.intellij.openapi.project.Project
 import com.intellij.util.messages.Topic
 
 interface CodeCompletionProgressNotifier {
@@ -10,5 +11,20 @@ interface CodeCompletionProgressNotifier {
         @JvmStatic
         val CODE_COMPLETION_PROGRESS_TOPIC =
             Topic.create("codeCompletionProgressTopic", CodeCompletionProgressNotifier::class.java)
+
+        fun startLoading(project: Project) {
+            handleLoading(project, true)
+        }
+
+        fun stopLoading(project: Project) {
+            handleLoading(project, false)
+        }
+
+        private fun handleLoading(project: Project, loading: Boolean) {
+            if (project.isDisposed) return
+            project.messageBus
+                .syncPublisher(CODE_COMPLETION_PROGRESS_TOPIC)
+                ?.loading(loading)
+        }
     }
 }
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt
@@ -117,9 +117,10 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
         }
 
         IS_FETCHING_COMPLETION.set(request.editor, true)
-        request.editor.project?.messageBus
-            ?.syncPublisher(CodeCompletionProgressNotifier.CODE_COMPLETION_PROGRESS_TOPIC)
-            ?.loading(true)
+
+        request.editor.project?.let {
+            CodeCompletionProgressNotifier.startLoading(it)
+        }
 
         return InlineCompletionSingleSuggestion.build(elements = channelFlow {
             val infillRequest = InfillRequestUtil.buildInfillRequest(request, completionType)

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt
@@ -3,21 +3,14 @@ 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.openapi.diagnostic.thisLogger
-import com.intellij.openapi.diff.impl.patch.IdeaTextPatchBuilder
-import com.intellij.openapi.diff.impl.patch.UnifiedDiffWriter
-import com.intellij.openapi.vcs.VcsException
-import com.intellij.openapi.vcs.changes.ChangeListManager
 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.settings.configuration.ConfigurationSettings
 import ee.carlrobert.codegpt.util.GitUtil
-import java.io.StringWriter
 
 object InfillRequestUtil {
-    private val logger = thisLogger()
 
     suspend fun buildInfillRequest(
         request: InlineCompletionRequest,
@@ -34,24 +27,9 @@ object InfillRequestUtil {
 
         val project = request.editor.project ?: return infillRequestBuilder.build()
         if (service<ConfigurationSettings>().state.codeCompletionSettings.gitDiffEnabled) {
-            GitUtil.getProjectRepository(project)?.let { repository ->
-                try {
-                    val repoRootPath = repository.root.toNioPath()
-                    val changes = ChangeListManager.getInstance(project).allChanges
-                        .sortedBy { it.virtualFile?.timeStamp }
-                    val patches = IdeaTextPatchBuilder.buildPatch(
-                        project, changes, repoRootPath, false, true
-                    )
-                    val diffWriter = StringWriter()
-                    UnifiedDiffWriter.write(
-                        null, repoRootPath, patches, diffWriter, "\n\n", null, null
-                    )
-                    val additionalContext =
-                        diffWriter.toString().cleanDiff().truncateText(1024, false)
-                    infillRequestBuilder.additionalContext(additionalContext)
-                } catch (e: VcsException) {
-                    logger.error("Failed to get git context", e)
-                }
+            val additionalContext = GitUtil.getCurrentChanges(project)
+            if (!additionalContext.isNullOrEmpty()) {
+                infillRequestBuilder.additionalContext(additionalContext)
             }
         }
 
@@ -62,18 +40,6 @@ object InfillRequestUtil {
         return infillRequestBuilder.build()
     }
 
-    private fun String.cleanDiff(showContext: Boolean = false): String =
-        lineSequence()
-            .filterNot { line ->
-                line.startsWith("index ") ||
-                        line.startsWith("diff --git") ||
-                        line.startsWith("---") ||
-                        line.startsWith("+++") ||
-                        line.startsWith("===") ||
-                        (!showContext && line.startsWith(" "))
-            }
-            .joinToString("\n")
-
     private fun getInfillContext(
         request: InlineCompletionRequest,
         caretOffset: Int

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/predictions/AcceptNextPredictionRevisionAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/predictions/AcceptNextPredictionRevisionAction.kt
@@ -0,0 +1,32 @@
+package ee.carlrobert.codegpt.predictions
+
+import com.intellij.codeInsight.hint.HintManagerImpl
+import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.editor.Caret
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.actionSystem.EditorAction
+import com.intellij.openapi.editor.actionSystem.EditorWriteActionHandler
+import ee.carlrobert.codegpt.CodeGPTKeys
+
+class AcceptNextPredictionRevisionAction : EditorAction(Handler()), HintManagerImpl.ActionToIgnore {
+
+    companion object {
+        const val ID = "codegpt.acceptNextPrediction"
+    }
+
+    private class Handler : EditorWriteActionHandler() {
+
+        override fun doExecute(editor: Editor, caret: Caret?, dataContext: DataContext?) {
+            val diffViewer = editor.getUserData(CodeGPTKeys.EDITOR_PREDICTION_DIFF_VIEWER)
+            if (diffViewer != null && diffViewer.isVisible()) {
+                diffViewer.applyChanges()
+                return
+            }
+        }
+
+        override fun isEnabledForCaret(editor: Editor, caret: Caret, dataContext: DataContext): Boolean {
+            val diffViewer = editor.getUserData(CodeGPTKeys.EDITOR_PREDICTION_DIFF_VIEWER)
+            return diffViewer != null && diffViewer.isVisible()
+        }
+    }
+}