From 8daf4f5f05c08b29af3a834ed338a55e744b4a80 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Wed, 13 Nov 2024 23:18:53 +0000 Subject: [PATCH] feat: improve code completion user experience (#763) --- .../llama/form/InfillPromptTemplatePanel.java | 2 +- .../CodeCompletionCompletionEventListener.kt | 113 ------- .../CodeCompletionEventListener.kt | 121 +++++++ .../CodeCompletionInsertAction.kt | 113 +++++++ .../CodeCompletionInsertHandler.kt | 58 +--- .../CodeCompletionTextElement.kt | 38 +++ .../DebouncedCodeCompletionProvider.kt | 82 ++++- .../codegpt/codecompletions/InfillRequest.kt | 11 +- .../CodeCompletionConfigurationForm.kt | 2 +- .../form/CustomServiceCodeCompletionForm.kt | 4 +- .../ee/carlrobert/codegpt/util/StringUtil.kt | 66 ++++ src/main/resources/META-INF/plugin.xml | 5 + .../CodeCompletionServiceTest.kt | 311 ++++++++++++++---- .../carlrobert/codegpt/util/StringUtilTest.kt | 54 +++ 14 files changed, 718 insertions(+), 262 deletions(-) delete mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCompletionEventListener.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt create mode 100644 src/test/kotlin/ee/carlrobert/codegpt/util/StringUtilTest.kt diff --git a/src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java b/src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java index 19873f99..99441f28 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java +++ b/src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java @@ -17,6 +17,6 @@ public class InfillPromptTemplatePanel extends BasePromptTemplatePanel { - - companion object { - private val logger = thisLogger() - } - - private val stringBuilder = StringBuilder() - private val isCancelled = AtomicBoolean(false) - private val isSending = AtomicBoolean(true) - - open fun onComplete(fullMessage: String) {} - open fun onMessage(message: String) {} - - override fun onOpen() { - setLoading(true) - } - - override fun onMessage(message: String, eventSource: EventSource) { - if (isCancelled.get()) return - - val processedMessage = if (infillRequest != null && stringBuilder.isEmpty()) { - message.tryTrimStart(infillRequest.prefix.lines()) - } else { - message - } - - val newLineCount = (stringBuilder.toString() + processedMessage).count { it == '\n' } - if (newLineCount >= MAX_LINES_TO_REQUEST) { - cancelStreaming(processedMessage, eventSource) - return - } - - stringBuilder.append(processedMessage) - - if (newLineCount <= MAX_LINES_TO_DISPLAY && isSending.get()) { - if (newLineCount == MAX_LINES_TO_DISPLAY && processedMessage.contains('\n')) { - isSending.set(false) - onMessage(processedMessage.substring(0, processedMessage.lastIndexOf('\n'))) - } else { - onMessage(processedMessage) - } - } - } - - override fun onComplete(messageBuilder: StringBuilder) { - setLoading(false) - onComplete(stringBuilder.trimEnd().toString()) - } - - override fun onCancelled(messageBuilder: StringBuilder) { - setLoading(false) - onComplete(stringBuilder.trimEnd().toString()) - } - - override fun onError(error: ErrorDetails, ex: Throwable) { - if (ex.message == null || (ex.message != null && ex.message != "Canceled")) { - showNotification(error.message, NotificationType.ERROR) - logger.error(error.message, ex) - } - } - - private fun setLoading(loading: Boolean) { - IS_FETCHING_COMPLETION.set(editor, loading) - editor.project?.messageBus - ?.syncPublisher(CodeCompletionProgressNotifier.CODE_COMPLETION_PROGRESS_TOPIC) - ?.loading(loading) - } - - private fun cancelStreaming(processedMessage: String, eventSource: EventSource) { - stringBuilder.append(processedMessage.substring(0, processedMessage.lastIndexOf('\n'))) - isCancelled.set(true) - isSending.set(false) - eventSource.cancel() - } - - private fun String.tryTrimStart(lines: List): String { - val whiteSpaces = this.takeWhitespaces() - if (lines.size >= 2 - && whiteSpaces.isNotEmpty() - && lines[lines.size - 1].trim().isEmpty() - ) { - return this.trimStart() - } - - if (lines.isNotEmpty()) { - val lastLine = lines[lines.size - 1] - if (lastLine.takeLastWhitespaces().isNotEmpty()) { - return this.trimStart() - } - } - return this - } -} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt new file mode 100644 index 00000000..df156af7 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt @@ -0,0 +1,121 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.intellij.notification.NotificationType +import com.intellij.openapi.application.runReadAction +import com.intellij.openapi.diagnostic.thisLogger +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.util.TextRange +import ee.carlrobert.codegpt.CodeGPTKeys.IS_FETCHING_COMPLETION +import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION +import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification +import ee.carlrobert.codegpt.util.StringUtil +import ee.carlrobert.llm.client.openai.completion.ErrorDetails +import ee.carlrobert.llm.completion.CompletionEventListener +import okhttp3.sse.EventSource + +abstract class CodeCompletionEventListener( + private val editor: Editor +) : CompletionEventListener { + + companion object { + private val logger = thisLogger() + } + + private var isFirstLine = true + private val currentLineBuffer = StringBuilder() + private val incomingTextBuffer = StringBuilder() + + open fun onLineReceived(completionLine: String) {} + + override fun onOpen() { + setLoading(true) + REMAINING_EDITOR_COMPLETION.set(editor, "") + } + + override fun onMessage(message: String, eventSource: EventSource) { + incomingTextBuffer.append(message) + + while (incomingTextBuffer.contains("\n")) { + val lineEndIndex = incomingTextBuffer.indexOf("\n") + val line = incomingTextBuffer.substring(0, lineEndIndex) + '\n' + processCompletionLine(line) + incomingTextBuffer.delete(0, lineEndIndex + 1) + } + } + + private fun processCompletionLine(line: String) { + currentLineBuffer.append(line) + + if (currentLineBuffer.trim().isNotEmpty()) { + val completionText = if (isFirstLine) { + line.adjustWhitespaces().also { + isFirstLine = false + onLineReceived(it) + } + } else { + currentLineBuffer.toString() + } + + appendRemainingCompletion(completionText) + currentLineBuffer.clear() + } + } + + override fun onComplete(messageBuilder: StringBuilder) { + handleCompleted(messageBuilder) + } + + override fun onCancelled(messageBuilder: StringBuilder) { + handleCompleted(messageBuilder) + } + + override fun onError(error: ErrorDetails, ex: Throwable) { + if (ex.message == null || (ex.message != null && ex.message != "Canceled")) { + showNotification(error.message, NotificationType.ERROR) + logger.error(error.message, ex) + } + setLoading(false) + } + + private fun String.adjustWhitespaces(): String { + val adjustedLine = runReadAction { + val lineNumber = editor.document.getLineNumber(editor.caretModel.offset) + val editorLine = editor.document.getText( + TextRange( + editor.document.getLineStartOffset(lineNumber), + editor.document.getLineEndOffset(lineNumber) + ) + ) + + StringUtil.adjustWhitespace(this, editorLine) + } + + return if (adjustedLine.length != this.length) adjustedLine else this + } + + private fun handleCompleted(messageBuilder: StringBuilder) { + setLoading(false) + + if (incomingTextBuffer.isNotEmpty()) { + appendRemainingCompletion(incomingTextBuffer.toString()) + } + + if (isFirstLine) { + val completionLine = messageBuilder.toString().adjustWhitespaces() + REMAINING_EDITOR_COMPLETION.set(editor, completionLine) + onLineReceived(completionLine) + } + } + + private fun appendRemainingCompletion(text: String) { + val previousRemainingText = REMAINING_EDITOR_COMPLETION.get(editor) ?: "" + REMAINING_EDITOR_COMPLETION.set(editor, previousRemainingText + text) + } + + private fun setLoading(loading: Boolean) { + IS_FETCHING_COMPLETION.set(editor, loading) + editor.project?.messageBus + ?.syncPublisher(CodeCompletionProgressNotifier.CODE_COMPLETION_PROGRESS_TOPIC) + ?.loading(loading) + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt new file mode 100644 index 00000000..c68b47fa --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt @@ -0,0 +1,113 @@ +package ee.carlrobert.codegpt.codecompletions + +import ai.grazie.nlp.utils.takeWhitespaces +import com.intellij.codeInsight.hint.HintManagerImpl +import com.intellij.codeInsight.inline.completion.InlineCompletionInsertEnvironment +import com.intellij.codeInsight.inline.completion.session.InlineCompletionContext +import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession +import com.intellij.codeInsight.lookup.LookupManager +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 com.intellij.psi.PsiDocumentManager +import com.intellij.util.concurrency.ThreadingAssertions +import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION + +class CodeCompletionInsertAction : + EditorAction(InsertInlineCompletionHandler()), HintManagerImpl.ActionToIgnore { + + class InsertInlineCompletionHandler : EditorWriteActionHandler() { + override fun executeWriteAction(editor: Editor, caret: Caret?, dataContext: DataContext) { + ThreadingAssertions.assertEventDispatchThread() + ThreadingAssertions.assertWriteAccess() + + val session = InlineCompletionSession.getOrNull(editor) ?: return + val context = session.context + val elements = context.state.elements + .filter { it.element is CodeCompletionTextElement } + .map { it.element as CodeCompletionTextElement } + + for (element in elements) { + val insertEnvironment = InlineCompletionInsertEnvironment( + editor, + session.request.file, + element.textRange + ) + context.copyUserDataTo(insertEnvironment) + + editor.document.insertString(element.textRange.startOffset, element.text) + + if (element.originalText == element.text) { + processStandardCompletionElement(element, editor) + } else { + processPartialCompletionElement(element, editor) + } + + PsiDocumentManager.getInstance(session.request.file.project) + .commitDocument(editor.document) + + session.provider.insertHandler.afterInsertion(insertEnvironment, elements) + + LookupManager.getActiveLookup(editor)?.hideLookup(false) + } + } + + override fun isEnabledForCaret( + editor: Editor, + caret: Caret, + dataContext: DataContext + ): Boolean { + val completionContext = InlineCompletionContext.getOrNull(editor) + val element = completionContext?.state?.elements?.firstOrNull()?.element + if (element is CodeCompletionTextElement) { + return completionContext.startOffset() == (caret.offset + element.offsetDelta) + } + return completionContext?.startOffset() == caret.offset + } + + private fun processStandardCompletionElement( + element: CodeCompletionTextElement, + editor: Editor + ) { + val endOffset = element.textRange.endOffset + editor.caretModel.moveToOffset(endOffset) + + val remainingCompletionLine = REMAINING_EDITOR_COMPLETION.get(editor) + .removePrefix(element.text) + + processRemainingCompletion(remainingCompletionLine, editor, endOffset) + } + + 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) + + val remainingText = REMAINING_EDITOR_COMPLETION.get(editor) + val remainingCompletionLine = if (element.originalText.length > remainingText.length) { + remainingText.removePrefix(element.text) + } else { + remainingText.removePrefix(element.originalText) + } + + processRemainingCompletion(remainingCompletionLine, editor, lineEndOffset + 1) + } + + private fun processRemainingCompletion( + remainingCompletion: String, + editor: Editor, + offset: Int + ) { + val whitespaces = remainingCompletion.takeWhitespaces() + if (whitespaces.isNotEmpty()) { + editor.document.insertString(offset, whitespaces) + editor.caretModel.moveToOffset(offset + whitespaces.length) + } + + val nextCompletionLine = remainingCompletion.removePrefix(whitespaces) + REMAINING_EDITOR_COMPLETION.set(editor, nextCompletionLine) + } + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt index daf6f545..e1756723 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt @@ -5,76 +5,20 @@ import com.intellij.codeInsight.inline.completion.InlineCompletionEvent import com.intellij.codeInsight.inline.completion.InlineCompletionInsertEnvironment import com.intellij.codeInsight.inline.completion.InlineCompletionInsertHandler import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement -import com.intellij.openapi.application.readAction -import com.intellij.openapi.components.service -import com.intellij.openapi.editor.Editor -import com.intellij.openapi.util.TextRange import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION -import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.SupervisorJob -import kotlinx.coroutines.launch - -private const val NEXT_COMPLETION_LINE_COUNT_THRESHOLD = 4 class CodeCompletionInsertHandler : InlineCompletionInsertHandler { - private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob()) - override fun afterInsertion( environment: InlineCompletionInsertEnvironment, elements: List ) { val editor = environment.editor - val appliedText = elements.joinToString("") { it.text } - val existingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) ?: "" - val remainingCompletion = - existingCompletion.substring(appliedText.length, existingCompletion.length) - - REMAINING_EDITOR_COMPLETION.set(editor, remainingCompletion) - + val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) if (remainingCompletion.isNotEmpty()) { InlineCompletion.getHandlerOrNull(editor)?.invoke( InlineCompletionEvent.DirectCall(editor, editor.caretModel.currentCaret) ) - - if (remainingCompletion.count { it == '\n' } <= NEXT_COMPLETION_LINE_COUNT_THRESHOLD) { - scope.launch { - fetchNextCompletion(editor, remainingCompletion) - } - } } } - - private suspend fun fetchNextCompletion(editor: Editor, remainingCompletion: String) { - val project = editor.project ?: return - project.service().getCodeCompletionAsync( - buildNextRequest(editor, remainingCompletion), - object : CodeCompletionCompletionEventListener(editor) { - override fun onComplete(fullMessage: String) { - val nextCompletion = - (REMAINING_EDITOR_COMPLETION.get(editor) ?: "") + fullMessage - REMAINING_EDITOR_COMPLETION.set(editor, nextCompletion) - } - } - ) - } - - private suspend fun buildNextRequest( - editor: Editor, - remainingCompletion: String - ): InfillRequest { - val caretOffset = readAction { editor.caretModel.offset } - val prefix = - (editor.document.getText(TextRange(0, caretOffset)) + remainingCompletion) - .truncateText(MAX_PROMPT_TOKENS, false) - val suffix = - editor.document.getText( - TextRange( - caretOffset, - editor.document.textLength - ) - ).truncateText(MAX_PROMPT_TOKENS) - return InfillRequest.Builder(prefix, suffix).build() - } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt new file mode 100644 index 00000000..25b5f1da --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt @@ -0,0 +1,38 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.intellij.codeInsight.inline.completion.InlineCompletionFontUtils +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionColorTextElement +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.util.TextRange + +class CodeCompletionTextElement( + text: String, + private val insertOffset: Int, + val textRange: TextRange, + val offsetDelta: Int = 0, + val originalText: String = text, +) : InlineCompletionColorTextElement(text, InlineCompletionFontUtils::color) { + + override fun toPresentable(): InlineCompletionElement.Presentable = + Presentable(this, insertOffset, textRange) + + open class Presentable( + element: InlineCompletionElement, + private val insertOffset: Int, + private val textRange: TextRange, + ) : InlineCompletionColorTextElement.Presentable(element, InlineCompletionFontUtils::color) { + + override fun render(editor: Editor, offset: Int) { + super.render(editor, insertOffset) + } + + override fun startOffset(): Int { + return textRange.startOffset + } + + override fun endOffset(): Int { + return textRange.endOffset + } + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt index b9489a4c..0bcac684 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt @@ -2,13 +2,13 @@ package ee.carlrobert.codegpt.codecompletions import com.intellij.codeInsight.inline.completion.* import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement -import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSingleSuggestion import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSuggestion import com.intellij.openapi.application.runInEdt import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger import com.intellij.openapi.editor.Editor +import com.intellij.openapi.util.TextRange import ee.carlrobert.codegpt.CodeGPTKeys.IS_FETCHING_COMPLETION import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION import ee.carlrobert.codegpt.settings.GeneralSettings @@ -18,6 +18,7 @@ import ee.carlrobert.codegpt.settings.service.custom.CustomServiceSettings import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings +import ee.carlrobert.codegpt.util.StringUtil.findCompletionParts import kotlinx.coroutines.channels.ProducerScope import kotlinx.coroutines.channels.awaitClose import kotlinx.coroutines.flow.channelFlow @@ -49,6 +50,14 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { override val providerPresentation: InlineCompletionProviderPresentation get() = CodeCompletionProviderPresentation() + private fun String.extractUntilNewline(): String { + val index = this.indexOf('\n') + if (index == -1) { + return this + } + return this.substring(0, index + 1) + } + override suspend fun getSuggestionDebounced(request: InlineCompletionRequest): InlineCompletionSuggestion { val editor = request.editor val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) @@ -56,7 +65,7 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { && remainingCompletion != null && remainingCompletion.isNotEmpty() ) { - return sendNextSuggestion(remainingCompletion) + return sendNextSuggestion(remainingCompletion.extractUntilNewline(), request) } val project = editor.project @@ -74,10 +83,12 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { ?.loading(true) val infillRequest = InfillRequestUtil.buildInfillRequest(request) - val call = project.service().getCodeCompletionAsync( - infillRequest, - getEventListener(request.editor, infillRequest) - ) + val call = project + .service() + .getCodeCompletionAsync( + infillRequest, + getEventListener(request.editor, infillRequest) + ) currentCallRef.set(call) awaitClose { currentCallRef.getAndSet(null)?.cancel() } }) @@ -114,27 +125,62 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { private fun ProducerScope.getEventListener( editor: Editor, infillRequest: InfillRequest - ) = object : CodeCompletionCompletionEventListener(editor, infillRequest) { - override fun onMessage(message: String) { - runInEdt { - trySend(InlineCompletionGrayTextElement(message)) - } - } + ) = object : CodeCompletionEventListener(editor) { - override fun onComplete(fullMessage: String) { - REMAINING_EDITOR_COMPLETION.set(editor, fullMessage) + override fun onLineReceived(completionLine: String) { + runInEdt { + val editorLineSuffix = editor.getLineSuffixAfterCaret() + if (editorLineSuffix.isEmpty()) { + trySend( + CodeCompletionTextElement( + completionLine, + infillRequest.caretOffset, + TextRange.from(infillRequest.caretOffset, completionLine.length), + ) + ) + } else { + var prevStartOffset = infillRequest.caretOffset + val completionParts = + findCompletionParts(editorLineSuffix, completionLine.trimEnd()) + + completionParts.forEach { (completionPart, offsetDelta) -> + val element = CodeCompletionTextElement( + completionPart, + infillRequest.caretOffset + offsetDelta, + TextRange.from(prevStartOffset + offsetDelta, completionPart.length), + offsetDelta, + completionLine + ) + prevStartOffset += completionPart.length + + trySend(element) + } + } + + } } } - private fun sendNextSuggestion(nextCompletion: String): InlineCompletionSingleSuggestion { + private fun Editor.getLineSuffixAfterCaret(): String { + val lineEndOffset = document.getLineEndOffset(document.getLineNumber(caretModel.offset)) + return document.getText(TextRange(caretModel.offset, lineEndOffset)) + } + + private fun sendNextSuggestion( + nextCompletion: String, + request: InlineCompletionRequest + ): InlineCompletionSingleSuggestion { + return InlineCompletionSingleSuggestion.build(elements = channelFlow { launch { trySend( - InlineCompletionGrayTextElement( - nextCompletion.lines().take(2).joinToString("\n") + CodeCompletionTextElement( + nextCompletion, + request.startOffset, + TextRange.from(request.startOffset, nextCompletion.length), ) ) } }) } -} +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt index 945a596c..9ceff084 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt @@ -13,13 +13,15 @@ const val MAX_PROMPT_TOKENS = 128 class InfillRequest private constructor( val prefix: String, val suffix: String, + val caretOffset: Int, val fileDetails: FileDetails?, val vcsDetails: VcsDetails?, val context: InfillContext? ) { companion object { - fun builder(prefix: String, suffix: String) = Builder(prefix, suffix) + fun builder(prefix: String, suffix: String, caretOffset: Int) = + Builder(prefix, suffix, caretOffset) } data class VcsDetails(val stagedDiff: String? = null, val unstagedDiff: String? = null) @@ -28,13 +30,15 @@ class InfillRequest private constructor( class Builder { private val prefix: String private val suffix: String + private val caretOffset: Int private var fileDetails: FileDetails? = null private var vcsDetails: VcsDetails? = null private var context: InfillContext? = null - constructor(prefix: String, suffix: String) { + constructor(prefix: String, suffix: String, caretOffset: Int) { this.prefix = prefix this.suffix = suffix + this.caretOffset = caretOffset } constructor(document: Document, caretOffset: Int) { @@ -48,6 +52,7 @@ class InfillRequest private constructor( document.textLength ) ).truncateText(MAX_PROMPT_TOKENS) + this.caretOffset = caretOffset } fun fileDetails(fileDetails: FileDetails) = apply { this.fileDetails = fileDetails } @@ -55,7 +60,7 @@ class InfillRequest private constructor( fun context(context: InfillContext) = apply { this.context = context } fun build() = - InfillRequest(prefix, suffix, fileDetails, vcsDetails, context) + InfillRequest(prefix, suffix, caretOffset, fileDetails, vcsDetails, context) } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt index ac381e65..3def827a 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt @@ -65,7 +65,7 @@ class CodeCompletionConfigurationForm( promptTemplateHelpText.setToolTipText(null) val description = StringEscapeUtils.escapeHtml4( - template.buildPrompt(InfillRequest.Builder("PREFIX", "SUFFIX").build()) + template.buildPrompt(InfillRequest.Builder("PREFIX", "SUFFIX", 0).build()) ) HelpTooltip() .setTitle(template.toString()) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt index ec6fd0c5..a38a3bba 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt @@ -163,7 +163,7 @@ class CustomServiceCodeCompletionForm( private fun testConnection() { CompletionRequestService.getInstance().getCustomOpenAICompletionAsync( CodeCompletionRequestFactory.buildCustomRequest( - InfillRequest.Builder("Hello", "!").build(), + InfillRequest.Builder("Hello", "!", 0).build(), urlField.text, tabbedPane.headers, tabbedPane.body, @@ -206,7 +206,7 @@ class CustomServiceCodeCompletionForm( val description = StringEscapeUtils.escapeHtml4( template.buildPrompt( - InfillRequest.Builder("PREFIX", "SUFFIX").build(), + InfillRequest.Builder("PREFIX", "SUFFIX", 0).build(), ) ) HelpTooltip() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt new file mode 100644 index 00000000..7878cbaa --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/StringUtil.kt @@ -0,0 +1,66 @@ +package ee.carlrobert.codegpt.util + +import ai.grazie.nlp.utils.takeWhitespaces +import com.intellij.util.diff.Diff + +object StringUtil { + + fun adjustWhitespace( + completionLine: String, + editorLine: String + ): String { + val editorWhitespaces = editorLine.takeWhitespaces() + + if (completionLine.isNotEmpty() && editorWhitespaces.isNotEmpty()) { + if (completionLine.startsWith(editorWhitespaces)) { + return completionLine.substring(editorWhitespaces.length) + } + if (editorLine.isBlank()) { + val completionWhitespaces = completionLine.takeWhitespaces() + return completionLine.substring(completionWhitespaces.length) + } + } + + return completionLine + } + + fun findCompletionParts( + editorLineSuffix: String, + completionLine: String + ): List> { + val nonOverlappingPart = findNonOverlappingPart(editorLineSuffix, completionLine) + if (nonOverlappingPart.length == completionLine.length) { + return listOf(Pair(completionLine, 0)) + } + + val result = ArrayList>() + val editorChars: IntArray = editorLineSuffix.chars().toArray() + val completionChars: IntArray = completionLine.chars().toArray() + val changes: List = + Diff.buildChanges(editorChars, completionChars)?.toList() ?: emptyList() + for (change in changes) { + val part = completionLine.substring(change.line1, change.line1 + change.inserted) + result.add(Pair(part, change.line0)) + } + + return result + } + + private fun findNonOverlappingPart( + editorLineSuffix: String, + completionLine: String + ): String { + var i = editorLineSuffix.length - 1 + var j = completionLine.length - 1 + while (i >= 0 && j >= 0 && editorLineSuffix[i] == completionLine[j]) { + i-- + j-- + } + + if (j >= 0) { + return completionLine.substring(0, j + 1) + } + + return "" + } +} diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index 1c28b195..0d859663 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -84,6 +84,11 @@ messages.codegpt + + + + + - assertThat(request.uri.path).isEqualTo("/v1/code/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.body) - .extracting("model", "prefix", "suffix", "fileExtension") - .containsExactly("TEST_CODE_MODEL", "p", "", "java") - listOf( - jsonMapResponse("choices", jsonArray(jsonMap("text", "ublic static"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", " void main(String"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", "[] args) {\n"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", " System.out.print"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", "ln(\"Hello, Worl"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", "d!\");\n"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", "}\n"))), - jsonMapResponse( - "choices", - jsonArray(jsonMap("text", "private static int getX() {\n")) - ) - ) - }) - myFixture.type('p') - assertInlineSuggestion("Failed to display initial inline suggestion.") { - "ublic static void main(String[] args) {\n System.out.println(\"Hello, World!\");" == it - } - assertThat(REMAINING_EDITOR_COMPLETION.get(myFixture.editor)).isEqualTo( - "ublic static void main(String[] args) {\n" + - " System.out.println(\"Hello, World!\");\n" + - "}\n" + - "private static int getX() {" - ) - expectCodeGPT(StreamHttpExchange { request: RequestEntity -> - assertThat(request.uri.path).isEqualTo("/v1/code/completions") - assertThat(request.method).isEqualTo("POST") - assertThat(request.body) - .extracting("model", "prefix", "suffix", "fileExtension") - .containsExactly( - "TEST_CODE_MODEL", - "public static void main(String[] args) {\n" + - " System.out.println(\"Hello, World!\");\n" + - "}\n" + - "private static int getX() {", - "", - "java" - ) - listOf( - jsonMapResponse("choices", jsonArray(jsonMap("text", "\n retur"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", "n 10;\n"))), - jsonMapResponse("choices", jsonArray(jsonMap("text", "}\n"))), - ) - }) - - myFixture.type('\t') - - PlatformTestUtil.waitWithEventsDispatching( - "Failed to retrieve next completion", - { - REMAINING_EDITOR_COMPLETION.get(myFixture.editor) == "\n}\nprivate static int getX() {" - }, - 10 - ) - } - fun `test apply next partial completion word`() { useLlamaService(true) myFixture.configureByText( @@ -178,7 +113,7 @@ class CodeCompletionServiceTest : IntegrationTest() { .extracting("prompt") .isEqualTo( InfillPromptTemplate.CODE_LLAMA.buildPrompt( - InfillRequest.Builder(prefix, suffix).build() + InfillRequest.Builder(prefix, suffix, 0).build() ) ) listOf( @@ -205,7 +140,249 @@ class CodeCompletionServiceTest : IntegrationTest() { } } - private fun assertInlineSuggestion(errorMessage: String, onAssert: (String) -> Boolean) { + fun `test apply inline suggestions without initial following text`() { + useCodeGPTService() + myFixture.configureByText( + "CompletionTest.java", + "class Node {\n " + ) + myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(1, 2)) + expectCodeGPT(StreamHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/code/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body) + .extracting("model", "prefix", "suffix", "fileExtension") + .containsExactly( + "TEST_CODE_MODEL", + "class Node {\n ", + "", + "java" + ) + listOf( + jsonMapResponse("choices", jsonArray(jsonMap("text", "\n int data;"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "\n Node lef"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "t;\n Node ri"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "ght;\n\n public"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", " Node(int data"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", ") {\n"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", " this.data ="))), + jsonMapResponse("choices", jsonArray(jsonMap("text", " data;\n }"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "\n}"))), + ) + }) + + myFixture.type(' ') + assertRemainingCompletion { + it == "int data;\n" + + " Node left;\n" + + " Node right;\n" + + "\n" + + " public Node(int data) {\n" + + " this.data = data;\n" + + " }\n" + + "}" + } + assertInlineSuggestion { + it == "int data;\n" + } + myFixture.type('\t') + assertRemainingCompletion { + it == "Node left;\n" + + " Node right;\n" + + "\n" + + " public Node(int data) {\n" + + " this.data = data;\n" + + " }\n" + + "}" + } + assertInlineSuggestion { + it == "Node left;\n" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(2, 3)) + myFixture.type('\t') + assertRemainingCompletion { + it == "Node right;\n" + + "\n" + + " public Node(int data) {\n" + + " this.data = data;\n" + + " }\n" + + "}" + } + assertInlineSuggestion("Failed to assert remaining completion.") { + it == "Node right;\n" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(3, 3)) + myFixture.type('\t') + assertRemainingCompletion { + it == "public Node(int data) {\n" + + " this.data = data;\n" + + " }\n" + + "}" + } + assertInlineSuggestion { + it == "public Node(int data) {\n" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(5, 3)) + myFixture.type('\t') + assertRemainingCompletion { + it == "this.data = data;\n" + + " }\n" + + "}" + } + assertInlineSuggestion { + it == "this.data = data;\n" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(6, 6)) + myFixture.type('\t') + assertRemainingCompletion { + it == "}\n" + + "}" + } + assertInlineSuggestion { + it == "}\n" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(7, 3)) + myFixture.type('\t') + assertRemainingCompletion { + it == "}" + } + assertInlineSuggestion { + it == "}" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(8, 0)) + myFixture.type('\t') + assertRemainingCompletion { + it == "" + } + } + + fun `test apply inline suggestions with initial following text`() { + useCodeGPTService() + myFixture.configureByText( + "CompletionTest.java", + "if () {\n \n} else {\n}" + ) + myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(0, 4)) + expectCodeGPT(StreamHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/code/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body) + .extracting("model", "prefix", "suffix", "fileExtension") + .containsExactly( + "TEST_CODE_MODEL", + "if (r", + ") {\n \n} else {\n}", + "java" + ) + listOf( + jsonMapResponse("choices", jsonArray(jsonMap("text", "oot == n"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "ull) {\n"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", " root = new Node"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "(data);\n"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", " return;"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "\n} else {"))), + ) + }) + myFixture.type('r') + assertRemainingCompletion { + it == "oot == null) {\n" + + " root = new Node(data);\n" + + " return;\n" + + "} else {" + } + assertInlineSuggestion { + it == "oot == null" + } + myFixture.type('\t') + assertRemainingCompletion { + it == "root = new Node(data);\n" + + " return;\n" + + "} else {" + } + assertInlineSuggestion { + it == "root = new Node(data);\n" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(1, 3)) + myFixture.type('\t') + assertRemainingCompletion { + it == "return;\n" + + "} else {" + } + assertInlineSuggestion("Failed to assert remaining completion.") { + it == "return;\n" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(2, 3)) + myFixture.type('\t') + assertRemainingCompletion { + it == "} else {" + } + assertInlineSuggestion { + it == "} else {" + } + assertThat(myFixture.editor.caretModel.visualPosition).isEqualTo(VisualPosition(3, 0)) + myFixture.type('\t') + assertRemainingCompletion { + it == "" + } + } + + fun `test adjust completion line whitespaces`() { + useCodeGPTService() + myFixture.configureByText( + "CompletionTest.java", + "class Node {\n" + + " \n" + + "}" + ) + myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(1, 3)) + expectCodeGPT(StreamHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/v1/code/completions") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body) + .extracting("model", "prefix", "suffix", "fileExtension") + .containsExactly( + "TEST_CODE_MODEL", + "class Node {\n ", + "\n}", + "java" + ) + listOf( + jsonMapResponse("choices", jsonArray(jsonMap("text", "\n int data;"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "\n Node"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", " left;\n N"))), + jsonMapResponse("choices", jsonArray(jsonMap("text", "ode right;\n"))), + ) + }) + myFixture.type(' ') + assertRemainingCompletion { + it == "int data;\n" + + " Node left;\n" + + " Node right;\n" + } + assertInlineSuggestion { + it == "int data;\n" + } + } + + private fun assertRemainingCompletion( + errorMessage: String = "Failed to assert remaining suggestion", + onAssert: (String) -> Boolean + ) { + PlatformTestUtil.waitWithEventsDispatching( + errorMessage, + { + val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(myFixture.editor) + ?: return@waitWithEventsDispatching false + onAssert(remainingCompletion) + }, + 5 + ) + } + + private fun assertInlineSuggestion( + errorMessage: String = "Failed to assert inline suggestion", + onAssert: (String) -> Boolean + ) { PlatformTestUtil.waitWithEventsDispatching( errorMessage, { diff --git a/src/test/kotlin/ee/carlrobert/codegpt/util/StringUtilTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/util/StringUtilTest.kt new file mode 100644 index 00000000..92d7db74 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/util/StringUtilTest.kt @@ -0,0 +1,54 @@ +package ee.carlrobert.codegpt.util + +import ee.carlrobert.codegpt.util.StringUtil.findCompletionParts +import org.assertj.core.api.Assertions.assertThat +import org.junit.Test + +class StringUtilTest { + + @Test + fun `should parse completion without brackets and braces`() { + val completionLine = "root != null" + val editorLineSuffix = ") {\n" + + val result = findCompletionParts(editorLineSuffix, completionLine) + + assertThat(result[0].second).isEqualTo(0) + assertThat(result[0].first).isEqualTo("root != null") + } + + @Test + fun `should parse completion with closing bracket and brace into separate parts`() { + val completionLine = "root != null) {\n" + val editorLineSuffix = ")\n" + + val result = findCompletionParts(editorLineSuffix, completionLine) + + assertThat(result[0].second).isEqualTo(0) + assertThat(result[0].first).isEqualTo("root != null") + assertThat(result[1].second).isEqualTo(1) + assertThat(result[1].first).isEqualTo(" {") + } + + @Test + fun `should parse completion when editor suffix contains closing bracket and brace`() { + val completionLine = "root != null) {\n" + val editorLineSuffix = ") {\n" + + val result = findCompletionParts(editorLineSuffix, completionLine) + + assertThat(result[0].second).isEqualTo(0) + assertThat(result[0].first).isEqualTo("root != null") + } + + @Test + fun `should parse completion between opening and closing brackets`() { + val completionLine = "(root != null) {\n" + val editorLineSuffix = "() {\n" + + val result = findCompletionParts(editorLineSuffix, completionLine) + + assertThat(result[0].second).isEqualTo(1) + assertThat(result[0].first).isEqualTo("root != null") + } +} \ No newline at end of file