feat: support partial code completion acceptance

This commit is contained in:
Carl-Robert Linnupuu 2024-12-10 22:58:13 +00:00
parent 453fa31c65
commit ba12865a72
9 changed files with 236 additions and 83 deletions

View file

@ -0,0 +1,23 @@
package ee.carlrobert.codegpt.actions
import com.intellij.codeInsight.inline.completion.session.InlineCompletionContext
import com.intellij.openapi.actionSystem.ActionPromoter
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.CommonDataKeys
import com.intellij.openapi.actionSystem.DataContext
import ee.carlrobert.codegpt.codecompletions.AcceptNextLineInlayAction
import ee.carlrobert.codegpt.codecompletions.AcceptNextWordInlayAction
class InlayActionPromoter : ActionPromoter {
override fun promote(actions: List<AnAction>, context: DataContext): List<AnAction> {
val editor = CommonDataKeys.EDITOR.getData(context) ?: return emptyList()
if (InlineCompletionContext.getOrNull(editor) == null) {
return emptyList()
}
actions.filterIsInstance<AcceptNextWordInlayAction>().takeIf { it.isNotEmpty() }?.let { return it }
actions.filterIsInstance<AcceptNextLineInlayAction>().takeIf { it.isNotEmpty() }?.let { return it }
return emptyList()
}
}

View file

@ -0,0 +1,68 @@
package ee.carlrobert.codegpt.codecompletions
import com.intellij.codeInsight.inline.completion.InlineCompletion
import com.intellij.codeInsight.inline.completion.InlineCompletionEvent
import com.intellij.codeInsight.inline.completion.InlineCompletionHandler
import com.intellij.codeInsight.inline.completion.InlineCompletionRequest
import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession
import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.application.runWriteAction
import com.intellij.openapi.editor.Caret
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.actionSystem.EditorWriteActionHandler
import com.intellij.openapi.util.TextRange
abstract class AbstractInlayActionHandler : EditorWriteActionHandler() {
abstract fun InlineCompletionSession.getTextToInsert(): String
abstract fun InlineCompletionHandler.invoke(request: InlineCompletionRequest, range: TextRange)
override fun doExecute(editor: Editor, caret: Caret?, dataContext: DataContext?) {
InlineCompletionSession.getOrNull(editor)?.let { session ->
if (!session.isActive() || !session.context.isCurrentlyDisplaying()) return
// TODO: Implement brace matching
InlineCompletion.getHandlerOrNull(editor)?.apply {
val startOffset = editor.caretModel.offset
val textToInsert = session.getTextToInsert()
withIgnoringDocumentChanges {
val suggestionTextRange =
TextRange(startOffset - textToInsert.length, startOffset)
invoke(session.request, suggestionTextRange)
runWriteAction {
editor.document.insertString(startOffset, textToInsert)
editor.caretModel.moveToOffset(startOffset + textToInsert.length)
}
}
}
}
}
override fun isEnabledForCaret(
editor: Editor,
caret: Caret,
dataContext: DataContext?
): Boolean = InlineCompletionSession.getOrNull(editor)?.let {
it.isActive() && it.context.isCurrentlyDisplaying()
} ?: false
}
internal abstract class ApplyNextInlaySuggestionEvent(
private val originalRequest: InlineCompletionRequest,
private val textRange: TextRange
) : InlineCompletionEvent {
override fun toRequest(): InlineCompletionRequest {
return InlineCompletionRequest(
this,
originalRequest.file,
originalRequest.editor,
originalRequest.document,
textRange.startOffset,
textRange.endOffset,
originalRequest.lookupElement
)
}
}

View file

@ -0,0 +1,45 @@
package ee.carlrobert.codegpt.codecompletions
import ai.grazie.nlp.utils.takeWhitespaces
import com.intellij.codeInsight.hint.HintManagerImpl
import com.intellij.codeInsight.inline.completion.InlineCompletionHandler
import com.intellij.codeInsight.inline.completion.InlineCompletionRequest
import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession
import com.intellij.openapi.editor.actionSystem.EditorAction
import com.intellij.openapi.util.TextRange
class AcceptNextLineInlayAction : EditorAction(Handler()), HintManagerImpl.ActionToIgnore {
companion object {
const val ID = "codegpt.acceptNextInlayLine"
}
private class Handler : AbstractInlayActionHandler() {
override fun InlineCompletionSession.getTextToInsert(): String {
val completionText = context.textToInsert()
val lineBreakIndex = completionText.indexOfFirst { it == '\n' }
return if (lineBreakIndex == -1) {
completionText
} else {
val nextLineStart = lineBreakIndex + 1
val whitespacesLength =
completionText.substring(nextLineStart).takeWhitespaces().length
completionText.substring(0, nextLineStart + whitespacesLength)
}
}
override fun InlineCompletionHandler.invoke(
request: InlineCompletionRequest,
range: TextRange
) {
invokeEvent(ApplyNextLineInlaySuggestionEvent(request, range))
}
}
}
internal class ApplyNextLineInlaySuggestionEvent(
originalRequest: InlineCompletionRequest,
textRange: TextRange
) : ApplyNextInlaySuggestionEvent(originalRequest, textRange)

View file

@ -0,0 +1,34 @@
package ee.carlrobert.codegpt.codecompletions
import com.intellij.codeInsight.hint.HintManagerImpl
import com.intellij.codeInsight.inline.completion.InlineCompletionHandler
import com.intellij.codeInsight.inline.completion.InlineCompletionRequest
import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession
import com.intellij.openapi.editor.actionSystem.EditorAction
import com.intellij.openapi.util.TextRange
class AcceptNextWordInlayAction : EditorAction(Handler()), HintManagerImpl.ActionToIgnore {
companion object {
const val ID = "codegpt.acceptNextInlayWord"
}
private class Handler : AbstractInlayActionHandler() {
override fun InlineCompletionSession.getTextToInsert(): String {
return CompletionSplitter.split(context.textToInsert())
}
override fun InlineCompletionHandler.invoke(
request: InlineCompletionRequest,
range: TextRange
) {
invokeEvent(ApplyNextWordInlaySuggestionEvent(request, range))
}
}
}
internal class ApplyNextWordInlaySuggestionEvent(
originalRequest: InlineCompletionRequest,
textRange: TextRange
) : ApplyNextInlaySuggestionEvent(originalRequest, textRange)

View file

@ -1,65 +0,0 @@
package ee.carlrobert.codegpt.codecompletions
import com.intellij.codeInsight.inline.completion.InlineCompletion
import com.intellij.codeInsight.inline.completion.InlineCompletionEvent
import com.intellij.codeInsight.inline.completion.InlineCompletionHandler
import com.intellij.codeInsight.inline.completion.InlineCompletionRequest
import com.intellij.codeInsight.inline.completion.session.InlineCompletionContext
import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession
import com.intellij.openapi.actionSystem.DataContext
import com.intellij.openapi.application.runWriteAction
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.openapi.util.TextRange
class ApplyNextWordInlayAction : EditorAction(Handler()) {
private class Handler : EditorWriteActionHandler() {
override fun doExecute(editor: Editor, caret: Caret?, dataContext: DataContext?) {
InlineCompletionSession.getOrNull(editor)?.let {
if (!it.isActive() || !it.context.isCurrentlyDisplaying()) return
InlineCompletion.getHandlerOrNull(editor)?.apply(
it.applyNextWordInlaySuggestion(it.context, editor)
)
}
}
private fun InlineCompletionSession.applyNextWordInlaySuggestion(
context: InlineCompletionContext,
editor: Editor
): InlineCompletionHandler.() -> Unit = {
val startOffset = editor.caretModel.offset
val textToInsert = CompletionSplitter.split(context.textToInsert())
withIgnoringDocumentChanges {
val suggestionTextRange = TextRange(startOffset - textToInsert.length, startOffset)
invokeEvent(ApplyNextWordInlaySuggestionEvent(request, suggestionTextRange))
runWriteAction {
editor.document.insertString(startOffset, textToInsert)
editor.caretModel.moveToOffset(startOffset + textToInsert.length)
}
}
}
}
}
internal class ApplyNextWordInlaySuggestionEvent(
private val originalRequest: InlineCompletionRequest,
private val textRange: TextRange
) : InlineCompletionEvent {
override fun toRequest(): InlineCompletionRequest {
return InlineCompletionRequest(
this,
originalRequest.file,
originalRequest.editor,
originalRequest.document,
textRange.startOffset,
textRange.endOffset,
originalRequest.lookupElement
)
}
}

View file

@ -1,10 +1,11 @@
package ee.carlrobert.codegpt.codecompletions
import com.intellij.codeInsight.inline.completion.InlineCompletionProviderPresentation
import com.intellij.openapi.actionSystem.ActionManager
import com.intellij.openapi.components.service
import com.intellij.openapi.keymap.KeymapUtil
import com.intellij.openapi.project.Project
import com.intellij.ui.components.JBLabel
import com.intellij.util.ui.JBFont
import ee.carlrobert.codegpt.Icons
import javax.swing.JComponent
import javax.swing.SwingConstants
@ -15,10 +16,21 @@ class CodeCompletionProviderPresentation : InlineCompletionProviderPresentation
val selectedModelCode =
project?.service<CodeCompletionService>()?.getSelectedModelCode() ?: ""
val text = if (selectedModelCode.isNotEmpty()) {
"CodeGPT: $selectedModelCode"
buildString {
append("<html>Model: (<strong>$selectedModelCode</strong>) | ")
append("Accept Word: (<strong>${getShortcutText(AcceptNextWordInlayAction.ID)}</strong>) | ")
append("Accept Line: (<strong>${getShortcutText(AcceptNextLineInlayAction.ID)}</strong>)</html>")
}
} else {
"CodeGPT"
}
return JBLabel(text, Icons.DefaultSmall, SwingConstants.LEADING).withFont(JBFont.small())
return JBLabel(text, Icons.DefaultSmall, SwingConstants.LEADING)
}
private fun getShortcutText(actionId: String): String {
return KeymapUtil.getFirstKeyboardShortcutText(
ActionManager.getInstance().getAction(actionId)
)
}
}

View file

@ -1,5 +1,6 @@
package ee.carlrobert.codegpt.codecompletions
import ai.grazie.nlp.utils.takeWhitespaces
import com.intellij.codeInsight.inline.completion.InlineCompletionEvent
import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement
import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSuggestionUpdateManager
@ -25,17 +26,50 @@ class CodeCompletionSuggestionUpdateAdapter :
event: InlineCompletionEvent,
variant: InlineCompletionVariant.Snapshot
): UpdateResult {
if (event !is ApplyNextWordInlaySuggestionEvent || variant.elements.isEmpty()) {
return Invalidated
}
if (variant.elements.isEmpty()) return Invalidated
val completionText = variant.elements.joinToString("") { it.text }
val textToInsert = event.toRequest().run {
CompletionSplitter.split(completionText)
return when (event) {
is ApplyNextWordInlaySuggestionEvent -> handleNextWordEvent(event, variant, completionText)
is ApplyNextLineInlaySuggestionEvent -> handleNextLineEvent(event, variant, completionText)
else -> Invalidated
}
}
updateRemainingCompletion(event.toRequest().editor, textToInsert)
private fun handleNextWordEvent(
event: ApplyNextWordInlaySuggestionEvent,
variant: InlineCompletionVariant.Snapshot,
completionText: String
): UpdateResult {
val textToInsert = CompletionSplitter.split(completionText)
return createUpdatedVariant(event.toRequest().editor, variant, completionText, textToInsert)
}
private fun handleNextLineEvent(
event: ApplyNextLineInlaySuggestionEvent,
variant: InlineCompletionVariant.Snapshot,
completionText: String
): UpdateResult {
val lineBreakIndex = completionText.indexOf('\n')
if (lineBreakIndex == -1) return Invalidated
val nextLineStart = lineBreakIndex + 1
val whitespacesLength = completionText
.substring(nextLineStart)
.takeWhitespaces()
.length
val textToInsert = completionText.substring(0, nextLineStart + whitespacesLength)
return createUpdatedVariant(event.toRequest().editor, variant, completionText, textToInsert)
}
private fun createUpdatedVariant(
editor: Editor,
variant: InlineCompletionVariant.Snapshot,
completionText: String,
textToInsert: String
): UpdateResult {
updateRemainingCompletion(editor, textToInsert)
return Changed(
variant.copy(
listOf(InlineCompletionGrayTextElement(completionText.removePrefix(textToInsert)))

View file

@ -57,6 +57,7 @@
instance="ee.carlrobert.codegpt.telemetry.ui.preferences.TelemetryConfigurable"
id="tools.preferences.codegpt.telemetry"
displayName="Telemetry"/>
<actionPromoter implementation="ee.carlrobert.codegpt.actions.InlayActionPromoter"/>
<applicationService
serviceImplementation="ee.carlrobert.codegpt.telemetry.core.service.TelemetryServiceFactory"/>
<applicationService serviceImplementation="ee.carlrobert.codegpt.settings.GeneralSettings"/>
@ -91,15 +92,16 @@
</action>
<action
id="codegpt.applyInlaysNextWord"
id="codegpt.acceptNextInlayWord"
text="Apply next word"
class="ee.carlrobert.codegpt.codecompletions.ApplyNextWordInlayAction">
<keyboard-shortcut first-keystroke="control meta alt RIGHT" keymap="$default"/>
<keyboard-shortcut first-keystroke="control meta alt RIGHT" keymap="Mac OS X" replace-all="true"/>
<keyboard-shortcut first-keystroke="control meta alt RIGHT" keymap="Mac OS X 10.5+" replace-all="true"/>
<override-text place="MainMenu" text="Accept Next Word"/>
<override-text place="EditorPopup" text="Accept Next Word"/>
</action>
use-shortcut-of="EditorNextWord"
class="ee.carlrobert.codegpt.codecompletions.AcceptNextWordInlayAction"/>
<action
id="codegpt.acceptNextInlayLine"
text="Apply next line"
use-shortcut-of="EditorLineEnd"
class="ee.carlrobert.codegpt.codecompletions.AcceptNextLineInlayAction"/>
<group id="CodeGPTEditorPopup">
<group id="action.editor.group.EditorActionGroup"

View file

@ -141,7 +141,7 @@ class CodeCompletionServiceTest : IntegrationTest() {
}
val offsetBeforeApply = myFixture.editor.caretModel.offset
myFixture.performEditorAction("codegpt.applyInlaysNextWord")
myFixture.performEditorAction(AcceptNextWordInlayAction.ID)
assertInlineSuggestion("Failed to display next partial inline suggestion.") {
myFixture.run {