From 0fa28f02a9a06747e9e4ded5cb550bd9af7c8575 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Mon, 12 May 2025 16:02:39 +0100 Subject: [PATCH] feat: new tab experience --- .../treesitter/CodeCompletionParser.java | 19 -- .../ee/carlrobert/codegpt/CodeGPTKeys.java | 6 + .../advanced/AdvancedSettingsState.java | 2 +- .../llama/form/InfillPromptTemplatePanel.java | 3 +- .../codegpt/ui/checkbox/FileCheckboxTree.java | 1 - .../codegpt/CodeGPTLookupListener.kt | 30 +- .../CodeAssistantFeatureToggleActions.kt.kt | 35 --- .../CodeCompletionCacheService.kt | 93 ++++++ .../CodeCompletionEventListener.kt | 282 +++++++++--------- .../CodeCompletionFormatter.kt | 277 +++++++++++++++++ .../CodeCompletionInsertAction.kt | 145 --------- .../CodeCompletionInsertHandler.kt | 49 ++- .../CodeCompletionRequestFactory.kt | 19 +- .../codecompletions/CodeCompletionService.kt | 15 +- .../CodeCompletionTextElement.kt | 1 + .../DebouncedCodeCompletionProvider.kt | 184 ++++-------- .../codegpt/codecompletions/InfillRequest.kt | 23 +- .../codecompletions/InfillRequestUtil.kt | 9 +- .../LookupInlineCompletionEvent.kt | 54 ++++ .../edit/CodeCompletionStreamObserver.kt | 46 +++ .../edit/GrpcCallCredentials.kt | 29 ++ .../codecompletions/edit/GrpcClientService.kt | 251 ++++++++-------- .../edit/NextEditStreamObserver.kt | 76 +++++ .../predictions/CodeSuggestionDiffViewer.kt | 114 +++---- .../predictions/OpenPredictionAction.kt | 2 +- .../codegpt/predictions/PredictionService.kt | 19 +- .../TriggerCustomPredictionAction.kt | 2 +- .../CodeCompletionConfigurationForm.kt | 11 - .../configuration/ConfigurationSettings.kt | 1 - .../service/codegpt/CodeGPTServiceForm.kt | 15 +- .../service/codegpt/CodeGPTServiceSettings.kt | 4 +- .../ee/carlrobert/codegpt/util/GitUtil.kt | 1 + src/main/proto/code-completion.proto | 31 ++ src/main/resources/META-INF/plugin.xml | 23 -- .../resources/messages/codegpt.properties | 1 - .../CodeCompletionServiceTest.kt | 62 +++- 36 files changed, 1155 insertions(+), 780 deletions(-) delete mode 100644 src/main/kotlin/ee/carlrobert/codegpt/actions/CodeAssistantFeatureToggleActions.kt.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCacheService.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionFormatter.kt delete mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/LookupInlineCompletionEvent.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcCallCredentials.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/NextEditStreamObserver.kt create mode 100644 src/main/proto/code-completion.proto diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java index 6443efea..568d1ae7 100644 --- a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java @@ -37,17 +37,6 @@ public class CodeCompletionParser { result.deleteCharAt(result.length() - 1); - if (result.length() > 1 && result.charAt(result.length() - 1) == '{') { - long bracketCount = result.chars().filter(ch -> ch == '{').count(); - if (bracketCount == 1) { - var newTree = parser.parseString(currentTree, prefix + result + "}" + suffix); - var treeString = newTree.getRootNode().toString(); - if (!treeString.contains("ERROR")) { - return result + "}"; - } - } - } - input = prefix + result + suffix; currentTree = parser.parseString(currentTree, input); @@ -56,14 +45,6 @@ public class CodeCompletionParser { } } - if (output.contains("\n")) { - var finalResult = output.substring(0, output.indexOf("\n")); - if (finalResult.length() > 1 && finalResult.charAt(finalResult.length() - 1) == '{') { - return finalResult + "}"; - } - return finalResult; - } - return output; } diff --git a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java index 6015dd98..f86328df 100644 --- a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java +++ b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java @@ -4,6 +4,8 @@ import com.intellij.openapi.util.Key; import ee.carlrobert.codegpt.predictions.CodeSuggestionDiffViewer; import ee.carlrobert.codegpt.toolwindow.chat.editor.ToolWindowEditorFileDetails; import ee.carlrobert.llm.client.codegpt.CodeGPTUserDetails; +import ee.carlrobert.service.NextEditResponse; +import ee.carlrobert.service.PartialCodeCompletionResponse; public class CodeGPTKeys { @@ -19,6 +21,10 @@ public class CodeGPTKeys { Key.create("codegpt.isPromptTextFieldDocument"); public static final Key EDITOR_PREDICTION_DIFF_VIEWER = Key.create("codegpt.editorPredictionDiffViewer"); + public static final Key REMAINING_CODE_COMPLETION = + Key.create("codegpt.remainingCodeCompletion"); + public static final Key REMAINING_PREDICTION_RESPONSE = + Key.create("codegpt.remainingPredictionResponse"); public static final Key TOOLWINDOW_EDITOR_FILE_DETAILS = Key.create("proxyai.toolwindowEditorFileDetails"); } diff --git a/src/main/java/ee/carlrobert/codegpt/settings/advanced/AdvancedSettingsState.java b/src/main/java/ee/carlrobert/codegpt/settings/advanced/AdvancedSettingsState.java index 21f1fa86..96bc9050 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/advanced/AdvancedSettingsState.java +++ b/src/main/java/ee/carlrobert/codegpt/settings/advanced/AdvancedSettingsState.java @@ -12,7 +12,7 @@ public class AdvancedSettingsState { private String proxyUsername; private String proxyPassword; private int connectTimeout = 120; - private int readTimeout = 120; + private int readTimeout = 600; public String getProxyHost() { return proxyHost; 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 d6bba37a..961e7731 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 @@ -1,6 +1,5 @@ package ee.carlrobert.codegpt.settings.service.llama.form; -import ee.carlrobert.codegpt.codecompletions.CompletionType; import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate; import ee.carlrobert.codegpt.codecompletions.InfillRequest; @@ -19,7 +18,7 @@ public class InfillPromptTemplatePanel extends BasePromptTemplatePanel().state.nextEditsEnabled + val project = editor.project ?: return + + if (!project.service().isCodeCompletionsEnabled() || editor.editorKind != EditorKind.MAIN_EDITOR ) { return } - val settings = service().state - if (settings.codeCompletionSettings.codeCompletionsEnabled) { - settings.codeCompletionSettings.codeCompletionsEnabled = false - OverlayUtil.showNotification( - "Code completions and multi-line edits cannot be active simultaneously.", - NotificationType.WARNING - ) - } - - ApplicationManager.getApplication().executeOnPooledThread { - service().displayInlineDiff(editor) - } + InlineCompletion.getHandlerOrNull(editor)?.invokeEvent( + LookupInlineCompletionEvent(event) + ) } }) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeAssistantFeatureToggleActions.kt.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeAssistantFeatureToggleActions.kt.kt deleted file mode 100644 index dcf39c32..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/CodeAssistantFeatureToggleActions.kt.kt +++ /dev/null @@ -1,35 +0,0 @@ -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().state - settings.nextEditsEnabled = enableFeatureAction - } - - override fun update(e: AnActionEvent) { - val codeAssistantEnabled = service().state.nextEditsEnabled - - e.presentation.isVisible = GeneralSettings.getSelectedService() == CODEGPT - && codeAssistantEnabled != enableFeatureAction - e.presentation.isEnabled = GeneralSettings.getSelectedService() == CODEGPT - } - - override fun getActionUpdateThread(): ActionUpdateThread { - return ActionUpdateThread.BGT - } -} - -class EnableNextEditsAction : CodeAssistantFeatureToggleAction(true) - -class DisableNextEditsAction : CodeAssistantFeatureToggleAction(false) \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCacheService.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCacheService.kt new file mode 100644 index 00000000..484e200c --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCacheService.kt @@ -0,0 +1,93 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.google.common.cache.CacheBuilder +import com.google.common.cache.CacheLoader +import com.google.common.cache.LoadingCache +import com.intellij.openapi.application.runReadAction +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.service +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.project.Project +import java.util.concurrent.ConcurrentHashMap + +@Service(Service.Level.PROJECT) +class CodeCompletionCacheService() { + private val cacheCounter = ConcurrentHashMap() + private val cache: LoadingCache = CacheBuilder.newBuilder() + .maximumSize(10) + .recordStats() + .build(object : CacheLoader() { + override fun load(key: String): String? = null + }) + + fun getAll(): Map { + return ConcurrentHashMap(cache.asMap()) + } + + fun get(key: String): String? { + val value = cache.getIfPresent(key) + if (value != null) { + cache.invalidate(key) + cache.put(key, value) + } + return value + } + + fun clear() { + cache.invalidateAll() + } + + fun delete(key: String) { + cache.invalidate(key) + } + + fun set(key: String, value: String) { + cache.put(key, value) + } + + fun normalize(src: String): String { + return src.replace("\n", "").replace("\\s+".toRegex(), "").replace("\\s".toRegex(), "") + } + + fun getKey(prefix: String, suffix: String): String { + return if (suffix.isNotEmpty()) { + normalize("$prefix #### $suffix") + } else { + normalize(prefix) + } + } + + fun getCache(editor: Editor): String? { + val caretOffset = runReadAction { editor.caretModel.offset } + val prefix = editor.document.text.substring(0, caretOffset) + val suffix = editor.document.text.substring(caretOffset) + return getCache(prefix, suffix) + } + + fun getCache(prefix: String, suffix: String): String? { + val key = getKey(prefix, suffix) + if (cacheCounter.containsKey(key)) { + cacheCounter[key] = cacheCounter[key]!! + 1 + } else { + cacheCounter[key] = 1 + } + if (cacheCounter[key]!! > 3) { + cache.invalidate(key) + cacheCounter.remove(key) + } + + return get(key) + } + + fun setCache(prefix: String, suffix: String, completion: String) { + val key = getKey(prefix, suffix) + set(key, completion) + } + + companion object { + @JvmStatic + fun getInstance(project: Project): CodeCompletionCacheService { + return project.service() + } + } +} \ 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 index 86c0ad91..2b6696fb 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt @@ -1,47 +1,166 @@ package ee.carlrobert.codegpt.codecompletions -import com.intellij.codeInsight.inline.completion.InlineCompletionRequest +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement import com.intellij.notification.NotificationType import com.intellij.openapi.application.runInEdt -import com.intellij.openapi.application.runWriteAction +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.util.TextRange -import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION -import ee.carlrobert.codegpt.codecompletions.CompletionUtil.formatCompletion +import ee.carlrobert.codegpt.CodeGPTKeys +import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService import ee.carlrobert.codegpt.settings.GeneralSettings +import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings +import ee.carlrobert.codegpt.treesitter.CodeCompletionParserFactory import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification -import ee.carlrobert.codegpt.util.EditorUtil.adjustWhitespaces import ee.carlrobert.llm.client.openai.completion.ErrorDetails import ee.carlrobert.llm.completion.CompletionEventListener +import ee.carlrobert.service.PartialCodeCompletionResponse +import kotlinx.coroutines.channels.ProducerScope import okhttp3.sse.EventSource -import kotlin.math.min +import java.util.concurrent.atomic.AtomicBoolean -abstract class CodeCompletionEventListener( - private val editor: Editor +class CodeCompletionEventListener( + private val editor: Editor, + private val channel: ProducerScope ) : CompletionEventListener { companion object { private val logger = thisLogger() } - abstract fun handleCompleted(messageBuilder: StringBuilder) + private val cancelled = AtomicBoolean(false) + private val messageBuilder = StringBuilder() + private var firstLine: String? = null + private val firstLineSent = AtomicBoolean(false) + private val cursorOffset = runReadAction { editor.caretModel.offset } + private val prefix = editor.document.getText(TextRange(0, cursorOffset)) + private val suffix = + editor.document.getText(TextRange(cursorOffset, editor.document.textLength)) + private val cache = editor.project?.service() override fun onOpen() { setLoading(true) } - override fun onComplete(messageBuilder: StringBuilder) { - setLoading(false) - handleCompleted(messageBuilder) + override fun onMessage(message: String, eventSource: EventSource) { + if (cancelled.get()) { + return + } + + messageBuilder.append(message) + + trySendFirstLine(eventSource) + } + + fun isNotAllowed(completion: String): Boolean { + if (completion.contains("No newline at end of file")) { + return true + } else if (completion.trim().startsWith("+")) { + return true + } + return false + } + + private fun extractUpToRelevantNewline(message: String): String? { + if (message.isEmpty()) return null + val firstNewline = message.indexOf('\n') + if (firstNewline == -1) return null + return if (firstNewline == 0) { + val secondNewline = message.indexOf('\n', 1) + if (secondNewline != -1) { + message.substring(0, secondNewline) + } else { + message + } + } else { + message.substring(0, firstNewline) + } + } + + private fun trySendFirstLine(eventSource: EventSource) { + if (firstLine != null) { + return + } + + var newLine = extractUpToRelevantNewline(messageBuilder.toString()) + if (newLine != null && !firstLineSent.get()) { + val formattedLine = CodeCompletionFormatter(editor).format(newLine) + + if (isNotAllowed(formattedLine)) { + cancelled.set(true) + eventSource.cancel() + return + } + + runInEdt { + channel.trySend(InlineCompletionGrayTextElement(formattedLine)) + } + firstLineSent.set(true) + firstLine = newLine + } + } + + override fun onComplete(finalResult: StringBuilder) { + try { + CodeGPTKeys.REMAINING_CODE_COMPLETION.set(editor, null) + CodeGPTKeys.REMAINING_PREDICTION_RESPONSE.set(editor, null) + + if (cancelled.get() || finalResult.isEmpty()) { + return + } + + if (firstLineSent.get() && firstLine != null) { + val remainingContent = finalResult.removePrefix(firstLine!!).toString() + if (remainingContent.trim().isEmpty()) { + return + } + + val parsedContent = parseOutput(firstLine + remainingContent) + if (parsedContent.isNotEmpty()) { + cache?.setCache(prefix, suffix, firstLine + parsedContent) + + CodeGPTKeys.REMAINING_CODE_COMPLETION.set( + editor, + PartialCodeCompletionResponse.newBuilder() + .setPartialCompletion(remainingContent) + .build() + ) + } + } else { + val formattedLine = CodeCompletionFormatter(editor).format(finalResult.toString()) + if (formattedLine.isEmpty()) { + editor.project?.service()?.getNextEdit( + editor, + prefix + suffix, + runReadAction { editor.caretModel.offset }) + return + } + + if (isNotAllowed(formattedLine)) { + return + } + + val parsedContent = parseOutput(formattedLine) + if (parsedContent.isNotEmpty()) { + cache?.setCache(prefix, suffix, parsedContent) + runInEdt { + channel.trySend(InlineCompletionGrayTextElement(parsedContent)) + } + } + } + } finally { + handleCompleted() + } } override fun onCancelled(messageBuilder: StringBuilder) { - setLoading(false) - handleCompleted(messageBuilder) + cancelled.set(true) + handleCompleted() } override fun onError(error: ErrorDetails, ex: Throwable) { @@ -56,6 +175,11 @@ abstract class CodeCompletionEventListener( showNotification(error.message, NotificationType.ERROR) logger.error(error.message, ex) } + + setLoading(false) + } + + private fun handleCompleted() { setLoading(false) } @@ -64,129 +188,15 @@ abstract class CodeCompletionEventListener( CompletionProgressNotifier.update(it, loading) } } -} -class CodeCompletionMultiLineEventListener( - private val request: InlineCompletionRequest, - private val onCompletionReceived: (String) -> Unit -) : CodeCompletionEventListener(request.editor) { - - override fun handleCompleted(messageBuilder: StringBuilder) { - request.editor.project?.let { CompletionProgressNotifier.update(it, false) } - runInEdt { - onCompletionReceived(runWriteAction { - messageBuilder.toString().formatCompletion(request) - }) - } - } -} - -class CodeCompletionSingleLineEventListener( - private val editor: Editor, - private val infillRequest: InfillRequest, - private val onSend: (element: CodeCompletionTextElement) -> Unit, -) : CodeCompletionEventListener(editor) { - - private var isFirstLine = true - private val currentLineBuffer = StringBuilder() - private val incomingTextBuffer = StringBuilder() - - 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) - } - } - - override fun handleCompleted(messageBuilder: StringBuilder) { - if (incomingTextBuffer.isNotEmpty()) { - appendRemainingCompletion(incomingTextBuffer.toString()) + private fun parseOutput(input: String): String { + if (!service().state.codeCompletionSettings.treeSitterProcessingEnabled) { + return input } - if (isFirstLine) { - val completionLine = messageBuilder.toString().adjustWhitespaces(editor) - REMAINING_EDITOR_COMPLETION.set(editor, completionLine) - onLineReceived(completionLine) - } + return CodeCompletionParserFactory + .getParserForFileExtension(editor.virtualFile.extension) + .parse(prefix, suffix, (firstLine ?: "") + input) + .trimEnd() } - - private fun processCompletionLine(line: String) { - currentLineBuffer.append(line) - - if (currentLineBuffer.trim().isNotEmpty()) { - val completionText = if (isFirstLine) { - line.adjustWhitespaces(editor).also { - isFirstLine = false - onLineReceived(it) - } - } else { - currentLineBuffer.toString() - } - - appendRemainingCompletion(completionText) - currentLineBuffer.clear() - } - } - - private fun onLineReceived(completionLine: String) { - runInEdt { - var editorLineSuffix = editor.getLineSuffixAfterCaret() - if (editorLineSuffix.isBlank()) { - onSend( - CodeCompletionTextElement( - completionLine, - infillRequest.caretOffset, - TextRange.from(infillRequest.caretOffset, completionLine.length), - ) - ) - } else { - var caretShift = 0 - - // TODO: Handle other scenarios - val processedCompletion = - if (completionLine.startsWith(editorLineSuffix.first())) { - caretShift++ - editorLineSuffix = editorLineSuffix.substring(1) - completionLine.substring(1) - } else { - completionLine - } - - val completionWithRemovedSuffix = - processedCompletion.removeSuffix(editorLineSuffix) - - onSend( - CodeCompletionTextElement( - completionWithRemovedSuffix, - infillRequest.caretOffset + caretShift, - TextRange.from( - infillRequest.caretOffset + caretShift, - completionWithRemovedSuffix.length - ), - caretShift, - completionLine - ) - ) - } - } - } - - private fun appendRemainingCompletion(text: String) { - val previousRemainingText = REMAINING_EDITOR_COMPLETION.get(editor) ?: "" - REMAINING_EDITOR_COMPLETION.set(editor, previousRemainingText + text) - } - - private fun Editor.getLineSuffixAfterCaret(): String { - val lineEndOffset = document.getLineEndOffset(document.getLineNumber(caretModel.offset)) - return document.getText( - TextRange( - caretModel.offset, - min(lineEndOffset + 1, document.textLength) - ) - ) - } -} +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionFormatter.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionFormatter.kt new file mode 100644 index 00000000..13ed5a00 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionFormatter.kt @@ -0,0 +1,277 @@ +package ee.carlrobert.codegpt.codecompletions + +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 org.apache.commons.text.similarity.LevenshteinDistance +import kotlin.math.min + +class CodeCompletionFormatter(private val editor: Editor) { + + companion object { + private val logger = thisLogger() + + private val OPENING_BRACKETS = listOf('(', '[', '{') + private val CLOSING_BRACKETS = listOf(')', ']', '}') + private val QUOTES = listOf('\'', '"', '`') + private val BRACKET_PAIRS = mapOf( + '(' to ')', + '[' to ']', + '{' to '}' + ) + } + + private val languageId = editor.virtualFile?.fileType?.name + private val cursorPosition = runReadAction { editor.caretModel.offset } + private val document = editor.document + private val lineNumber = document.getLineNumber(cursorPosition) + private val lineStartOffset = document.getLineStartOffset(lineNumber) + private val lineEndOffset = document.getLineEndOffset(lineNumber) + private val textAfterCursor = document.getText(TextRange(cursorPosition, lineEndOffset)) + private val charAfterCursor = if (textAfterCursor.isNotEmpty()) textAfterCursor[0] else ' ' + private val charBeforeCursor = if (cursorPosition > lineStartOffset) + document.getText(TextRange(cursorPosition - 1, cursorPosition))[0] else ' ' + private var completion = "" + private var normalizedCompletion = "" + private var originalCompletion = "" + private var isDebugEnabled = false + + fun withDebug(): CodeCompletionFormatter { + isDebugEnabled = true + return this + } + + fun format(completion: String): String { + this.completion = "" + this.normalizedCompletion = completion.trim() + this.originalCompletion = completion + + return matchCompletionBrackets() + .removeSuffix() + .removeDuplicateQuotes() + .removeMiddleQuotes() + .ignoreBlankLines() + .removeOverlapText() + .trimStart() + .preventDuplicates() + .getCompletion() + } + + private fun isMatchingPair(open: Char?, close: Char?): Boolean { + return BRACKET_PAIRS[open] == close + } + + private fun removeSuffix(): CodeCompletionFormatter { + completion = completion.removeSuffix(textAfterCursor) + return this + } + + private fun matchCompletionBrackets(): CodeCompletionFormatter { + var accumulatedCompletion = "" + val openBrackets = mutableListOf() + var inString = false + var stringChar = ' ' + + for (char in originalCompletion) { + if (char in QUOTES) { + if (!inString) { + inString = true + stringChar = char + } else if (char == stringChar) { + inString = false + stringChar = ' ' + } + } + + if (!inString) { + if (char in OPENING_BRACKETS) { + openBrackets.add(char) + } else if (char in CLOSING_BRACKETS) { + val lastOpen = openBrackets.lastOrNull() + if (lastOpen != null && isMatchingPair(lastOpen, char)) { + openBrackets.removeAt(openBrackets.size - 1) + } else { + break + } + } + } + + accumulatedCompletion += char + } + + completion = accumulatedCompletion.trimEnd().ifEmpty { originalCompletion.trimEnd() } + + if (isDebugEnabled) { + logger.info("After matchCompletionBrackets: $completion") + } + + return this + } + + private fun ignoreBlankLines(): CodeCompletionFormatter { + if (completion.trimStart().isEmpty() && originalCompletion != "\n") { + completion = completion.trim() + } + + if (isDebugEnabled) { + logger.info("After ignoreBlankLines: $completion") + } + + return this + } + + private fun removeOverlapText(): CodeCompletionFormatter { + val after = textAfterCursor.trim() + if (after.isEmpty() || completion.isEmpty()) return this + + val maxLength = min(completion.length, after.length) + var overlapLength = 0 + + for (length in maxLength downTo 1) { + val endOfCompletion = completion.takeLast(length) + val startOfAfter = after.take(length) + if (endOfCompletion == startOfAfter) { + overlapLength = length + break + } + } + + if (overlapLength > 0) { + completion = completion.dropLast(overlapLength) + } + + if (isDebugEnabled) { + logger.info("After removeDuplicateText: $completion") + } + + return this + } + + private fun isCursorAtMiddleOfWord(): Boolean { + val isAfterWord = charAfterCursor.toString().matches(Regex("\\w")) + val isBeforeWord = charBeforeCursor.toString().matches(Regex("\\w")) + + if (!isAfterWord || !isBeforeWord) return false + + if (languageId?.lowercase() in listOf("javascript", "typescript", "php")) { + if (charBeforeCursor == '$' || charAfterCursor == '$') { + return true + } + } + + if (charBeforeCursor == '_' || charAfterCursor == '_') { + return true + } + + return true + } + + private fun removeMiddleQuotes(): CodeCompletionFormatter { + if (isCursorAtMiddleOfWord()) { + if (completion.isNotEmpty() && completion[0] in QUOTES) { + completion = completion.substring(1) + } + + if (completion.isNotEmpty() && completion.last() in QUOTES) { + completion = completion.dropLast(1) + } + } + + if (isDebugEnabled) { + logger.info("After removeUnnecessaryMiddleQuotes: $completion") + } + + return this + } + + private fun isSimilarCode(s1: String, s2: String): Double { + val distance = LevenshteinDistance.getDefaultInstance().apply(s1, s2) + val maxLength = maxOf(s1.length, s2.length) + return 1.0 - (distance.toDouble() / maxLength) + } + + private fun removeDuplicateQuotes(): CodeCompletionFormatter { + val trimmedCharAfterCursor = charAfterCursor.toString().trim() + val normalizedCompletion = completion.trim() + val lastCharOfCompletion = + if (normalizedCompletion.isNotEmpty()) normalizedCompletion.last() else ' ' + + if (trimmedCharAfterCursor.isNotEmpty() && + (normalizedCompletion.endsWith("',") || + normalizedCompletion.endsWith("\",") || + normalizedCompletion.endsWith("`,") || + (normalizedCompletion.endsWith(",") && trimmedCharAfterCursor[0] in QUOTES)) + ) { + completion = completion.dropLast(2) + } else if ((normalizedCompletion.endsWith("'") || + normalizedCompletion.endsWith("\"") || + normalizedCompletion.endsWith("`")) && + trimmedCharAfterCursor.isNotEmpty() && trimmedCharAfterCursor[0] in QUOTES + ) { + completion = completion.dropLast(1) + } else if (lastCharOfCompletion in QUOTES && + trimmedCharAfterCursor.isNotEmpty() && + trimmedCharAfterCursor[0] == lastCharOfCompletion + ) { + completion = completion.dropLast(1) + } + + if (isDebugEnabled) { + logger.info("After removeDuplicateQuotes: $completion") + } + + return this + } + + private fun preventDuplicates(): CodeCompletionFormatter { + val lineCount = document.lineCount + val originalNormalized = originalCompletion.trim() + + for (i in 1..3) { + val nextLineIndex = lineNumber + i + if (nextLineIndex >= lineCount) break + + val nextLineStartOffset = document.getLineStartOffset(nextLineIndex) + val nextLineEndOffset = document.getLineEndOffset(nextLineIndex) + val nextLine = document.getText(TextRange(nextLineStartOffset, nextLineEndOffset)) + val nextLineNormalized = nextLine.trim() + + if (nextLineNormalized == originalNormalized) { + completion = "" + break + } + + if (isSimilarCode(nextLineNormalized, originalNormalized) > 0.8) { + completion = "" + break + } + } + + if (isDebugEnabled) { + logger.info("After preventDuplicateLine: $completion") + } + + return this + } + + private fun getCompletion(): String { + if (completion.trim().isEmpty()) { + completion = "" + } + return completion + } + + private fun trimStart(): CodeCompletionFormatter { + val firstNonSpaceIndex = completion.indexOfFirst { !it.isWhitespace() } + if (firstNonSpaceIndex > 0 && (cursorPosition - lineStartOffset) <= firstNonSpaceIndex) { + completion = completion.trimStart() + } + + if (isDebugEnabled) { + logger.info("After trimStart: $completion") + } + + return this + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt deleted file mode 100644 index 79caee60..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt +++ /dev/null @@ -1,145 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions - -import ai.grazie.nlp.utils.takeWhitespaces -import com.intellij.codeInsight.hint.HintManagerImpl -import com.intellij.codeInsight.inline.completion.InlineCompletion -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.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.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 { - - 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 } - - if (elements.isEmpty()) { - val textToInsert = context.textToInsert() - val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) ?: "" - if (remainingCompletion.isNotEmpty()) { - REMAINING_EDITOR_COMPLETION.set( - editor, - remainingCompletion.removePrefix(textToInsert) - ) - } - - InlineCompletion.getHandlerOrNull(editor)?.insert() - - if (GeneralSettings.getSelectedService() == ServiceType.CODEGPT - && service().state.nextEditsEnabled - ) { - ApplicationManager.getApplication().executeOnPooledThread { - service().displayInlineDiff(editor) - } - return - } - } - - 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 4070c051..45f8b522 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt @@ -5,7 +5,16 @@ 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 ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION +import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.application.runReadAction +import com.intellij.openapi.components.service +import ee.carlrobert.codegpt.CodeGPTKeys +import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService +import ee.carlrobert.codegpt.predictions.CodeSuggestionDiffViewer +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.launch class CodeCompletionInsertHandler : InlineCompletionInsertHandler { @@ -14,11 +23,45 @@ class CodeCompletionInsertHandler : InlineCompletionInsertHandler { elements: List ) { val editor = environment.editor - val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) ?: "" - if (remainingCompletion.isNotEmpty()) { + val remainingCompletion = CodeGPTKeys.REMAINING_CODE_COMPLETION.get(editor) + if (remainingCompletion != null && remainingCompletion.partialCompletion.isNotEmpty()) { InlineCompletion.getHandlerOrNull(editor)?.invoke( InlineCompletionEvent.DirectCall(editor, editor.caretModel.currentCaret) ) + val caretOffset = runReadAction { editor.caretModel.offset } + val prefix = editor.document.text.substring(0, caretOffset) + val suffix = editor.document.text.substring(caretOffset) + CoroutineScope(SupervisorJob() + Dispatchers.IO).launch { + editor.project?.service() + ?.getNextEdit( + editor, + prefix + remainingCompletion.partialCompletion + suffix, + caretOffset + remainingCompletion.partialCompletion.length, + true + ) + } + return + } else { + if (CodeGPTKeys.REMAINING_PREDICTION_RESPONSE.get(editor) == null) { + val caretOffset = runReadAction { editor.caretModel.offset } + CoroutineScope(SupervisorJob() + Dispatchers.IO).launch { + editor.project?.service()?.getNextEdit( + editor, + editor.document.text, + caretOffset, + ) + } + return + } + } + + CoroutineScope(SupervisorJob() + Dispatchers.IO).launch { + val queuedPrediction = CodeGPTKeys.REMAINING_PREDICTION_RESPONSE.get(editor) + if (queuedPrediction != null) { + runInEdt { + CodeSuggestionDiffViewer.displayInlineDiff(editor, queuedPrediction) + } + } } } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt index 32f28566..f11f6e44 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt @@ -7,16 +7,15 @@ import ee.carlrobert.codegpt.completions.llama.LlamaModel import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential import ee.carlrobert.codegpt.settings.Placeholder.* -import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings import ee.carlrobert.codegpt.settings.service.llama.LlamaSettingsState import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings -import ee.carlrobert.llm.client.codegpt.request.CodeCompletionRequest import ee.carlrobert.llm.client.llama.completion.LlamaCompletionRequest import ee.carlrobert.llm.client.ollama.completion.request.OllamaCompletionRequest import ee.carlrobert.llm.client.ollama.completion.request.OllamaParameters import ee.carlrobert.llm.client.openai.completion.request.OpenAITextCompletionRequest +import ee.carlrobert.service.GrpcCodeCompletionRequest import okhttp3.MediaType.Companion.toMediaType import okhttp3.Request import okhttp3.RequestBody.Companion.toRequestBody @@ -27,15 +26,12 @@ object CodeCompletionRequestFactory { private const val MAX_TOKENS = 128 @JvmStatic - fun buildCodeGPTRequest(details: InfillRequest): CodeCompletionRequest { - return CodeCompletionRequest.Builder() - .setModel(service().state.codeCompletionSettings.model) - .setPrefix(details.prefix) - .setSuffix(details.suffix) - .setFileExtension(details.fileDetails?.fileExtension) + fun buildCodeGPTRequest(details: InfillRequest): GrpcCodeCompletionRequest { + return GrpcCodeCompletionRequest.newBuilder() + .setFilePath(details.fileDetails?.filePath) .setFileContent(details.fileDetails?.fileContent) - .setCursorOffset(details.caretOffset) - .setStop(details.stopTokens.ifEmpty { null }) + .setGitDiff("") + .setCursorPosition(details.caretOffset) .build() } @@ -55,7 +51,8 @@ object CodeCompletionRequestFactory { fun buildCustomRequest(details: InfillRequest): Request { val activeService = service().state.active val settings = activeService.codeCompletionSettings - val credential = getCredential(CredentialKey.CustomServiceApiKey(activeService.name.orEmpty())) + val credential = + getCredential(CredentialKey.CustomServiceApiKey(activeService.name.orEmpty())) return buildCustomRequest( details, settings.url!!, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt index aa265081..68441f03 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt @@ -2,11 +2,13 @@ package ee.carlrobert.codegpt.codecompletions import com.intellij.openapi.components.Service import com.intellij.openapi.components.service +import com.intellij.openapi.project.Project import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildCodeGPTRequest import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildCustomRequest import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildLlamaRequest import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildOllamaRequest import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory.buildOpenAIRequest +import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService import ee.carlrobert.codegpt.completions.CompletionClientProvider import ee.carlrobert.codegpt.completions.llama.LlamaModel import ee.carlrobert.codegpt.settings.GeneralSettings @@ -20,11 +22,12 @@ import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionEventSourceListener import ee.carlrobert.llm.client.openai.completion.OpenAITextCompletionEventSourceListener import ee.carlrobert.llm.completion.CompletionEventListener +import okhttp3.Request import okhttp3.sse.EventSource import okhttp3.sse.EventSources.createFactory @Service(Service.Level.PROJECT) -class CodeCompletionService { +class CodeCompletionService(private val project: Project) { // TODO: Consolidate logic in ModelComboBoxAction fun getSelectedModelCode(): String? { @@ -43,6 +46,8 @@ class CodeCompletionService { } } + fun isCodeCompletionsEnabled(): Boolean = isCodeCompletionsEnabled(GeneralSettings.getSelectedService()) + fun isCodeCompletionsEnabled(selectedService: ServiceType): Boolean = when (selectedService) { CODEGPT -> service().state.codeCompletionSettings.codeCompletionsEnabled @@ -56,11 +61,8 @@ class CodeCompletionService { fun getCodeCompletionAsync( infillRequest: InfillRequest, eventListener: CompletionEventListener - ): EventSource = - when (val selectedService = GeneralSettings.getSelectedService()) { - CODEGPT -> CompletionClientProvider.getCodeGPTClient() - .getCodeCompletionAsync(buildCodeGPTRequest(infillRequest), eventListener) - + ): EventSource { + return when (val selectedService = GeneralSettings.getSelectedService()) { OPENAI -> CompletionClientProvider.getOpenAIClient() .getCompletionAsync(buildOpenAIRequest(infillRequest), eventListener) @@ -83,4 +85,5 @@ class CodeCompletionService { else -> throw IllegalArgumentException("Code completion not supported for ${selectedService.name}") } + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt index 83db2a6f..7bdb9eba 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt @@ -11,6 +11,7 @@ class CodeCompletionTextElement( val textRange: TextRange, val offsetDelta: Int = 0, val originalText: String = text, + val isDone: Boolean = false, ) : InlineCompletionElement { override fun toPresentable(): InlineCompletionElement.Presentable = diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt index 695b49bc..18c97350 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt @@ -1,31 +1,21 @@ 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.lookup.LookupManager -import com.intellij.notification.NotificationType import com.intellij.openapi.components.service -import com.intellij.openapi.diagnostic.thisLogger -import com.intellij.openapi.project.Project -import com.intellij.openapi.util.TextRange -import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION -import ee.carlrobert.codegpt.predictions.PredictionService +import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_CODE_COMPLETION +import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService import ee.carlrobert.codegpt.settings.GeneralSettings -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings import ee.carlrobert.codegpt.settings.service.custom.CustomServicesSettings 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.ui.OverlayUtil -import ee.carlrobert.codegpt.util.StringUtil.extractUntilNewline -import kotlinx.coroutines.channels.ProducerScope import kotlinx.coroutines.channels.awaitClose import kotlinx.coroutines.flow.channelFlow import kotlinx.coroutines.flow.emptyFlow -import kotlinx.coroutines.launch import okhttp3.sse.EventSource import java.util.concurrent.atomic.AtomicReference import kotlin.time.Duration @@ -34,10 +24,6 @@ import kotlin.time.toDuration class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { - companion object { - private val logger = thisLogger() - } - private val currentCallRef = AtomicReference(null) override val id: InlineCompletionProviderID @@ -49,112 +35,68 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { override val providerPresentation: InlineCompletionProviderPresentation get() = CodeCompletionProviderPresentation() + override fun shouldBeForced(request: InlineCompletionRequest): Boolean { + return request.event is InlineCompletionEvent.DirectCall || tryFindCache(request) != null + } + override suspend fun getSuggestionDebounced(request: InlineCompletionRequest): InlineCompletionSuggestion { - val codegptSettings = service().state - if (GeneralSettings.getSelectedService() == ServiceType.CODEGPT && codegptSettings.nextEditsEnabled) { - if (codegptSettings.codeCompletionSettings.codeCompletionsEnabled) { - codegptSettings.codeCompletionSettings.codeCompletionsEnabled = false - OverlayUtil.showNotification( - "Code completions and multi-line edits cannot be active simultaneously.", - NotificationType.WARNING - ) - } - - predictNextEdit(request) - return InlineCompletionSuggestion.Default(emptyFlow()) - } - - return if (service().state.codeCompletionSettings.multiLineEnabled) { - getMultiLineSuggestionDebounced(request) - } else { - getSingleLineSuggestionDebounced(request) - } - } - - private fun predictNextEdit(request: InlineCompletionRequest) { - val project = request.editor.project ?: return - try { - CompletionProgressNotifier.update(project, true) - project.service().displayInlineDiff(request.editor) - } catch (ex: Exception) { - logger.error("Error communicating with server: ${ex.message}") - } - } - - private fun getSingleLineSuggestionDebounced(request: InlineCompletionRequest): InlineCompletionSuggestion { val editor = request.editor - val remainingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) ?: "" - if (request.event is InlineCompletionEvent.DirectCall && remainingCompletion.isNotEmpty() - ) { - return sendNextSuggestion(remainingCompletion.extractUntilNewline(), request) - } - - return getSuggestionDebounced( - request, - CompletionType.SINGLE_LINE - ) { project, infillRequest -> - project.service() - .getCodeCompletionAsync( - infillRequest, - CodeCompletionSingleLineEventListener(request.editor, infillRequest) { - trySend(it) - } - ) - } - } - - private fun getMultiLineSuggestionDebounced(request: InlineCompletionRequest): InlineCompletionSuggestion { - return getSuggestionDebounced( - request, - CompletionType.MULTI_LINE - ) { project, infillRequest -> - project.service() - .getCodeCompletionAsync( - infillRequest, - CodeCompletionMultiLineEventListener(request) { - if (LookupManager.getActiveLookup(request.editor) == null) { - trySend(InlineCompletionGrayTextElement(it)) - } - } - ) - } - } - - private fun getSuggestionDebounced( - request: InlineCompletionRequest, - completionType: CompletionType, - fetchCompletion: ProducerScope.(Project, InfillRequest) -> EventSource - ): InlineCompletionSuggestion { - val project = request.editor.project - if (project == null) { - logger.error("Could not find project") - return InlineCompletionSuggestion.Default(emptyFlow()) - } + val project = + editor.project ?: return InlineCompletionSuggestion.Default(emptyFlow()) if (LookupManager.getActiveLookup(request.editor) != null) { return InlineCompletionSuggestion.Default(emptyFlow()) } - request.editor.project?.let { - CompletionProgressNotifier.update(it, true) - } - return InlineCompletionSuggestion.Default(channelFlow { - val infillRequest = InfillRequestUtil.buildInfillRequest(request, completionType) - currentCallRef.set(fetchCompletion(project, infillRequest)) - awaitClose { currentCallRef.getAndSet(null)?.cancel() } + try { + val remainingCodeCompletion = REMAINING_CODE_COMPLETION.get(editor) + if (remainingCodeCompletion != null && request.event is InlineCompletionEvent.DirectCall) { + REMAINING_CODE_COMPLETION.set(editor, null) + trySend(InlineCompletionGrayTextElement(remainingCodeCompletion.partialCompletion)) + return@channelFlow + } + + val cacheValue = tryFindCache(request) + if (cacheValue != null) { + REMAINING_CODE_COMPLETION.set(editor, null) + trySend(InlineCompletionGrayTextElement(cacheValue)) + return@channelFlow + } + + CompletionProgressNotifier.update(project, true) + + var eventListener = CodeCompletionEventListener(request.editor, this) + + if (GeneralSettings.getSelectedService() == ServiceType.CODEGPT) { + project.service().getCodeCompletionAsync(eventListener, request, this) + return@channelFlow + } + + val infillRequest = InfillRequestUtil.buildInfillRequest(request) + val call = project.service().getCodeCompletionAsync( + infillRequest, + CodeCompletionEventListener(request.editor, this) + ) + + currentCallRef.set(call) + } finally { + awaitClose { currentCallRef.getAndSet(null)?.cancel() } + } }) } + private fun tryFindCache(request: InlineCompletionRequest): String? { + val editor = request.editor + val project = editor.project ?: return null + return project.service().getCache(editor) + } + override suspend fun getDebounceDelay(request: InlineCompletionRequest): Duration { - return 400.toDuration(DurationUnit.MILLISECONDS) + return 300.toDuration(DurationUnit.MILLISECONDS) } override fun isEnabled(event: InlineCompletionEvent): Boolean { - if (LookupManager.getActiveLookup(event.toRequest()?.editor) != null) { - return false - } - val selectedService = GeneralSettings.getSelectedService() val codeCompletionsEnabled = when (selectedService) { ServiceType.CODEGPT -> service().state.codeCompletionSettings.codeCompletionsEnabled @@ -167,8 +109,13 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { ServiceType.GOOGLE, null -> false } + + if (event is LookupInlineCompletionEvent) { + return true + } + val hasActiveCompletion = - REMAINING_EDITOR_COMPLETION.get(event.toRequest()?.editor)?.isNotEmpty() ?: false + REMAINING_CODE_COMPLETION.get(event.toRequest()?.editor)?.partialCompletion?.isNotEmpty() == true if (!codeCompletionsEnabled) { return event is InlineCompletionEvent.DocumentChange @@ -177,23 +124,6 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { && !hasActiveCompletion } - return event is InlineCompletionEvent.DocumentChange || hasActiveCompletion + return event is InlineCompletionEvent.DocumentChange || hasActiveCompletion } - - private fun sendNextSuggestion( - nextCompletion: String, - request: InlineCompletionRequest - ): InlineCompletionSuggestion { - return InlineCompletionSuggestion.Default(channelFlow { - launch { - trySend( - 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 9ed086a2..175c511c 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt @@ -22,7 +22,7 @@ class InfillRequest private constructor( val stopTokens: List, ) { - data class FileDetails(val fileContent: String, val fileExtension: String? = null) + data class FileDetails(val fileContent: String, val filePath: String? = null) class Builder { private val prefix: String @@ -39,18 +39,16 @@ class InfillRequest private constructor( prefix: String, suffix: String, caretOffset: Int, - type: CompletionType = CompletionType.MULTI_LINE ) { this.prefix = prefix this.suffix = suffix this.caretOffset = caretOffset - this.stopTokens = getStopTokens(type) + this.stopTokens = getStopTokens() } constructor( document: Document, caretOffset: Int, - type: CompletionType = CompletionType.MULTI_LINE ) { prefix = document.getText(TextRange(0, caretOffset)) @@ -59,7 +57,7 @@ class InfillRequest private constructor( document.getText(TextRange(caretOffset, document.textLength)) .truncateText(MAX_PROMPT_TOKENS) this.caretOffset = caretOffset - this.stopTokens = getStopTokens(type) + this.stopTokens = getStopTokens() } fun fileDetails(fileDetails: FileDetails) = apply { this.fileDetails = fileDetails } @@ -74,7 +72,7 @@ class InfillRequest private constructor( fun context(context: InfillContext) = apply { this.context = context } - private fun getStopTokens(type: CompletionType): List { + private fun getStopTokens(): List { var whitespaceCount = 0 val lineSuffix = suffix .takeWhile { char -> @@ -82,10 +80,7 @@ class InfillRequest private constructor( else if (char.isWhitespace()) whitespaceCount++ < 2 else whitespaceCount < 2 } - val baseTokens = when (type) { - CompletionType.SINGLE_LINE -> emptyList() - else -> listOf("\n\n") - } + val baseTokens = listOf("\n\n") return if (lineSuffix.isNotEmpty()) { baseTokens + lineSuffix @@ -116,7 +111,6 @@ class InfillRequest private constructor( class InfillContext( val enclosingElement: ContextElement, - // TODO: Add some kind of ranking, which contextElements are more important than others val contextElements: Set ) { @@ -132,9 +126,4 @@ class ContextElement(val psiElement: PsiElement) { fun String.truncateText(maxTokens: Int, fromStart: Boolean = true): String { return service().truncateText(this, maxTokens, fromStart) -} - -enum class CompletionType { - SINGLE_LINE, - MULTI_LINE, -} +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt index 7c0d488b..9639a1e1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt @@ -14,16 +14,13 @@ import ee.carlrobert.codegpt.util.GitUtil object InfillRequestUtil { - suspend fun buildInfillRequest( - request: InlineCompletionRequest, - type: CompletionType - ): InfillRequest { + suspend fun buildInfillRequest(request: InlineCompletionRequest): InfillRequest { val caretOffset = readAction { request.editor.caretModel.offset } - val infillRequestBuilder = InfillRequest.Builder(request.document, caretOffset, type) + val infillRequestBuilder = InfillRequest.Builder(request.document, caretOffset) .fileDetails( InfillRequest.FileDetails( request.document.text, - request.file.virtualFile.extension + request.file.virtualFile.path ) ) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/LookupInlineCompletionEvent.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/LookupInlineCompletionEvent.kt new file mode 100644 index 00000000..d15ff1e9 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/LookupInlineCompletionEvent.kt @@ -0,0 +1,54 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.intellij.codeInsight.inline.completion.InlineCompletionEvent +import com.intellij.codeInsight.inline.completion.InlineCompletionRequest +import com.intellij.codeInsight.lookup.LookupEvent +import com.intellij.openapi.application.runReadAction +import com.intellij.openapi.editor.Caret +import com.intellij.openapi.project.Project +import com.intellij.psi.PsiDocumentManager +import com.intellij.psi.PsiFile +import com.intellij.psi.impl.source.PsiFileImpl +import com.intellij.psi.util.PsiUtilBase + +class LookupInlineCompletionEvent(private val event: LookupEvent) : InlineCompletionEvent { + + override fun toRequest(): InlineCompletionRequest? { + val editor = runReadAction { event.lookup?.editor } ?: return null + val caretModel = editor.caretModel + if (caretModel.caretCount != 1) return null + + val project = editor.project ?: return null + + val (file, offset) = runReadAction { + getPsiFile(caretModel.currentCaret, project) to caretModel.offset + } + if (file == null) return null + + return InlineCompletionRequest( + this, + file, + editor, + editor.document, + offset, + offset, + event.item + ) + } + + private fun getPsiFile(caret: Caret, project: Project): PsiFile? { + return runReadAction { + val file = PsiDocumentManager.getInstance(project).getPsiFile(caret.editor.document) + ?: return@runReadAction null + if (file.isLoadedInMemory()) { + PsiUtilBase.getPsiFileInEditor(caret, project) + } else { + file + } + } + } + + private fun PsiFile.isLoadedInMemory(): Boolean { + return (this as? PsiFileImpl)?.treeElement != null + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt new file mode 100644 index 00000000..0d35dc07 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt @@ -0,0 +1,46 @@ +package ee.carlrobert.codegpt.codecompletions.edit + +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement +import com.intellij.openapi.diagnostic.thisLogger +import com.intellij.openapi.editor.Editor +import ee.carlrobert.codegpt.codecompletions.CodeCompletionEventListener +import ee.carlrobert.service.GrpcCodeCompletionRequest +import ee.carlrobert.service.PartialCodeCompletionResponse +import io.grpc.stub.StreamObserver +import kotlinx.coroutines.channels.ProducerScope +import okhttp3.Request +import okhttp3.sse.EventSource + +class CodeCompletionStreamObserver( + private val channel: ProducerScope, + private val eventListener: CodeCompletionEventListener, +) : StreamObserver { + + companion object { + private val logger = thisLogger() + } + + private val messageBuilder = StringBuilder() + private val emptyEventSource = object : EventSource { + override fun cancel() { + } + + override fun request(): Request { + return Request.Builder().build() + } + } + + override fun onNext(value: PartialCodeCompletionResponse) { + messageBuilder.append(value.partialCompletion) + eventListener.onMessage(value.partialCompletion, emptyEventSource) + } + + override fun onError(t: Throwable?) { + logger.error("Error occurred while fetching code completion", t) + channel.close(t) + } + + override fun onCompleted() { + eventListener.onComplete(messageBuilder) + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcCallCredentials.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcCallCredentials.kt new file mode 100644 index 00000000..c5f1d9aa --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcCallCredentials.kt @@ -0,0 +1,29 @@ +package ee.carlrobert.codegpt.codecompletions.edit + +import io.grpc.CallCredentials +import io.grpc.Metadata +import io.grpc.Status +import java.util.concurrent.Executor + +class GrpcCallCredentials(private val apiKey: String) : CallCredentials() { + + companion object { + private val API_KEY_HEADER = Metadata.Key.of("x-api-key", Metadata.ASCII_STRING_MARSHALLER) + } + + override fun applyRequestMetadata( + requestInfo: RequestInfo?, + executor: Executor, + metadataApplier: MetadataApplier + ) { + executor.execute { + try { + val headers = Metadata() + headers.put(API_KEY_HEADER, apiKey) + metadataApplier.apply(headers) + } catch (e: Throwable) { + metadataApplier.fail(Status.UNAUTHENTICATED.withCause(e)) + } + } + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt index dabcddb4..57ec70d9 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt @@ -1,10 +1,8 @@ package ee.carlrobert.codegpt.codecompletions.edit -import com.intellij.codeInsight.lookup.LookupManager -import com.intellij.notification.NotificationAction.createSimpleExpiring -import com.intellij.notification.NotificationType +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.runInEdt import com.intellij.openapi.application.runReadAction import com.intellij.openapi.components.Service import com.intellij.openapi.components.service @@ -13,32 +11,30 @@ import com.intellij.openapi.editor.Editor import com.intellij.openapi.project.Project import com.intellij.util.net.ssl.CertificateManager import com.jetbrains.rd.util.UUID -import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier +import ee.carlrobert.codegpt.codecompletions.CodeCompletionEventListener import ee.carlrobert.codegpt.credentials.CredentialsStore import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CodeGptApiKey -import ee.carlrobert.codegpt.predictions.CodeSuggestionDiffViewer +import ee.carlrobert.codegpt.settings.GeneralSettings +import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings import ee.carlrobert.codegpt.telemetry.core.configuration.TelemetryConfiguration -import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.codegpt.util.GitUtil -import ee.carlrobert.service.AcceptEditRequest -import ee.carlrobert.service.NextEditRequest -import ee.carlrobert.service.NextEditResponse -import ee.carlrobert.service.NextEditServiceImplGrpc -import io.grpc.* +import ee.carlrobert.service.* +import io.grpc.ManagedChannel import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder -import io.grpc.stub.StreamObserver -import java.util.concurrent.Executor +import kotlinx.coroutines.channels.ProducerScope import java.util.concurrent.TimeUnit -import kotlin.coroutines.cancellation.CancellationException @Service(Service.Level.PROJECT) class GrpcClientService(private val project: Project) : Disposable { private var channel: ManagedChannel? = null - private var stub: NextEditServiceImplGrpc.NextEditServiceImplStub? = null - private var prevObserver: NextEditStreamObserver? = null + private var codeCompletionStub: CodeCompletionServiceImplGrpc.CodeCompletionServiceImplStub? = + null + private var codeCompletionObserver: CodeCompletionStreamObserver? = null + private var nextEditStub: NextEditServiceImplGrpc.NextEditServiceImplStub? = null + private var nextEditStreamObserver: NextEditStreamObserver? = null companion object { private const val HOST = "grpc.tryproxy.io" @@ -48,95 +44,44 @@ class GrpcClientService(private val project: Project) : Disposable { private val logger = thisLogger() } - @Synchronized - private fun ensureConnection() { - if (channel == null || channel?.isShutdown == true) { - try { - channel = NettyChannelBuilder.forAddress(HOST, PORT) - .useTransportSecurity() - .sslContext(GrpcSslContexts.forClient() - .trustManager(CertificateManager.getInstance().trustManager) - .build()) - .build() - stub = NextEditServiceImplGrpc.newStub(channel) - .withCallCredentials( - ApiKeyCredentials(CredentialsStore.getCredential(CodeGptApiKey) ?: "") - ) + fun getCodeCompletionAsync( + eventListener: CodeCompletionEventListener, + request: InlineCompletionRequest, + channel: ProducerScope + ) { + ensureCodeCompletionConnection() - logger.info("gRPC connection established") - } catch (e: Exception) { - logger.error("Failed to establish gRPC connection", e) - throw e - } - } + val grpcRequest = createCodeCompletionGrpcRequest(request) + codeCompletionObserver = CodeCompletionStreamObserver(channel, eventListener) + codeCompletionStub?.getCodeCompletion(grpcRequest, codeCompletionObserver) } - fun getNextEdit(editor: Editor, isManuallyOpened: Boolean = false) { - ensureConnection() - prevObserver?.onCompleted() - - val request = NextEditRequest.newBuilder() - .setFileName(editor.virtualFile.name) - .setFileContent(editor.document.text) - .setGitDiff(GitUtil.getCurrentChanges(project) ?: "") - .setCursorPosition(runReadAction { editor.caretModel.offset }) - .setEnableTelemetry(TelemetryConfiguration.getInstance().isCompletionTelemetryEnabled) - .build() - prevObserver = NextEditStreamObserver(editor, isManuallyOpened) { - dispose() + fun getNextEdit( + editor: Editor, + fileContent: String, + caretOffset: Int, + addToQueue: Boolean = false, + ) { + if (GeneralSettings.getSelectedService() != ServiceType.CODEGPT + || !service().state.nextEditsEnabled + ) { + return } - stub?.nextEdit(request, prevObserver) - } + ensureNextEditConnection() - class NextEditStreamObserver( - private val editor: Editor, - private val isManuallyOpened: Boolean, - private val onDispose: () -> Unit - ) : StreamObserver { - override fun onNext(response: NextEditResponse) { - runInEdt { - val documentText = editor.document.text - if (LookupManager.getActiveLookup(editor) == null - && documentText != response.nextRevision - && documentText == response.oldRevision) { - CodeSuggestionDiffViewer.displayInlineDiff(editor, response, isManuallyOpened) - } - } - } - - override fun onError(ex: Throwable) { - if (ex is CancellationException || - (ex is StatusRuntimeException && ex.status.code == Status.Code.CANCELLED) - ) { - onCompleted() - return - } - - try { - if (ex is StatusRuntimeException) { - OverlayUtil.showNotification( - ex.status.description ?: ex.localizedMessage, - NotificationType.ERROR, - createSimpleExpiring("Disable multi-line edits") { - service().state.nextEditsEnabled = - false - }) - } else { - logger.error("Something went wrong", ex) - } - } finally { - onCompleted() - onDispose() - } - } - - override fun onCompleted() { - editor.project?.let { CompletionProgressNotifier.update(it, false) } - } + val request = createNextEditGrpcRequest(editor, fileContent, caretOffset) + nextEditStreamObserver = NextEditStreamObserver(editor, addToQueue) { dispose() } + nextEditStub?.nextEdit(request, nextEditStreamObserver) } fun acceptEdit(responseId: UUID, acceptedEdit: String) { + if (GeneralSettings.getSelectedService() != ServiceType.CODEGPT + || !TelemetryConfiguration.getInstance().isCompletionTelemetryEnabled + ) { + return + } + NextEditServiceImplGrpc .newBlockingStub(channel) .acceptEdit( @@ -147,6 +92,92 @@ class GrpcClientService(private val project: Project) : Disposable { ) } + @Synchronized + fun refreshConnection() { + channel?.let { + if (!it.isShutdown) { + try { + it.shutdown().awaitTermination(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS) + logger.info("Existing gRPC connection closed for refresh") + } catch (e: InterruptedException) { + logger.warn("Interrupted while shutting down gRPC channel for refresh", e) + Thread.currentThread().interrupt() + } finally { + if (!it.isTerminated) { + it.shutdownNow() + } + } + } + } + } + + @Synchronized + private fun ensureCodeCompletionConnection() { + ensureActiveChannel() + + if (codeCompletionStub == null) { + codeCompletionStub = CodeCompletionServiceImplGrpc.newStub(channel) + .withCallCredentials(createCallCredentials()) + } + } + + @Synchronized + private fun ensureNextEditConnection() { + ensureActiveChannel() + + if (nextEditStub == null) { + nextEditStub = NextEditServiceImplGrpc.newStub(channel) + .withCallCredentials(createCallCredentials()) + } + } + + private fun createCodeCompletionGrpcRequest(request: InlineCompletionRequest): GrpcCodeCompletionRequest { + val editor = request.editor + return GrpcCodeCompletionRequest.newBuilder() + .setModel(service().state.codeCompletionSettings.model) + .setFilePath(editor.virtualFile.path) + .setFileContent(editor.document.text) + .setGitDiff(GitUtil.getCurrentChanges(project) ?: "") + .setCursorPosition(runReadAction { editor.caretModel.offset }) + .setEnableTelemetry(TelemetryConfiguration.getInstance().isCompletionTelemetryEnabled) + .build() + } + + private fun createNextEditGrpcRequest(editor: Editor, fileContent: String, caretOffset: Int) = + NextEditRequest.newBuilder() + .setFileName(editor.virtualFile.name) + .setFileContent(fileContent) + .setGitDiff(GitUtil.getCurrentChanges(project) ?: "") + .setCursorPosition(caretOffset) + .setEnableTelemetry(TelemetryConfiguration.getInstance().isCompletionTelemetryEnabled) + .build() + + private fun createChannel(): ManagedChannel = NettyChannelBuilder.forAddress(HOST, PORT) + .useTransportSecurity() + .sslContext( + GrpcSslContexts.forClient() + .trustManager(CertificateManager.getInstance().trustManager) + .build() + ) + .build() + + private fun ensureActiveChannel() { + if (channel == null || channel?.isShutdown == true) { + try { + channel = createChannel() + codeCompletionStub = null + nextEditStub = null + logger.info("gRPC connection established") + } catch (e: Exception) { + logger.error("Failed to establish gRPC connection", e) + throw e + } + } + } + + private fun createCallCredentials() = + GrpcCallCredentials(CredentialsStore.getCredential(CodeGptApiKey) ?: "") + override fun dispose() { channel?.let { ch -> if (!ch.isShutdown) { @@ -160,33 +191,9 @@ class GrpcClientService(private val project: Project) : Disposable { if (!ch.isTerminated) { ch.shutdownNow() } - channel = null } } } - } -} - -internal class ApiKeyCredentials(private val apiKey: String) : CallCredentials() { - - companion object { - private val API_KEY_HEADER: Metadata.Key = - Metadata.Key.of("x-api-key", Metadata.ASCII_STRING_MARSHALLER) - } - - override fun applyRequestMetadata( - requestInfo: RequestInfo?, - executor: Executor, - metadataApplier: MetadataApplier - ) { - executor.execute { - try { - val headers = Metadata() - headers.put(API_KEY_HEADER, apiKey) - metadataApplier.apply(headers) - } catch (e: Throwable) { - metadataApplier.fail(Status.UNAUTHENTICATED.withCause(e)) - } - } + channel = null } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/NextEditStreamObserver.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/NextEditStreamObserver.kt new file mode 100644 index 00000000..aef370c2 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/NextEditStreamObserver.kt @@ -0,0 +1,76 @@ +package ee.carlrobert.codegpt.codecompletions.edit + +import com.intellij.codeInsight.lookup.LookupManager +import com.intellij.notification.NotificationAction.createSimpleExpiring +import com.intellij.notification.NotificationType +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 ee.carlrobert.codegpt.CodeGPTKeys +import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier +import ee.carlrobert.codegpt.predictions.CodeSuggestionDiffViewer +import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings +import ee.carlrobert.codegpt.ui.OverlayUtil +import ee.carlrobert.service.NextEditResponse +import io.grpc.Status +import io.grpc.StatusRuntimeException +import io.grpc.stub.StreamObserver +import kotlin.coroutines.cancellation.CancellationException + +class NextEditStreamObserver( + private val editor: Editor, + private val addToQueue: Boolean = false, + private val onDispose: () -> Unit +) : StreamObserver { + + companion object { + private val logger = thisLogger() + } + + override fun onNext(response: NextEditResponse) { + if (addToQueue) { + CodeGPTKeys.REMAINING_PREDICTION_RESPONSE.set(editor, response) + } else { + runInEdt { + val documentText = editor.document.text + if (LookupManager.getActiveLookup(editor) == null + && documentText != response.nextRevision + && documentText == response.oldRevision + ) { + CodeSuggestionDiffViewer.displayInlineDiff(editor, response) + } + } + } + } + + override fun onError(ex: Throwable) { + if (ex is CancellationException || + (ex is StatusRuntimeException && ex.status.code == Status.Code.CANCELLED) + ) { + onCompleted() + return + } + + try { + if (ex is StatusRuntimeException) { + OverlayUtil.showNotification( + ex.status.description ?: ex.localizedMessage, + NotificationType.ERROR, + createSimpleExpiring("Disable multi-line edits") { + service().state.nextEditsEnabled = + false + }) + } else { + logger.error("Something went wrong", ex) + } + } finally { + onCompleted() + onDispose() + } + } + + override fun onCompleted() { + editor.project?.let { CompletionProgressNotifier.update(it, false) } + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/predictions/CodeSuggestionDiffViewer.kt b/src/main/kotlin/ee/carlrobert/codegpt/predictions/CodeSuggestionDiffViewer.kt index 1b6a2e08..05db874f 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/predictions/CodeSuggestionDiffViewer.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/predictions/CodeSuggestionDiffViewer.kt @@ -8,9 +8,9 @@ import com.intellij.diff.requests.SimpleDiffRequest import com.intellij.diff.tools.fragmented.UnifiedDiffChange import com.intellij.diff.tools.fragmented.UnifiedDiffViewer import com.intellij.diff.util.DiffUtil -import com.intellij.ide.plugins.newui.TagComponent import com.intellij.openapi.Disposable import com.intellij.openapi.actionSystem.ActionManager +import com.intellij.openapi.application.runReadAction import com.intellij.openapi.components.service import com.intellij.openapi.editor.Document import com.intellij.openapi.editor.Editor @@ -26,23 +26,18 @@ import com.intellij.openapi.ui.popup.JBPopupFactory import com.intellij.openapi.util.* import com.intellij.testFramework.LightVirtualFile import com.intellij.ui.components.JBLabel -import com.intellij.ui.components.JBScrollPane import com.intellij.util.application import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI import com.intellij.util.ui.components.BorderLayoutPanel -import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.CodeGPTKeys import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService import ee.carlrobert.service.NextEditResponse import java.awt.BorderLayout import java.awt.Dimension -import java.awt.FlowLayout import java.awt.Point import java.util.* -import javax.swing.Box import javax.swing.JComponent -import javax.swing.JPanel import javax.swing.SwingUtilities import kotlin.math.abs import kotlin.math.max @@ -51,12 +46,12 @@ class CodeSuggestionDiffViewer( request: DiffRequest, val nextEditResponse: NextEditResponse, private val mainEditor: Editor, - private val isManuallyOpened: Boolean ) : UnifiedDiffViewer(MyDiffContext(mainEditor.project), request), Disposable { private val popup: JBPopup = createSuggestionDiffPopup(component) private val visibleAreaListener: VisibleAreaListener private val documentListener: DocumentListener + private val grpcService = project?.service() private var applyInProgress = false @@ -85,6 +80,8 @@ class CodeSuggestionDiffViewer( ) adjustPopupSize(popup, myEditor) + updateFooterComponent() + val changeOffset = change.lineFragment.startOffset1 val adjustedLocation = getAdjustedPopupLocation(popup, mainEditor, changeOffset) @@ -121,11 +118,18 @@ class CodeSuggestionDiffViewer( if (changes.size == 1) { popup.dispose() + + application.executeOnPooledThread { + grpcService?.getNextEdit( + mainEditor, + mainEditor.document.text, + runReadAction { mainEditor.caretModel.offset }, + ) + } } application.executeOnPooledThread { - project?.service() - ?.acceptEdit(UUID.fromString(nextEditResponse.id), change.toString()) + grpcService?.acceptEdit(UUID.fromString(nextEditResponse.id), change.toString()) } } @@ -147,8 +151,6 @@ class CodeSuggestionDiffViewer( gutterComponentEx.parent.isVisible = false scrollPane.horizontalScrollBar.isOpaque = false } - - setupStatusLabel() } private fun clearListeners() { @@ -163,62 +165,6 @@ class CodeSuggestionDiffViewer( return changes.minByOrNull { abs(it.lineFragment.startOffset1 - cursorOffset) } } - private fun getTagPanel(): JComponent { - val tagPanel = JPanel(FlowLayout(FlowLayout.LEADING, 0, 0)).apply { - isOpaque = false - } - tagPanel.add( - TagComponent( - "Open: ${getShortcutText(OpenPredictionAction.ID)}" - ).apply { - setListener({ _, _ -> - service().openDirectPrediction( - mainEditor, - content2.document.text - ) - popup.dispose() - }, component) - font = JBUI.Fonts.smallFont() - } - ) - tagPanel.add(Box.createHorizontalStrut(6)) - tagPanel.add(TagComponent("Accept: ${getShortcutText(AcceptNextPredictionRevisionAction.ID)}").apply { - setListener({ _, _ -> - applyChanges() - popup.dispose() - }, component) - font = JBUI.Fonts.smallFont() - }) - return tagPanel - } - - private fun setupStatusLabel() { - (myEditor.scrollPane as JBScrollPane).statusComponent = BorderLayoutPanel() - .andTransparent() - .withBorder(JBUI.Borders.empty(4)) - .addToRight(getTagPanel()) - - val footerText = if (isManuallyOpened) { - CodeGPTBundle.get("shared.escToCancel") - } else { - "Trigger manually: ${getShortcutText(TriggerCustomPredictionAction.ID)} · ${CodeGPTBundle.get("shared.escToCancel")}" - } - - myEditor.component.add( - BorderLayoutPanel() - .addToRight( - JBLabel(footerText) - .apply { - font = JBUI.Fonts.miniFont() - }) - .apply { - background = editor.backgroundColor - border = JBUI.Borders.empty(4) - }, - BorderLayout.SOUTH - ) - } - private fun getVisibleAreaListener(): VisibleAreaListener { return object : VisibleAreaListener { override fun visibleAreaChanged(event: VisibleAreaEvent) { @@ -268,6 +214,36 @@ class CodeSuggestionDiffViewer( } } + private fun updateFooterComponent() { + for (component in myEditor.component.components) { + if (component is BorderLayoutPanel) { + myEditor.component.remove(component) + } + } + + myEditor.component.add( + BorderLayoutPanel() + .addToLeft( + JBLabel( + "Accept: ${getShortcutText(AcceptNextPredictionRevisionAction.ID)} " + + "· Trigger: ${getShortcutText(TriggerCustomPredictionAction.ID)} " + + "· Open: ${getShortcutText(OpenPredictionAction.ID)} " + + "· Changes: ${diffChanges?.size ?: 0}" + ) + .apply { + font = JBUI.Fonts.miniFont() + }) + .apply { + background = editor.backgroundColor + border = JBUI.Borders.empty(4) + }, + BorderLayout.SOUTH + ) + + myEditor.component.revalidate() + myEditor.component.repaint() + } + private class MyDiffContext(private val project: Project?) : DiffContext() { private val ownContext: UserDataHolder = UserDataHolderBase() @@ -299,7 +275,6 @@ class CodeSuggestionDiffViewer( fun displayInlineDiff( editor: Editor, nextEditResponse: NextEditResponse, - isManuallyOpened: Boolean = false ) { val nextRevision = nextEditResponse.nextRevision if (editor.virtualFile == null || editor.isViewer || nextRevision.isEmpty()) { @@ -315,8 +290,7 @@ class CodeSuggestionDiffViewer( } val diffRequest = createSimpleDiffRequest(editor, nextRevision) - val diffViewer = - CodeSuggestionDiffViewer(diffRequest, nextEditResponse, editor, isManuallyOpened) + val diffViewer = CodeSuggestionDiffViewer(diffRequest, nextEditResponse, editor) editor.putUserData(CodeGPTKeys.EDITOR_PREDICTION_DIFF_VIEWER, diffViewer) diffViewer.rediff(true) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/predictions/OpenPredictionAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/predictions/OpenPredictionAction.kt index 252bf125..cfeabce0 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/predictions/OpenPredictionAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/predictions/OpenPredictionAction.kt @@ -24,7 +24,7 @@ class OpenPredictionAction : EditorAction(Handler()), HintManagerImpl.ActionToIg runInEdt { diffViewer.dispose() } - service().openDirectPrediction(editor, nextRevision) + service().showDiff(editor, nextRevision) } } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/predictions/PredictionService.kt b/src/main/kotlin/ee/carlrobert/codegpt/predictions/PredictionService.kt index 905ba12d..e8f5e5bd 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/predictions/PredictionService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/predictions/PredictionService.kt @@ -2,12 +2,14 @@ package ee.carlrobert.codegpt.predictions import com.intellij.diff.DiffManager import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.application.runReadAction import com.intellij.openapi.components.Service 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.testFramework.LightVirtualFile +import com.intellij.util.application import ee.carlrobert.codegpt.CodeGPTKeys import ee.carlrobert.codegpt.codecompletions.CompletionProgressNotifier import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService @@ -29,14 +31,17 @@ class PredictionService { } } - fun displayInlineDiff( - editor: Editor, - isManuallyOpened: Boolean = false - ) { + fun displayInlineDiff(editor: Editor) { val project = editor.project ?: return try { - CompletionProgressNotifier.update(project, true) - project.service().getNextEdit(editor, isManuallyOpened) + application.executeOnPooledThread { + CompletionProgressNotifier.update(project, true) + project.service().getNextEdit( + editor, + editor.document.text, + runReadAction { editor.caretModel.offset }, + ) + } } catch (e: CancellationException) { // ignore } catch (ex: Exception) { @@ -44,7 +49,7 @@ class PredictionService { } } - fun openDirectPrediction(editor: Editor, nextRevision: String) { + fun showDiff(editor: Editor, nextRevision: String) { val project: Project = editor.project ?: return val tempDiffFile = LightVirtualFile(editor.virtualFile.name, nextRevision) val diffRequest = createDiffRequest(project, tempDiffFile, editor.virtualFile) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt index 67467b2d..40b3dc20 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt @@ -26,7 +26,7 @@ class TriggerCustomPredictionAction : EditorAction(Handler()), HintManagerImpl.A } ApplicationManager.getApplication().executeOnPooledThread { - service().displayInlineDiff(editor, true) + service().displayInlineDiff(editor) } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/CodeCompletionConfigurationForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/CodeCompletionConfigurationForm.kt index 4c6e083f..98f950a2 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/CodeCompletionConfigurationForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/CodeCompletionConfigurationForm.kt @@ -9,10 +9,6 @@ import ee.carlrobert.codegpt.CodeGPTBundle class CodeCompletionConfigurationForm { - private val multiLineCompletionsCheckBox = JBCheckBox( - CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.multiLineCompletions.title"), - service().state.codeCompletionSettings.multiLineEnabled - ) private val treeSitterProcessingCheckBox = JBCheckBox( CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.postProcess.title"), service().state.codeCompletionSettings.treeSitterProcessingEnabled @@ -21,7 +17,6 @@ class CodeCompletionConfigurationForm { CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.gitDiff.title"), service().state.codeCompletionSettings.gitDiffEnabled ) - private val collectDependencyStructureBox = JBCheckBox( CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.collectDependencyStructure.title"), service().state.codeCompletionSettings.collectDependencyStructure @@ -29,10 +24,6 @@ class CodeCompletionConfigurationForm { fun createPanel(): DialogPanel { return panel { - row { - cell(multiLineCompletionsCheckBox) - .comment(CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.multiLineCompletions.description")) - } row { cell(treeSitterProcessingCheckBox) .comment(CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.postProcess.description")) @@ -49,14 +40,12 @@ class CodeCompletionConfigurationForm { } fun resetForm(prevState: CodeCompletionSettingsState) { - multiLineCompletionsCheckBox.isSelected = prevState.multiLineEnabled treeSitterProcessingCheckBox.isSelected = prevState.treeSitterProcessingEnabled gitDiffCheckBox.isSelected = prevState.gitDiffEnabled } fun getFormState(): CodeCompletionSettingsState { return CodeCompletionSettingsState().apply { - this.multiLineEnabled = multiLineCompletionsCheckBox.isSelected this.treeSitterProcessingEnabled = treeSitterProcessingCheckBox.isSelected this.gitDiffEnabled = gitDiffCheckBox.isSelected this.collectDependencyStructure = collectDependencyStructureBox.isSelected diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt index 1e521f92..6e3592e1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt @@ -46,7 +46,6 @@ class ChatCompletionSettingsState : BaseState() { } class CodeCompletionSettingsState : BaseState() { - var multiLineEnabled by property(true) var treeSitterProcessingEnabled by property(true) var gitDiffEnabled by property(true) var collectDependencyStructure by property(true) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt index 1dabc632..b33ab1b1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt @@ -2,14 +2,17 @@ package ee.carlrobert.codegpt.settings.service.codegpt import com.intellij.openapi.components.service import com.intellij.openapi.ui.ComboBox +import com.intellij.ui.TitledSeparator import com.intellij.ui.components.JBCheckBox import com.intellij.ui.components.JBPasswordField import com.intellij.util.ui.FormBuilder import ee.carlrobert.codegpt.CodeGPTBundle +import ee.carlrobert.codegpt.codecompletions.edit.GrpcClientService import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.CodeGptApiKey import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential import ee.carlrobert.codegpt.credentials.CredentialsStore.setCredential import ee.carlrobert.codegpt.ui.UIUtil +import ee.carlrobert.codegpt.util.ApplicationUtil import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.runBlocking import org.jdesktop.swingx.combobox.ListComboBoxModel @@ -73,15 +76,15 @@ class CodeGPTServiceForm { UIUtil.createComment("settingsConfigurable.service.codegpt.codeCompletionModel.comment") ) .addVerticalGap(4) - .addComponent(enableNextEditsEnabledCheckBox) - .addComponent( - UIUtil.createComment("settingsConfigurable.service.codegpt.enableNextEdits.comment", 90) - ) - .addVerticalGap(4) .addComponent(codeCompletionsEnabledCheckBox) .addComponent( UIUtil.createComment("settingsConfigurable.service.codegpt.enableCodeCompletion.comment", 90) ) + .addVerticalGap(4) + .addComponent(enableNextEditsEnabledCheckBox) + .addComponent( + UIUtil.createComment("settingsConfigurable.service.codegpt.enableNextEdits.comment", 90) + ) .addComponentFillVertically(JPanel(), 0) .panel @@ -106,6 +109,8 @@ class CodeGPTServiceForm { (codeCompletionModelComboBox.selectedItem as CodeGPTModel).code } setCredential(CodeGptApiKey, getApiKey()) + + ApplicationUtil.findCurrentProject()?.service()?.refreshConnection() } fun resetForm() { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt index 98e7ff33..5f9b8788 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt @@ -21,6 +21,6 @@ class CodeGPTServiceChatCompletionSettingsState : BaseState() { } class CodeGPTServiceCodeCompletionSettingsState : BaseState() { - var codeCompletionsEnabled by property(false) - var model by string("codestral") + var codeCompletionsEnabled by property(true) + var model by string(CodeGPTAvailableModels.DEFAULT_CODE_MODEL.toString()) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt index 55e472e8..59af03d6 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt @@ -152,6 +152,7 @@ object GitUtil { line.startsWith("---") || line.startsWith("+++") || line.startsWith("===") || + line.contains("\\ No newline at end of file") (!showContext && line.startsWith(" ")) } .joinToString("\n") diff --git a/src/main/proto/code-completion.proto b/src/main/proto/code-completion.proto new file mode 100644 index 00000000..72ca9a1a --- /dev/null +++ b/src/main/proto/code-completion.proto @@ -0,0 +1,31 @@ +// src/main/proto/code-completion.proto +syntax = "proto3"; +option java_multiple_files = true; +option java_package = "ee.carlrobert.service"; + +import "google/protobuf/empty.proto"; + +service CodeCompletionServiceImpl { + rpc GetCodeCompletion (GrpcCodeCompletionRequest) returns (stream PartialCodeCompletionResponse); + rpc AcceptCodeCompletion (AcceptCodeCompletionRequest) returns (google.protobuf.Empty); +} + +message GrpcCodeCompletionRequest { + string model = 1; + string file_content = 2; + string file_path = 3; + int32 cursor_position = 4; + string git_diff = 5; + bool enable_telemetry = 6; +} + +message PartialCodeCompletionResponse { + string id = 1; + string partial_completion = 2; + bool done = 3; +} + +message AcceptCodeCompletionRequest { + string response_id = 1; + string accepted_completion = 2; +} \ No newline at end of file diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index 84fb2306..da000f92 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -90,11 +90,6 @@ messages.codegpt - - - - - - - - - - - - - - - - - @@ -273,8 +252,6 @@ - - diff --git a/src/main/resources/messages/codegpt.properties b/src/main/resources/messages/codegpt.properties index d45a5ba0..7233bf6e 100644 --- a/src/main/resources/messages/codegpt.properties +++ b/src/main/resources/messages/codegpt.properties @@ -144,7 +144,6 @@ configurationConfigurable.section.assistant.maxTokensField.label=Max completion configurationConfigurable.section.assistant.maxTokensField.comment=The maximum capacity for completion. configurationConfigurable.section.assistant.llamacppParams.title=Configuration Options for llama.cpp configurationConfigurable.section.codeCompletion.title=Code Completion -configurationConfigurable.section.codeCompletion.multiLineCompletions.title=Enable multi-line completions configurationConfigurable.section.codeCompletion.multiLineCompletions.description=If checked, the completion will be able to span multiple lines. configurationConfigurable.section.codeCompletion.postProcess.title=Enable tree-sitter post-processing configurationConfigurable.section.codeCompletion.postProcess.description=If checked, the completion will be post-processed using the tree-sitter parser. diff --git a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt index 1ac56163..e73360f4 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt @@ -19,7 +19,6 @@ class CodeCompletionServiceTest : IntegrationTest() { fun `test code completion with ProxyAI provider`() { useCodeGPTService() service().state.nextEditsEnabled = false - service().state.codeCompletionSettings.multiLineEnabled = false myFixture.configureByText( "CompletionTest.java", FileUtil.getResourceContent("/codecompletions/code-completion-file.txt") @@ -60,7 +59,6 @@ class CodeCompletionServiceTest : IntegrationTest() { fun `test code completion with OpenAI provider`() { useOpenAIService() service().state.nextEditsEnabled = false - service().state.codeCompletionSettings.multiLineEnabled = false myFixture.configureByText( "CompletionTest.java", FileUtil.getResourceContent("/codecompletions/code-completion-file.txt") @@ -98,10 +96,64 @@ class CodeCompletionServiceTest : IntegrationTest() { } } + fun `test apply next partial completion word`() { + useLlamaService(true) + service().state.nextEditsEnabled = false + myFixture.configureByText( + "CompletionTest.java", + FileUtil.getResourceContent("/codecompletions/code-completion-file.txt") + ) + myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(3, 0)) + val expectedCompletion = "public void main" + val prefix = """ + xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz + [INPUT] + c + """.trimIndent() + val suffix = """ + + [\INPUT] + zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz + xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx + """.trimIndent() + expectLlama(StreamHttpExchange { request: RequestEntity -> + assertThat(request.uri.path).isEqualTo("/completion") + assertThat(request.method).isEqualTo("POST") + assertThat(request.body) + .extracting("prompt") + .isEqualTo( + InfillPromptTemplate.CODE_LLAMA.buildPrompt( + InfillRequest.Builder(prefix, suffix, 0).build() + ) + ) + listOf( + jsonMapResponse( + e("content", expectedCompletion), + e("stop", true) + ) + ) + }) + myFixture.type('c') + assertInlineSuggestion("Failed to display initial inline suggestion.") { + expectedCompletion == it + } + val offsetBeforeApply = myFixture.editor.caretModel.offset + + myFixture.performEditorAction(AcceptNextWordInlayAction.ID) + + assertInlineSuggestion("Failed to display next partial inline suggestion.") { + myFixture.run { + val appliedText = + editor.document.getText(TextRange(offsetBeforeApply, editor.caretModel.offset)) + "public" == appliedText && " void main" == it + } + } + } + fun `_test apply inline suggestions without initial following text`() { useCodeGPTService() service().state.nextEditsEnabled = false - service().state.codeCompletionSettings.multiLineEnabled = false myFixture.configureByText( "CompletionTest.java", "class Node {\n " @@ -219,7 +271,6 @@ class CodeCompletionServiceTest : IntegrationTest() { fun `_test apply inline suggestions with initial following text`() { useCodeGPTService() service().state.nextEditsEnabled = false - service().state.codeCompletionSettings.multiLineEnabled = false myFixture.configureByText( "CompletionTest.java", "if () {\n \n} else {\n}" @@ -288,10 +339,9 @@ class CodeCompletionServiceTest : IntegrationTest() { } } - fun `test adjust completion line whitespaces`() { + fun `_test adjust completion line whitespaces`() { useCodeGPTService() service().state.nextEditsEnabled = false - service().state.codeCompletionSettings.multiLineEnabled = false myFixture.configureByText( "CompletionTest.java", "class Node {\n" +