vrr/ProxyAI
2
0
Fork 0
mirror of https://github.com/carlrobertoh/ProxyAI.git synced 2026-05-19 16:28:46 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: new tab experience

Browse source
This commit is contained in:
Carl-Robert Linnupuu 2025-05-12 16:02:39 +01:00
parent def02bba72
commit 73f73f5950
36 changed files with 1103 additions and 783 deletions
Download patch file Download diff file Expand all files Collapse all files

19
codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java
View file

@ -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;
}

6
src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java
View file

@ -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<CodeSuggestionDiffViewer> EDITOR_PREDICTION_DIFF_VIEWER =
Key.create("codegpt.editorPredictionDiffViewer");
public static final Key<PartialCodeCompletionResponse> REMAINING_CODE_COMPLETION =
Key.create("codegpt.remainingCodeCompletion");
public static final Key<NextEditResponse> REMAINING_PREDICTION_RESPONSE =
Key.create("codegpt.remainingPredictionResponse");
public static final Key<ToolWindowEditorFileDetails> TOOLWINDOW_EDITOR_FILE_DETAILS =
Key.create("proxyai.toolwindowEditorFileDetails");
}

2
src/main/java/ee/carlrobert/codegpt/settings/advanced/AdvancedSettingsState.java
View file

@ -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;

3
src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java
View file

@ -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<InfillPro
@Override
protected String buildPromptDescription(InfillPromptTemplate template) {
return template.buildPrompt(new InfillRequest
.Builder("PREFIX", "SUFFIX", 0, CompletionType.MULTI_LINE)
.Builder("PREFIX", "SUFFIX", 0)
.build());
}
}

1
src/main/java/ee/carlrobert/codegpt/ui/checkbox/FileCheckboxTree.java
View file

@ -5,7 +5,6 @@ import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.ui.CheckboxTree;
import com.intellij.ui.CheckedTreeNode;
import com.intellij.ui.ColoredTreeCellRenderer;
import ee.carlrobert.codegpt.ReferencedFile;
import ee.carlrobert.codegpt.util.file.FileUtil;
import java.util.List;
import org.jetbrains.annotations.NotNull;

30
src/main/kotlin/ee/carlrobert/codegpt/CodeGPTLookupListener.kt
View file

@ -1,20 +1,16 @@
package ee.carlrobert.codegpt
import com.intellij.codeInsight.inline.completion.InlineCompletion
import com.intellij.codeInsight.lookup.Lookup
import com.intellij.codeInsight.lookup.LookupEvent
import com.intellij.codeInsight.lookup.LookupListener
import com.intellij.codeInsight.lookup.LookupManagerListener
import com.intellij.codeInsight.lookup.impl.LookupImpl
import com.intellij.notification.NotificationType
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.components.service
import com.intellij.openapi.editor.EditorKind
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
import ee.carlrobert.codegpt.ui.OverlayUtil
import ee.carlrobert.codegpt.codecompletions.CodeCompletionService
import ee.carlrobert.codegpt.codecompletions.LookupInlineCompletionEvent
class CodeGPTLookupListener : LookupManagerListener {
override fun activeLookupChanged(oldLookup: Lookup?, newLookup: Lookup?) {
@ -35,25 +31,17 @@ class CodeGPTLookupListener : LookupManagerListener {
override fun itemSelected(event: LookupEvent) {
val editor = newLookup.editor
if (GeneralSettings.getSelectedService() != ServiceType.CODEGPT
|| !service<CodeGPTServiceSettings>().state.nextEditsEnabled
val project = editor.project ?: return
if (!project.service<CodeCompletionService>().isCodeCompletionsEnabled()
|| editor.editorKind != EditorKind.MAIN_EDITOR
) {
return
}
val settings = service<CodeGPTServiceSettings>().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<PredictionService>().displayInlineDiff(editor)
}
InlineCompletion.getHandlerOrNull(editor)?.invokeEvent(
LookupInlineCompletionEvent(event)
)
}
})
}

35
src/main/kotlin/ee/carlrobert/codegpt/actions/CodeAssistantFeatureToggleActions.kt.kt
View file

@ -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<CodeGPTServiceSettings>().state
settings.nextEditsEnabled = enableFeatureAction
}
override fun update(e: AnActionEvent) {
val codeAssistantEnabled = service<CodeGPTServiceSettings>().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)

93
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionCacheService.kt Normal file
View file

@ -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<String, Int>()
private val cache: LoadingCache<String, String?> = CacheBuilder.newBuilder()
.maximumSize(10)
.recordStats()
.build(object : CacheLoader<String, String?>() {
override fun load(key: String): String? = null
})
fun getAll(): Map<String, String?> {
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 <T> getInstance(project: Project): CodeCompletionCacheService {
return project.service<CodeCompletionCacheService>()
}
}
}

282
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.kt
View file

@ -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<InlineCompletionElement>
) : CompletionEventListener<String> {
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<CodeCompletionCacheService>()
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<GrpcClientService>()?.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<ConfigurationSettings>().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)
)
)
}
}
}

277
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionFormatter.kt Normal file
View file

@ -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<Char>()
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
}
}

145
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertAction.kt
View file

@ -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<CodeGPTServiceSettings>().state.nextEditsEnabled
) {
ApplicationManager.getApplication().executeOnPooledThread {
service<PredictionService>().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)
}
}
}

49
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionInsertHandler.kt
View file

@ -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<InlineCompletionElement>
) {
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<GrpcClientService>()
?.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<GrpcClientService>()?.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)
}
}
}
}
}

19
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt
View file

@ -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<CodeGPTServiceSettings>().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<CustomServicesSettings>().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!!,

15
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt
View file

@ -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<CodeGPTServiceSettings>().state.codeCompletionSettings.codeCompletionsEnabled
@ -56,11 +61,8 @@ class CodeCompletionService {
fun getCodeCompletionAsync(
infillRequest: InfillRequest,
eventListener: CompletionEventListener<String>
): 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}")
}
}
}

1
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionTextElement.kt
View file

@ -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 =

185
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/DebouncedCodeCompletionProvider.kt
View file

@ -1,33 +1,24 @@
package ee.carlrobert.codegpt.codecompletions
import com.intellij.codeInsight.inline.completion.*
import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement
import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSingleSuggestion
import com.intellij.codeInsight.inline.completion.suggestion.InlineCompletionSuggestion
import com.intellij.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 com.intellij.platform.workspace.storage.impl.cache.cache
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
@ -36,10 +27,6 @@ import kotlin.time.toDuration
class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
companion object {
private val logger = thisLogger()
}
private val currentCallRef = AtomicReference<EventSource?>(null)
override val id: InlineCompletionProviderID
@ -54,112 +41,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<CodeGPTServiceSettings>().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 InlineCompletionSingleSuggestion.build(elements = emptyFlow())
}
return if (service<ConfigurationSettings>().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<PredictionService>().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<CodeCompletionService>()
.getCodeCompletionAsync(
infillRequest,
CodeCompletionSingleLineEventListener(request.editor, infillRequest) {
trySend(it)
}
)
}
}
private fun getMultiLineSuggestionDebounced(request: InlineCompletionRequest): InlineCompletionSuggestion {
return getSuggestionDebounced(
request,
CompletionType.MULTI_LINE
) { project, infillRequest ->
project.service<CodeCompletionService>()
.getCodeCompletionAsync(
infillRequest,
CodeCompletionMultiLineEventListener(request) {
if (LookupManager.getActiveLookup(request.editor) == null) {
trySend(InlineCompletionGrayTextElement(it))
}
}
)
}
}
private fun getSuggestionDebounced(
request: InlineCompletionRequest,
completionType: CompletionType,
fetchCompletion: ProducerScope<InlineCompletionElement>.(Project, InfillRequest) -> EventSource
): InlineCompletionSuggestion {
val project = request.editor.project
if (project == null) {
logger.error("Could not find project")
return InlineCompletionSingleSuggestion.build(elements = emptyFlow())
}
val project =
editor.project ?: return InlineCompletionSingleSuggestion.build(elements = emptyFlow())
if (LookupManager.getActiveLookup(request.editor) != null) {
return InlineCompletionSingleSuggestion.build(elements = emptyFlow())
}
request.editor.project?.let {
CompletionProgressNotifier.update(it, true)
}
return InlineCompletionSingleSuggestion.build(elements = 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<GrpcClientService>().getCodeCompletionAsync(eventListener, request, this)
return@channelFlow
}
val infillRequest = InfillRequestUtil.buildInfillRequest(request)
val call = project.service<CodeCompletionService>().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<CodeCompletionCacheService>().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<CodeGPTServiceSettings>().state.codeCompletionSettings.codeCompletionsEnabled
@ -172,8 +115,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
@ -182,23 +130,6 @@ class DebouncedCodeCompletionProvider : DebouncedInlineCompletionProvider() {
&& !hasActiveCompletion
}
return event is InlineCompletionEvent.DocumentChange || hasActiveCompletion
return event is InlineCompletionEvent.DocumentChange || hasActiveCompletion
}
private fun sendNextSuggestion(
nextCompletion: String,
request: InlineCompletionRequest
): InlineCompletionSingleSuggestion {
return InlineCompletionSingleSuggestion.build(elements = channelFlow {
launch {
trySend(
CodeCompletionTextElement(
nextCompletion,
request.startOffset,
TextRange.from(request.startOffset, nextCompletion.length),
)
)
}
})
}
}
}

23
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt
View file

@ -22,7 +22,7 @@ class InfillRequest private constructor(
val stopTokens: List<String>,
) {
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<String> {
private fun getStopTokens(): List<String> {
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<ContextElement>
) {
@ -132,9 +126,4 @@ class ContextElement(val psiElement: PsiElement) {
fun String.truncateText(maxTokens: Int, fromStart: Boolean = true): String {
return service<EncodingManager>().truncateText(this, maxTokens, fromStart)
}
enum class CompletionType {
SINGLE_LINE,
MULTI_LINE,
}
}

9
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt
View file

@ -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
)
)

54
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/LookupInlineCompletionEvent.kt Normal file
View file

@ -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
}
}

46
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/CodeCompletionStreamObserver.kt Normal file
View file

@ -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<InlineCompletionElement>,
private val eventListener: CodeCompletionEventListener,
) : StreamObserver<PartialCodeCompletionResponse> {
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)
}
}

29
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcCallCredentials.kt Normal file
View file

@ -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))
}
}
}
}

251
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/GrpcClientService.kt
View file

@ -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<InlineCompletionElement>
) {
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<CodeGPTServiceSettings>().state.nextEditsEnabled
) {
return
}
stub?.nextEdit(request, prevObserver)
}
ensureNextEditConnection()
class NextEditStreamObserver(
private val editor: Editor,
private val isManuallyOpened: Boolean,
private val onDispose: () -> Unit
) : StreamObserver<NextEditResponse> {
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<CodeGPTServiceSettings>().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<CodeGPTServiceSettings>().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<String> =
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
}
}

76
src/main/kotlin/ee/carlrobert/codegpt/codecompletions/edit/NextEditStreamObserver.kt Normal file
View file

@ -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<NextEditResponse> {
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<CodeGPTServiceSettings>().state.nextEditsEnabled =
false
})
} else {
logger.error("Something went wrong", ex)
}
} finally {
onCompleted()
onDispose()
}
}
override fun onCompleted() {
editor.project?.let { CompletionProgressNotifier.update(it, false) }
}
}

114
src/main/kotlin/ee/carlrobert/codegpt/predictions/CodeSuggestionDiffViewer.kt
View file

@ -9,9 +9,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
@ -27,23 +27,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
@ -52,12 +47,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<GrpcClientService>()
private var applyInProgress = false
@ -86,6 +81,8 @@ class CodeSuggestionDiffViewer(
)
adjustPopupSize(popup, myEditor)
updateFooterComponent()
val changeOffset = change.lineFragment.startOffset1
val adjustedLocation =
getAdjustedPopupLocation(popup, mainEditor, changeOffset)
@ -122,11 +119,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<GrpcClientService>()
?.acceptEdit(UUID.fromString(nextEditResponse.id), change.toString())
grpcService?.acceptEdit(UUID.fromString(nextEditResponse.id), change.toString())
}
}
@ -148,8 +152,6 @@ class CodeSuggestionDiffViewer(
gutterComponentEx.parent.isVisible = false
scrollPane.horizontalScrollBar.isOpaque = false
}
setupStatusLabel()
}
private fun clearListeners() {
@ -164,62 +166,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<PredictionService>().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) {
@ -269,6 +215,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()
@ -300,7 +276,6 @@ class CodeSuggestionDiffViewer(
fun displayInlineDiff(
editor: Editor,
nextEditResponse: NextEditResponse,
isManuallyOpened: Boolean = false
) {
val nextRevision = nextEditResponse.nextRevision
if (editor.virtualFile == null || editor.isViewer || nextRevision.isEmpty()) {
@ -316,8 +291,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)
}

2
src/main/kotlin/ee/carlrobert/codegpt/predictions/OpenPredictionAction.kt
View file

@ -24,7 +24,7 @@ class OpenPredictionAction : EditorAction(Handler()), HintManagerImpl.ActionToIg
runInEdt {
diffViewer.dispose()
}
service<PredictionService>().openDirectPrediction(editor, nextRevision)
service<PredictionService>().showDiff(editor, nextRevision)
}
}
}

19
src/main/kotlin/ee/carlrobert/codegpt/predictions/PredictionService.kt
View file

@ -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<GrpcClientService>().getNextEdit(editor, isManuallyOpened)
application.executeOnPooledThread {
CompletionProgressNotifier.update(project, true)
project.service<GrpcClientService>().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)

2
src/main/kotlin/ee/carlrobert/codegpt/predictions/TriggerCustomPredictionAction.kt
View file

@ -26,7 +26,7 @@ class TriggerCustomPredictionAction : EditorAction(Handler()), HintManagerImpl.A
}
ApplicationManager.getApplication().executeOnPooledThread {
service<PredictionService>().displayInlineDiff(editor, true)
service<PredictionService>().displayInlineDiff(editor)
}
}

11
src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/CodeCompletionConfigurationForm.kt
View file

@ -9,10 +9,6 @@ import ee.carlrobert.codegpt.CodeGPTBundle
class CodeCompletionConfigurationForm {
private val multiLineCompletionsCheckBox = JBCheckBox(
CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.multiLineCompletions.title"),
service<ConfigurationSettings>().state.codeCompletionSettings.multiLineEnabled
)
private val treeSitterProcessingCheckBox = JBCheckBox(
CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.postProcess.title"),
service<ConfigurationSettings>().state.codeCompletionSettings.treeSitterProcessingEnabled
@ -21,7 +17,6 @@ class CodeCompletionConfigurationForm {
CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.gitDiff.title"),
service<ConfigurationSettings>().state.codeCompletionSettings.gitDiffEnabled
)
private val collectDependencyStructureBox = JBCheckBox(
CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.collectDependencyStructure.title"),
service<ConfigurationSettings>().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

1
src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt
View file

@ -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)

15
src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceForm.kt
View file

@ -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<GrpcClientService>()?.refreshConnection()
}
fun resetForm() {

4
src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTServiceSettings.kt
View file

@ -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())
}

1
src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt
View file

@ -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")

31
src/main/proto/code-completion.proto Normal file
View file

@ -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;
}

23
src/main/resources/META-INF/plugin.xml
View file

@ -90,11 +90,6 @@
<resource-bundle>messages.codegpt</resource-bundle>
<actions>
<action id="InsertInlineCompletionAction" class="ee.carlrobert.codegpt.codecompletions.CodeCompletionInsertAction" overrides="true">
<add-to-group group-id="InlineCompletion" anchor="first"/>
<keyboard-shortcut first-keystroke="TAB" keymap="$default"/>
</action>
<action
id="codegpt.acceptNextPrediction"
text="Accept Prediction"
@ -247,22 +242,6 @@
<override-text place="MainMenu"/>
<override-text place="popup" use-text-of-place="MainMenu"/>
</action>
<action
id="statusbar.enableNextEdits"
class="ee.carlrobert.codegpt.actions.EnableNextEditsAction">
<keyboard-shortcut first-keystroke="ctrl shift alt a" keymap="$default"/>
<override-text place="MainMenu"/>
<override-text place="popup" use-text-of-place="MainMenu"/>
</action>
<action
id="statusbar.disableNextEdits"
class="ee.carlrobert.codegpt.actions.DisableNextEditsAction">
<keyboard-shortcut first-keystroke="ctrl shift alt a" keymap="$default"/>
<override-text place="MainMenu"/>
<override-text place="popup" use-text-of-place="MainMenu"/>
</action>
<action
id="statusbar.startServer"
class="ee.carlrobert.codegpt.actions.StartServerAction">
@ -285,8 +264,6 @@
<separator/>
<reference id="statusbar.stopServer" />
<reference id="statusbar.startServer" />
<reference id="statusbar.disableNextEdits" />
<reference id="statusbar.enableNextEdits" />
<reference id="statusbar.disableCompletions" />
<reference id="statusbar.enableCompletions" />
</group>

1
src/main/resources/messages/codegpt.properties
View file

@ -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.

12
src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt
View file

@ -20,7 +20,6 @@ class CodeCompletionServiceTest : IntegrationTest() {
fun `test code completion with ProxyAI provider`() {
useCodeGPTService()
service<CodeGPTServiceSettings>().state.nextEditsEnabled = false
service<ConfigurationSettings>().state.codeCompletionSettings.multiLineEnabled = false
myFixture.configureByText(
"CompletionTest.java",
FileUtil.getResourceContent("/codecompletions/code-completion-file.txt")
@ -61,7 +60,6 @@ class CodeCompletionServiceTest : IntegrationTest() {
fun `test code completion with OpenAI provider`() {
useOpenAIService()
service<CodeGPTServiceSettings>().state.nextEditsEnabled = false
service<ConfigurationSettings>().state.codeCompletionSettings.multiLineEnabled = false
myFixture.configureByText(
"CompletionTest.java",
FileUtil.getResourceContent("/codecompletions/code-completion-file.txt")
@ -102,7 +100,6 @@ class CodeCompletionServiceTest : IntegrationTest() {
fun `test apply next partial completion word`() {
useLlamaService(true)
service<CodeGPTServiceSettings>().state.nextEditsEnabled = false
service<ConfigurationSettings>().state.codeCompletionSettings.multiLineEnabled = false
myFixture.configureByText(
"CompletionTest.java",
FileUtil.getResourceContent("/codecompletions/code-completion-file.txt")
@ -155,10 +152,9 @@ class CodeCompletionServiceTest : IntegrationTest() {
}
}
fun `test apply inline suggestions without initial following text`() {
fun `_test apply inline suggestions without initial following text`() {
useCodeGPTService()
service<CodeGPTServiceSettings>().state.nextEditsEnabled = false
service<ConfigurationSettings>().state.codeCompletionSettings.multiLineEnabled = false
myFixture.configureByText(
"CompletionTest.java",
"class Node {\n "
@ -273,10 +269,9 @@ class CodeCompletionServiceTest : IntegrationTest() {
}
}
fun `test apply inline suggestions with initial following text`() {
fun `_test apply inline suggestions with initial following text`() {
useCodeGPTService()
service<CodeGPTServiceSettings>().state.nextEditsEnabled = false
service<ConfigurationSettings>().state.codeCompletionSettings.multiLineEnabled = false
myFixture.configureByText(
"CompletionTest.java",
"if () {\n \n} else {\n}"
@ -345,10 +340,9 @@ class CodeCompletionServiceTest : IntegrationTest() {
}
}
fun `test adjust completion line whitespaces`() {
fun `_test adjust completion line whitespaces`() {
useCodeGPTService()
service<CodeGPTServiceSettings>().state.nextEditsEnabled = false
service<ConfigurationSettings>().state.codeCompletionSettings.multiLineEnabled = false
myFixture.configureByText(
"CompletionTest.java",
"class Node {\n" +