diff --git a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java index 43c139f3..6851dcc1 100644 --- a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java +++ b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java @@ -20,4 +20,8 @@ public class CodeGPTKeys { Key.create("codegpt.addedDocumentation"); public static final Key ADDED_PERSONA = Key.create("codegpt.addedPersona"); + public static final Key REMAINING_EDITOR_COMPLETION = + Key.create("codegpt.editorCompletionLines"); + public static final Key IS_FETCHING_COMPLETION = + Key.create("codegpt.isFetchingCompletion"); } diff --git a/src/main/java/ee/carlrobert/codegpt/Icons.java b/src/main/java/ee/carlrobert/codegpt/Icons.java index 5cd2dc58..bc618476 100644 --- a/src/main/java/ee/carlrobert/codegpt/Icons.java +++ b/src/main/java/ee/carlrobert/codegpt/Icons.java @@ -1,6 +1,7 @@ package ee.carlrobert.codegpt; import com.intellij.openapi.util.IconLoader; +import com.intellij.ui.AnimatedIcon; import javax.swing.Icon; public final class Icons { @@ -27,4 +28,5 @@ public final class Icons { public static final Icon Upload = IconLoader.getIcon("/icons/upload.svg", Icons.class); public static final Icon GreenCheckmark = IconLoader.getIcon("/icons/greenCheckmark.svg", Icons.class); + public static final Icon StatusBarCompletionInProgress = new AnimatedIcon.Default(); } diff --git a/src/main/java/ee/carlrobert/codegpt/statusbar/CodeGPTStatusBarWidget.java b/src/main/java/ee/carlrobert/codegpt/statusbar/CodeGPTStatusBarWidget.java index 9352e6a0..a3d9db2d 100644 --- a/src/main/java/ee/carlrobert/codegpt/statusbar/CodeGPTStatusBarWidget.java +++ b/src/main/java/ee/carlrobert/codegpt/statusbar/CodeGPTStatusBarWidget.java @@ -1,5 +1,7 @@ package ee.carlrobert.codegpt.statusbar; +import static ee.carlrobert.codegpt.CodeGPTKeys.IS_FETCHING_COMPLETION; + import com.intellij.openapi.actionSystem.ActionGroup; import com.intellij.openapi.actionSystem.ActionManager; import com.intellij.openapi.actionSystem.DataContext; @@ -8,24 +10,43 @@ import com.intellij.openapi.ui.popup.JBPopupFactory; import com.intellij.openapi.ui.popup.JBPopupFactory.ActionSelectionAid; import com.intellij.openapi.ui.popup.ListPopup; import com.intellij.openapi.vfs.VirtualFile; +import com.intellij.openapi.wm.StatusBar; import com.intellij.openapi.wm.StatusBarWidget; +import com.intellij.openapi.wm.WindowManager; import com.intellij.openapi.wm.impl.status.EditorBasedStatusBarPopup; import ee.carlrobert.codegpt.CodeGPTBundle; import ee.carlrobert.codegpt.Icons; +import ee.carlrobert.codegpt.codecompletions.CodeCompletionProgressNotifier; import org.jetbrains.annotations.NonNls; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; public class CodeGPTStatusBarWidget extends EditorBasedStatusBarPopup { + private static final String ID = "ee.carlrobert.codegpt.statusbar.widget"; + public CodeGPTStatusBarWidget(Project project) { super(project, false); + + project.getMessageBus() + .connect(this) + .subscribe( + CodeCompletionProgressNotifier.Companion.getCODE_COMPLETION_PROGRESS_TOPIC(), + (CodeCompletionProgressNotifier) loading -> { + CodeGPTStatusBarWidget widget = findWidget(project); + if (widget != null && widget.myStatusBar != null) { + widget.update(() -> widget.myStatusBar.updateWidget(ID)); + } + }); } @Override protected @NotNull WidgetState getWidgetState(@Nullable VirtualFile file) { var state = new WidgetState(CodeGPTBundle.get("statusBar.widget.tooltip"), "", true); - state.setIcon(Icons.DefaultSmall); + var fetchingCompletion = IS_FETCHING_COMPLETION.get(getEditor()); + var loading = fetchingCompletion != null && fetchingCompletion; + + state.setIcon(loading ? Icons.StatusBarCompletionInProgress : Icons.DefaultSmall); return state; } @@ -47,6 +68,18 @@ public class CodeGPTStatusBarWidget extends EditorBasedStatusBarPopup { @Override public @NonNls @NotNull String ID() { - return "ee.carlrobert.codegpt.statusbar.widget"; + return ID; + } + + private static @Nullable CodeGPTStatusBarWidget findWidget(@NotNull Project project) { + StatusBar bar = WindowManager.getInstance().getStatusBar(project); + if (bar != null) { + StatusBarWidget widget = bar.getWidget(ID); + if (widget instanceof CodeGPTStatusBarWidget) { + return (CodeGPTStatusBarWidget) widget; + } + } + + return null; } } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowType.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowType.java new file mode 100644 index 00000000..3394bf1f --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowType.java @@ -0,0 +1,13 @@ +package ee.carlrobert.codegpt.toolwindow.chat; + +public enum ChatToolWindowType { + CODEGPT_CHAT("CodeGPT Chat"), + CODEGPT_CHAT_WITHOUT_PERSONA("CodeGPT Chat without Persona"), + CODEGPT_CHAT_WITH_PERSONA("CodeGPT Chat with Persona"); + + private final String name; + + ChatToolWindowType(String name) { + this.name = name; + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCompletionEventListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCompletionEventListener.kt new file mode 100644 index 00000000..4124ef24 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCompletionEventListener.kt @@ -0,0 +1,113 @@ +package ee.carlrobert.codegpt.codecompletions + +import ai.grazie.nlp.utils.takeLastWhitespaces +import ai.grazie.nlp.utils.takeWhitespaces +import com.intellij.notification.NotificationType +import com.intellij.openapi.diagnostic.thisLogger +import com.intellij.openapi.editor.Editor +import ee.carlrobert.codegpt.CodeGPTKeys.IS_FETCHING_COMPLETION +import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification +import ee.carlrobert.llm.client.openai.completion.ErrorDetails +import ee.carlrobert.llm.completion.CompletionEventListener +import okhttp3.sse.EventSource +import java.util.concurrent.atomic.AtomicBoolean + +private const val MAX_LINES_TO_REQUEST = 4 +private const val MAX_LINES_TO_DISPLAY = 2 + +abstract class CodeCompletionCompletionEventListener( + private val editor: Editor, + private val infillRequest: InfillRequest? = null, +) : CompletionEventListener { + + companion object { + private val logger = thisLogger() + } + + private val stringBuilder = StringBuilder() + private val isCancelled = AtomicBoolean(false) + private val isSending = AtomicBoolean(true) + + open fun onComplete(fullMessage: String) {} + open fun onMessage(message: String) {} + + override fun onOpen() { + setLoading(true) + } + + override fun onMessage(message: String, eventSource: EventSource) { + if (isCancelled.get()) return + + val processedMessage = if (infillRequest != null && stringBuilder.isEmpty()) { + message.tryTrimStart(infillRequest.prefix.lines()) + } else { + message + } + + val newLineCount = (stringBuilder.toString() + processedMessage).count { it == '\n' } + if (newLineCount >= MAX_LINES_TO_REQUEST) { + cancelStreaming(processedMessage, eventSource) + return + } + + stringBuilder.append(processedMessage) + + if (newLineCount <= MAX_LINES_TO_DISPLAY && isSending.get()) { + if (newLineCount == MAX_LINES_TO_DISPLAY && processedMessage.contains('\n')) { + isSending.set(false) + onMessage(processedMessage.substring(0, processedMessage.lastIndexOf('\n'))) + } else { + onMessage(processedMessage) + } + } + } + + override fun onComplete(messageBuilder: StringBuilder) { + setLoading(false) + onComplete(stringBuilder.trimEnd().toString()) + } + + override fun onCancelled(messageBuilder: StringBuilder) { + setLoading(false) + onComplete(stringBuilder.trimEnd().toString()) + } + + override fun onError(error: ErrorDetails, ex: Throwable) { + if (ex.message == null || (ex.message != null && ex.message != "Canceled")) { + showNotification(error.message, NotificationType.ERROR) + logger.error(error.message, ex) + } + } + + private fun setLoading(loading: Boolean) { + IS_FETCHING_COMPLETION.set(editor, loading) + editor.project?.messageBus + ?.syncPublisher(CodeCompletionProgressNotifier.CODE_COMPLETION_PROGRESS_TOPIC) + ?.loading(loading) + } + + private fun cancelStreaming(processedMessage: String, eventSource: EventSource) { + stringBuilder.append(processedMessage.substring(0, processedMessage.lastIndexOf('\n'))) + isCancelled.set(true) + isSending.set(false) + eventSource.cancel() + } + + private fun String.tryTrimStart(lines: List): String { + val whiteSpaces = this.takeWhitespaces() + if (lines.size >= 2 + && whiteSpaces.isNotEmpty() + && lines[lines.size - 1].trim().isEmpty() + ) { + return this.trimStart() + } + + if (lines.isNotEmpty()) { + val lastLine = lines[lines.size - 1] + if (lastLine.takeLastWhitespaces().isNotEmpty()) { + return this.trimStart() + } + } + return this + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt new file mode 100644 index 00000000..355b7778 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt @@ -0,0 +1,81 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.intellij.codeInsight.inline.completion.InlineCompletion +import com.intellij.codeInsight.inline.completion.InlineCompletionEvent +import com.intellij.codeInsight.inline.completion.InlineCompletionInsertEnvironment +import com.intellij.codeInsight.inline.completion.InlineCompletionInsertHandler +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement +import com.intellij.openapi.application.readAction +import com.intellij.openapi.components.service +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.util.TextRange +import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.launch + +private const val NEXT_COMPLETION_LINE_COUNT_THRESHOLD = 4 + +class CodeCompletionInsertHandler : InlineCompletionInsertHandler { + + private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob()) + + override fun afterInsertion( + environment: InlineCompletionInsertEnvironment, + elements: List + ) { + val editor = environment.editor + val appliedText = elements.joinToString("") { it.text } + val existingCompletion = REMAINING_EDITOR_COMPLETION.get(editor) ?: "" + val remainingCompletion = + existingCompletion.substring(appliedText.length, existingCompletion.length) + + REMAINING_EDITOR_COMPLETION.set(editor, remainingCompletion) + + if (remainingCompletion.isNotEmpty()) { + val handler = InlineCompletion.getHandlerOrNull(editor) + handler?.invoke( + InlineCompletionEvent.DirectCall(editor, editor.caretModel.currentCaret) + ) + + if (remainingCompletion.count { it == '\n' } <= NEXT_COMPLETION_LINE_COUNT_THRESHOLD) { + scope.launch { + fetchNextCompletion(editor, remainingCompletion) + } + } + } + } + + private suspend fun fetchNextCompletion(editor: Editor, remainingCompletion: String) { + val project = editor.project ?: return + project.service().getCodeCompletionAsync( + buildNextRequest(editor, remainingCompletion), + object : CodeCompletionCompletionEventListener(editor) { + override fun onComplete(fullMessage: String) { + val nextCompletion = + (REMAINING_EDITOR_COMPLETION.get(editor) ?: "") + fullMessage + REMAINING_EDITOR_COMPLETION.set(editor, nextCompletion) + } + } + ) + } + + private suspend fun buildNextRequest( + editor: Editor, + remainingCompletion: String + ): InfillRequest { + val caretOffset = readAction { editor.caretModel.offset } + val prefix = + (editor.document.getText(TextRange(0, caretOffset)) + remainingCompletion) + .truncateText(MAX_PROMPT_TOKENS, false) + val suffix = + editor.document.getText( + TextRange( + caretOffset, + editor.document.textLength + ) + ).truncateText(MAX_PROMPT_TOKENS) + return InfillRequest.Builder(prefix, suffix).build() + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionProgressNotifier.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionProgressNotifier.kt new file mode 100644 index 00000000..c8cbc402 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionProgressNotifier.kt @@ -0,0 +1,14 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.intellij.util.messages.Topic + +interface CodeCompletionProgressNotifier { + + fun loading(loading: Boolean) + + companion object { + @JvmStatic + val CODE_COMPLETION_PROGRESS_TOPIC = + Topic.create("codeCompletionProgressTopic", CodeCompletionProgressNotifier::class.java) + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionProviderPresentation.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionProviderPresentation.kt new file mode 100644 index 00000000..0b4fb0ed --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionProviderPresentation.kt @@ -0,0 +1,24 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.intellij.codeInsight.inline.completion.InlineCompletionProviderPresentation +import com.intellij.openapi.components.service +import com.intellij.openapi.project.Project +import com.intellij.ui.components.JBLabel +import com.intellij.util.ui.JBFont +import ee.carlrobert.codegpt.Icons +import javax.swing.JComponent +import javax.swing.SwingConstants + +class CodeCompletionProviderPresentation : InlineCompletionProviderPresentation { + + override fun getTooltip(project: Project?): JComponent { + val selectedModelCode = + project?.service()?.getSelectedModelCode() ?: "" + val text = if (selectedModelCode.isNotEmpty()) { + "CodeGPT: $selectedModelCode" + } else { + "CodeGPT" + } + return JBLabel(text, Icons.Sparkle, SwingConstants.LEADING).withFont(JBFont.small()) + } +} \ 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 89c2e219..722444b2 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt @@ -3,6 +3,7 @@ package ee.carlrobert.codegpt.codecompletions import com.fasterxml.jackson.core.JsonProcessingException import com.fasterxml.jackson.databind.ObjectMapper import com.intellij.openapi.components.service +import ee.carlrobert.codegpt.EncodingManager import ee.carlrobert.codegpt.completions.llama.LlamaModel import ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey import ee.carlrobert.codegpt.credentials.CredentialsStore.getCredential @@ -42,8 +43,9 @@ object CodeCompletionRequestFactory { @JvmStatic fun buildOpenAIRequest(details: InfillRequest): OpenAITextCompletionRequest { - return OpenAITextCompletionRequest.Builder(details.prefix) - .setSuffix(details.suffix) + val (prefix, suffix) = getCompletionContext(details) + return OpenAITextCompletionRequest.Builder(prefix) + .setSuffix(suffix) .setStream(true) .setMaxTokens(MAX_TOKENS) .setTemperature(0.4) @@ -116,6 +118,7 @@ object CodeCompletionRequestFactory { settings.model, settings.fimTemplate.buildPrompt(details) ) + .setStream(true) .setOptions( OllamaParameters.Builder() .stop(settings.fimTemplate.stopTokens) @@ -144,15 +147,41 @@ object CodeCompletionRequestFactory { ): Any { if (value !is String) return value + val (prefix, suffix) = getCompletionContext(details) return when (value) { FIM_PROMPT.code -> template.buildPrompt(details) - PREFIX.code -> details.prefix - SUFFIX.code -> details.suffix + PREFIX.code -> prefix + SUFFIX.code -> suffix else -> { return value.takeIf { it.contains(PREFIX.code) || it.contains(SUFFIX.code) } - ?.replace(PREFIX.code, details.prefix) - ?.replace(SUFFIX.code, details.suffix) ?: value + ?.replace(PREFIX.code, prefix) + ?.replace(SUFFIX.code, suffix) ?: value } } } + + private fun getCompletionContext(request: InfillRequest): Pair { + val encodingManager = EncodingManager.getInstance() + val truncatedPrefix = encodingManager.truncateText(request.prefix, 128, false) + val truncatedSuffix = encodingManager.truncateText(request.suffix, 128, true) + val vcsDetails = request.vcsDetails ?: return truncatedPrefix to truncatedSuffix + + val stagedDiff = if (vcsDetails.stagedDiff != null) + encodingManager.truncateText(vcsDetails.stagedDiff, 200, true) + else + "" + val unstagedDiff = if (vcsDetails.unstagedDiff != null) + encodingManager.truncateText(vcsDetails.unstagedDiff, 200, true) + else + "" + val prompt: String = if (vcsDetails.stagedDiff != null) + """ + ${"/*\n${stagedDiff + unstagedDiff}\n\n*/"} + $truncatedPrefix + """.trimIndent() + else + truncatedPrefix + + return prompt to truncatedSuffix + } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt index 40526045..ede79420 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt @@ -24,6 +24,22 @@ import okhttp3.sse.EventSources.createFactory @Service(Service.Level.PROJECT) class CodeCompletionService { + // TODO: Consolidate logic in ModelComboBoxAction + fun getSelectedModelCode(): String? { + return when (service().state.selectedService) { + CODEGPT -> service().state.codeCompletionSettings.model + OPENAI -> "gpt-3.5-turbo-instruct" + CUSTOM_OPENAI -> service().state + .codeCompletionSettings + .body + .getOrDefault("model", null) as String + + LLAMA_CPP -> "LlamaSettings.getCurrentState()." + OLLAMA -> service().state.model + else -> null + } + } + fun isCodeCompletionsEnabled(selectedService: ServiceType): Boolean = when (selectedService) { CODEGPT -> service().state.codeCompletionSettings.codeCompletionsEnabled diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionProvider.kt deleted file mode 100644 index e0310413..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionProvider.kt +++ /dev/null @@ -1,55 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions - -import com.intellij.codeInsight.inline.completion.* -import com.intellij.openapi.components.service -import com.intellij.openapi.diagnostic.thisLogger -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.settings.service.custom.CustomServiceSettings -import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings -import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings -import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings -import kotlin.time.Duration -import kotlin.time.DurationUnit -import kotlin.time.toDuration - -class CodeGPTInlineCompletionProvider : DebouncedInlineCompletionProvider() { - - companion object { - private val logger = thisLogger() - } - - override val id: InlineCompletionProviderID - get() = InlineCompletionProviderID("CodeGPTInlineCompletionProvider") - - override suspend fun getSuggestionDebounced(request: InlineCompletionRequest): InlineCompletionSuggestion { - val editor = request.editor - val project = editor.project - if (project == null) { - logger.error("Could not find project") - return InlineCompletionSuggestion.empty() - } - return CodeGPTInlineCompletionSuggestion(project, request) - } - - override suspend fun getDebounceDelay(request: InlineCompletionRequest): Duration { - return 600.toDuration(DurationUnit.MILLISECONDS) - } - - override fun isEnabled(event: InlineCompletionEvent): Boolean { - val selectedService = GeneralSettings.getSelectedService() - val codeCompletionsEnabled = when (selectedService) { - ServiceType.CODEGPT -> service().state.codeCompletionSettings.codeCompletionsEnabled - ServiceType.OPENAI -> OpenAISettings.getCurrentState().isCodeCompletionsEnabled - ServiceType.CUSTOM_OPENAI -> service().state.codeCompletionSettings.codeCompletionsEnabled - ServiceType.LLAMA_CPP -> LlamaSettings.isCodeCompletionsPossible() - ServiceType.OLLAMA -> service().state.codeCompletionsEnabled - ServiceType.ANTHROPIC, - ServiceType.AZURE, - ServiceType.GOOGLE, - null -> false - } - return event is InlineCompletionEvent.DocumentChange && codeCompletionsEnabled - } -} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionSuggestion.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionSuggestion.kt deleted file mode 100644 index d75b9005..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionSuggestion.kt +++ /dev/null @@ -1,172 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions - -import com.intellij.codeInsight.inline.completion.InlineCompletionRequest -import com.intellij.codeInsight.inline.completion.InlineCompletionSuggestion -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.EDT -import com.intellij.openapi.application.runInEdt -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 com.intellij.openapi.vcs.VcsException -import com.intellij.refactoring.suggested.startOffset -import ee.carlrobert.codegpt.EncodingManager -import ee.carlrobert.codegpt.codecompletions.psi.CompletionContextService -import ee.carlrobert.codegpt.codecompletions.psi.readText -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings -import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification -import ee.carlrobert.codegpt.util.GitUtil -import ee.carlrobert.llm.client.openai.completion.ErrorDetails -import ee.carlrobert.llm.completion.CompletionEventListener -import git4idea.repo.GitRepositoryManager -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.channels.awaitClose -import kotlinx.coroutines.flow.Flow -import kotlinx.coroutines.flow.channelFlow -import kotlinx.coroutines.launch -import kotlinx.coroutines.withContext -import okhttp3.sse.EventSource -import java.util.concurrent.atomic.AtomicReference - -class CodeGPTInlineCompletionSuggestion( - private val project: Project, - private val inlineCompletionRequest: InlineCompletionRequest -) : InlineCompletionSuggestion() { - - private val currentCall = AtomicReference(null) - - companion object { - private val logger = thisLogger() - private const val MAX_PROMPT_TOKENS = 128 - } - - override val suggestionFlow: Flow - get() = channelFlow { - val infillRequest = buildInfillRequest(inlineCompletionRequest) - launch { - val completionCall = - project.service().getCodeCompletionAsync( - infillRequest, - CodeCompletionEventListener { - try { - runInEdt { - trySend(InlineCompletionGrayTextElement(it.toString())) - } - } catch (e: Exception) { - logger.error("Failed to send inline completion suggestion", e) - } - } - ) - currentCall.set(completionCall) - } - awaitClose { currentCall.getAndSet(null)?.cancel() } - } - - private class CodeCompletionEventListener( - private val completed: (StringBuilder) -> Unit - ) : CompletionEventListener { - - override fun onComplete(messageBuilder: StringBuilder) { - completed(messageBuilder) - } - - override fun onCancelled(messageBuilder: StringBuilder) { - completed(messageBuilder) - } - - override fun onError(error: ErrorDetails, ex: Throwable) { - if (ex.message == null || (ex.message != null && ex.message != "Canceled")) { - showNotification(error.message, NotificationType.ERROR) - logger.error(error.message, ex) - } - } - } - - private suspend fun buildInfillRequest(request: InlineCompletionRequest): InfillRequest { - val caretOffset = withContext(Dispatchers.EDT) { request.editor.caretModel.offset } - val configurationState = service().state - val (prefix, suffix) = withContext(Dispatchers.EDT) { - val prefix = - request.document.getText(TextRange(0, caretOffset)) - val suffix = - request.document.getText( - TextRange( - caretOffset, - request.document.textLength - ) - ) - Pair( - prefix.truncateText(MAX_PROMPT_TOKENS, false), - suffix.truncateText(MAX_PROMPT_TOKENS) - ) - } - val fileExtension = request.file.virtualFile.extension - val fileContent = request.document.text - val infillRequestBuilder = InfillRequest.Builder(prefix, suffix) - .fileDetails(InfillRequest.FileDetails(fileContent, fileExtension)) - val project = request.editor.project ?: return infillRequestBuilder.build() - - val gitRepository = - project.service().getRepositoryForFile(project.workspaceFile) - if (configurationState.autocompletionGitContextEnabled && gitRepository != null) { - try { - val stagedDiff = GitUtil.getStagedDiff(project, gitRepository) - val unstagedDiff = GitUtil.getUnstagedDiff(project, gitRepository) - if (stagedDiff.isNotEmpty() || unstagedDiff.isNotEmpty()) { - infillRequestBuilder.vcsDetails( - InfillRequest.VcsDetails( - stagedDiff.joinToString("\n"), - unstagedDiff.joinToString("\n") - ) - ) - } - } catch (e: VcsException) { - logger.error("Failed to get git context", e) - } - } - - getInfillContext(request, caretOffset)?.let { infillRequestBuilder.context(it) } - - return infillRequestBuilder.build() - } - - private fun getInfillContext( - request: InlineCompletionRequest, - caretOffset: Int - ): InfillContext? { - val infillContext = - if (service().state.autocompletionContextAwareEnabled) - service().findContext(request.editor, caretOffset) - else null - - if (infillContext == null) { - return null - } - - val caretInEnclosingElement = - caretOffset - infillContext.enclosingElement.psiElement.startOffset - val entireText = infillContext.enclosingElement.psiElement.readText() - val prefix = entireText.take(caretInEnclosingElement) - val suffix = - if (entireText.length < caretInEnclosingElement) "" else entireText.takeLast( - entireText.length - caretInEnclosingElement - ) - return truncateContext(prefix + suffix, infillContext) - } - - private fun String.truncateText(maxTokens: Int, fromStart: Boolean = true): String { - return service().truncateText(this, maxTokens, fromStart) - } - - private fun truncateContext(prompt: String, infillContext: InfillContext): InfillContext { - var promptTokens = EncodingManager.getInstance().countTokens(prompt) - val truncatedContextElements = infillContext.contextElements.takeWhile { - promptTokens += it.tokens - promptTokens <= MAX_PROMPT_TOKENS - }.toSet() - return InfillContext(infillContext.enclosingElement, truncatedContextElements) - } -} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CompletionSplitter.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CompletionSplitter.kt index d786150f..c1eeaa79 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CompletionSplitter.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CompletionSplitter.kt @@ -22,7 +22,7 @@ class CompletionSplitter { fullCompletion: String, isBoundaryCharacter: BoundaryFinder ): Int { - if (!isBoundaryCharacter.isBoundaryCharacter(fullCompletion[0])) { + if (fullCompletion.isEmpty() || !isBoundaryCharacter.isBoundaryCharacter(fullCompletion[0])) { return -1 } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt new file mode 100644 index 00000000..bd89b0d4 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt @@ -0,0 +1,127 @@ +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.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.IS_FETCHING_COMPLETION +import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION +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.settings.service.custom.CustomServiceSettings +import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings +import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings +import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings +import 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 +import kotlin.time.DurationUnit +import kotlin.time.toDuration + +class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() { + + companion object { + private val logger = thisLogger() + } + + private val currentCallRef = AtomicReference(null) + + override val id: InlineCompletionProviderID + get() = InlineCompletionProviderID("CodeGPTInlineCompletionProvider") + + override val insertHandler: InlineCompletionInsertHandler + get() = CodeCompletionInsertHandler() + + override val providerPresentation: InlineCompletionProviderPresentation + get() = CodeCompletionProviderPresentation() + + override suspend fun getSuggestionDebounced(request: InlineCompletionRequest): InlineCompletionSuggestion { + val editor = request.editor + if (request.event is InlineCompletionEvent.DirectCall) { + val activeCompletionLine = REMAINING_EDITOR_COMPLETION.get(editor) + if (activeCompletionLine != null && activeCompletionLine.isNotEmpty()) { + return InlineCompletionSuggestion.Default(channelFlow { + launch { + trySend( + InlineCompletionGrayTextElement( + activeCompletionLine.lines().take(2).joinToString("\n") + ) + ) + } + }) + } + } + + val project = editor.project + if (project == null) { + logger.error("Could not find project") + return InlineCompletionSuggestion.Default(emptyFlow()) + } + + return InlineCompletionSuggestion.Default(channelFlow { + REMAINING_EDITOR_COMPLETION.set(request, "") + IS_FETCHING_COMPLETION.set(request.editor, true) + + request.editor.project?.messageBus + ?.syncPublisher(CodeCompletionProgressNotifier.CODE_COMPLETION_PROGRESS_TOPIC) + ?.loading(true) + + val infillRequest = InfillRequestUtil.buildInfillRequest(request) + val call = project.service().getCodeCompletionAsync( + infillRequest, + getEventListener(request.editor, infillRequest) + ) + currentCallRef.set(call) + awaitClose { currentCallRef.getAndSet(null)?.cancel() } + }) + } + + override suspend fun getDebounceDelay(request: InlineCompletionRequest): Duration { + return 600.toDuration(DurationUnit.MILLISECONDS) + } + + override fun isEnabled(event: InlineCompletionEvent): Boolean { + val selectedService = GeneralSettings.getSelectedService() + val codeCompletionsEnabled = when (selectedService) { + ServiceType.CODEGPT -> service().state.codeCompletionSettings.codeCompletionsEnabled + ServiceType.OPENAI -> OpenAISettings.getCurrentState().isCodeCompletionsEnabled + ServiceType.CUSTOM_OPENAI -> service().state.codeCompletionSettings.codeCompletionsEnabled + ServiceType.LLAMA_CPP -> LlamaSettings.isCodeCompletionsPossible() + ServiceType.OLLAMA -> service().state.codeCompletionsEnabled + ServiceType.ANTHROPIC, + ServiceType.AZURE, + ServiceType.GOOGLE, + null -> false + } + val containsActiveCompletion = + REMAINING_EDITOR_COMPLETION.get(event.toRequest()?.editor)?.isNotEmpty() ?: false + + return (event is InlineCompletionEvent.DocumentChange && codeCompletionsEnabled) + || containsActiveCompletion + } + + private fun ProducerScope.getEventListener( + editor: Editor, + infillRequest: InfillRequest + ) = object : CodeCompletionCompletionEventListener(editor, infillRequest) { + override fun onMessage(message: String) { + runInEdt { + trySend(InlineCompletionGrayTextElement(message)) + } + } + + override fun onComplete(fullMessage: String) { + REMAINING_EDITOR_COMPLETION.set(editor, fullMessage) + } + } +} + diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt index 0f81de33..e9fbd23b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt @@ -1,8 +1,8 @@ package ee.carlrobert.codegpt.codecompletions -enum class InfillPromptTemplate(val label: String, val stopTokens: List?) { +enum class InfillPromptTemplate(val label: String, val stopTokens: List? = listOf("\n\n")) { - OPENAI("OpenAI", null) { + OPENAI("OpenAI") { override fun buildPrompt(infillDetails: InfillRequest): String { val infillPrompt = "<|fim_prefix|> ${infillDetails.prefix} <|fim_suffix|>${infillDetails.suffix} <|fim_middle|>" @@ -15,10 +15,7 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List? return createDefaultMultiFilePrompt(infillDetails, infillPrompt) } }, - CODE_GEMMA( - "CodeGemma Instruct", - listOf("<|file_separator|>", "<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>", "") - ) { + CODE_GEMMA("CodeGemma Instruct") { override fun buildPrompt(infillDetails: InfillRequest): String { // see https://huggingface.co/google/codegemma-7b#for-code-completion val infillPrompt = @@ -60,7 +57,7 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List? return createDefaultMultiFilePrompt(infillDetails, infillPrompt) } }, - DEEPSEEK_CODER("DeepSeek Coder", listOf("<|EOT|>")) { + DEEPSEEK_CODER("DeepSeek Coder", listOf("<|EOT|>", "<｜end▁of▁sentence｜>")) { override fun buildPrompt(infillDetails: InfillRequest): String { // see https://github.com/deepseek-ai/DeepSeek-Coder?tab=readme-ov-file#2-code-insertion val infillPrompt = diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt index dd6d028a..945a596c 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt @@ -1,9 +1,15 @@ package ee.carlrobert.codegpt.codecompletions +import com.intellij.openapi.components.service +import com.intellij.openapi.editor.Document +import com.intellij.openapi.util.TextRange import com.intellij.psi.PsiElement +import ee.carlrobert.codegpt.EncodingManager import ee.carlrobert.codegpt.codecompletions.psi.filePath import ee.carlrobert.codegpt.codecompletions.psi.readText +const val MAX_PROMPT_TOKENS = 128 + class InfillRequest private constructor( val prefix: String, val suffix: String, @@ -12,14 +18,38 @@ class InfillRequest private constructor( val context: InfillContext? ) { + companion object { + fun builder(prefix: String, suffix: String) = Builder(prefix, suffix) + } + data class VcsDetails(val stagedDiff: String? = null, val unstagedDiff: String? = null) data class FileDetails(val fileContent: String, val fileExtension: String? = null) - class Builder(private val prefix: String, private val suffix: String) { + class Builder { + private val prefix: String + private val suffix: String private var fileDetails: FileDetails? = null private var vcsDetails: VcsDetails? = null private var context: InfillContext? = null + constructor(prefix: String, suffix: String) { + this.prefix = prefix + this.suffix = suffix + } + + constructor(document: Document, caretOffset: Int) { + prefix = + document.getText(TextRange(0, caretOffset)) + .truncateText(MAX_PROMPT_TOKENS, false) + suffix = + document.getText( + TextRange( + caretOffset, + document.textLength + ) + ).truncateText(MAX_PROMPT_TOKENS) + } + fun fileDetails(fileDetails: FileDetails) = apply { this.fileDetails = fileDetails } fun vcsDetails(vcsDetails: VcsDetails) = apply { this.vcsDetails = vcsDetails } fun context(context: InfillContext) = apply { this.context = context } @@ -27,10 +57,6 @@ class InfillRequest private constructor( fun build() = InfillRequest(prefix, suffix, fileDetails, vcsDetails, context) } - - companion object { - fun builder(prefix: String, suffix: String) = Builder(prefix, suffix) - } } class InfillContext( @@ -48,3 +74,7 @@ class ContextElement(val psiElement: PsiElement) { fun filePath() = this.psiElement.filePath() fun text() = this.psiElement.readText() } + +fun String.truncateText(maxTokens: Int, fromStart: Boolean = true): String { + return service().truncateText(this, maxTokens, fromStart) +} \ 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 new file mode 100644 index 00000000..4225374e --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt @@ -0,0 +1,86 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.intellij.codeInsight.inline.completion.InlineCompletionRequest +import com.intellij.openapi.application.readAction +import com.intellij.openapi.components.service +import com.intellij.openapi.diagnostic.thisLogger +import com.intellij.openapi.vcs.VcsException +import com.intellij.refactoring.suggested.startOffset +import ee.carlrobert.codegpt.EncodingManager +import ee.carlrobert.codegpt.codecompletions.psi.CompletionContextService +import ee.carlrobert.codegpt.codecompletions.psi.readText +import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings +import ee.carlrobert.codegpt.util.GitUtil +import git4idea.repo.GitRepositoryManager + +object InfillRequestUtil { + private val logger = thisLogger() + + suspend fun buildInfillRequest(request: InlineCompletionRequest): InfillRequest { + val caretOffset = readAction { request.editor.caretModel.offset } + val infillRequestBuilder = InfillRequest.Builder(request.document, caretOffset) + .fileDetails( + InfillRequest.FileDetails( + request.document.text, + request.file.virtualFile.extension + ) + ) + val project = request.editor.project ?: return infillRequestBuilder.build() + + val gitRepository = + project.service().getRepositoryForFile(project.workspaceFile) + if (service().state.autocompletionGitContextEnabled && gitRepository != null) { + try { + val stagedDiff = GitUtil.getStagedDiff(project, gitRepository) + val unstagedDiff = GitUtil.getUnstagedDiff(project, gitRepository) + if (stagedDiff.isNotEmpty() || unstagedDiff.isNotEmpty()) { + infillRequestBuilder.vcsDetails( + InfillRequest.VcsDetails( + stagedDiff.joinToString("\n"), + unstagedDiff.joinToString("\n") + ) + ) + } + } catch (e: VcsException) { + logger.error("Failed to get git context", e) + } + } + + getInfillContext(request, caretOffset)?.let { infillRequestBuilder.context(it) } + + return infillRequestBuilder.build() + } + + private fun getInfillContext( + request: InlineCompletionRequest, + caretOffset: Int + ): InfillContext? { + val infillContext = + if (service().state.autocompletionContextAwareEnabled) + service().findContext(request.editor, caretOffset) + else null + + if (infillContext == null) { + return null + } + + val caretInEnclosingElement = + caretOffset - infillContext.enclosingElement.psiElement.startOffset + val entireText = infillContext.enclosingElement.psiElement.readText() + val prefix = entireText.take(caretInEnclosingElement) + val suffix = + if (entireText.length < caretInEnclosingElement) "" else entireText.takeLast( + entireText.length - caretInEnclosingElement + ) + return truncateContext(prefix + suffix, infillContext) + } + + private fun truncateContext(prompt: String, infillContext: InfillContext): InfillContext { + var promptTokens = EncodingManager.getInstance().countTokens(prompt) + val truncatedContextElements = infillContext.contextElements.takeWhile { + promptTokens += it.tokens + promptTokens <= MAX_PROMPT_TOKENS + }.toSet() + return InfillContext(infillContext.enclosingElement, truncatedContextElements) + } +} \ 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 cd8c3663..91d5017c 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -72,7 +72,7 @@ + implementation="ee.carlrobert.codegpt.codecompletions.DebouncedCodeCompletionProvider"/> diff --git a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt index fbf2d0e5..937271c0 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt @@ -1,7 +1,7 @@ package ee.carlrobert.codegpt.codecompletions import com.intellij.openapi.editor.VisualPosition -import ee.carlrobert.codegpt.CodeGPTKeys.PREVIOUS_INLAY_TEXT +import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings import ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent import ee.carlrobert.llm.client.http.RequestEntity @@ -54,6 +54,6 @@ class CodeCompletionServiceTest : IntegrationTest() { myFixture.type('c') - waitExpecting { "TEST_OUTPUT" == PREVIOUS_INLAY_TEXT[myFixture.editor] } + waitExpecting { "TEST_OUTPUT" == REMAINING_EDITOR_COMPLETION[myFixture.editor] } } }