diff --git a/src/main/kotlin/ai/devchat/plugin/completion/agent/Agent.kt b/src/main/kotlin/ai/devchat/plugin/completion/agent/Agent.kt index ee4c5c5..2f92afe 100644 --- a/src/main/kotlin/ai/devchat/plugin/completion/agent/Agent.kt +++ b/src/main/kotlin/ai/devchat/plugin/completion/agent/Agent.kt @@ -62,6 +62,8 @@ class Agent(val scope: CoroutineScope) { @SerializedName("completion_id") val completionId: String, @SerializedName("lines") val lines: Int, @SerializedName("length") val length: Int, + @SerializedName("ide") val ide: String, + @SerializedName("language") val language: String ) { enum class EventType { @SerializedName("view") VIEW, @@ -143,7 +145,7 @@ class Agent(val scope: CoroutineScope) { val devChatEndpoint = CONFIG["providers.devchat.api_base"] as? String val devChatAPIKey = CONFIG["providers.devchat.api_key"] as? String val endpoint = "$devChatEndpoint/completions" - val endingChunk = "data:[DONE]" + val endingChunk = "[DONE]" val payload = mapOf( "model" to ((CONFIG["complete_model"] as? String) ?: defaultCompletionModel), "prompt" to prompt, @@ -161,8 +163,10 @@ class Agent(val scope: CoroutineScope) { response.body?.charStream()?.buffered()?.use {reader -> reader.lineSequence().asFlow() .filter {it.isNotEmpty()} - .takeWhile { it.startsWith("data:") && it != endingChunk} - .map { gson.fromJson(it.drop(5).trim(), CompletionResponseChunk::class.java) } + .takeWhile { it.startsWith("data:") } + .map { it.drop(5).trim() } + .takeWhile { it.uppercase() != endingChunk } + .map { gson.fromJson(it, CompletionResponseChunk::class.java) } .takeWhile {it != null} .collect { emit(CodeCompletionChunk(it.id, it.choices[0].text!!)) } } @@ -172,11 +176,7 @@ class Agent(val scope: CoroutineScope) { private fun toLines(chunks: Flow): Flow = flow { var ongoingLine = "" var latestId = "" - chunks.catch { - if (ongoingLine.isNotEmpty()) { - emit(CodeCompletionChunk(latestId, ongoingLine)) - } - }.collect { chunk -> + chunks.catch { logger.warn(it) }.collect { chunk -> var remaining = chunk.text while (remaining.contains(LINE_SEPARATOR)) { val parts = remaining.split(LINE_SEPARATOR, limit = 2) @@ -187,6 +187,7 @@ class Agent(val scope: CoroutineScope) { ongoingLine += remaining latestId = chunk.id } + if (ongoingLine.isNotEmpty()) emit(CodeCompletionChunk(latestId, ongoingLine)) } private fun stopAtFirstBrace(chunks: Flow): Flow = flow { @@ -264,7 +265,7 @@ class Agent(val scope: CoroutineScope) { completionRequest.filepath, completionRequest.text, completionRequest.position - ).createPrompt() + ).createPrompt(CONFIG["complete_model"] as? String) scope.launch { val chunks = request(prompt) diff --git a/src/main/kotlin/ai/devchat/plugin/completion/agent/ContextBuilder.kt b/src/main/kotlin/ai/devchat/plugin/completion/agent/ContextBuilder.kt index 8107489..b38d6ac 100644 --- a/src/main/kotlin/ai/devchat/plugin/completion/agent/ContextBuilder.kt +++ b/src/main/kotlin/ai/devchat/plugin/completion/agent/ContextBuilder.kt @@ -80,8 +80,11 @@ class ContextBuilder(val filepath: String, val content: String, val offset: Int) ) } - fun createPrompt(): String { + fun createPrompt(model: String?): String { val (prefix, suffix) = buildFileContext() - return "$filepath\n\n$prefix$suffix" + return if (!model.isNullOrEmpty() && model.contains("deepseek")) + "<|fim▁begin|>$filepath\n\n$prefix<|fim▁hole|>$suffix<|fim▁end|>" + else + "$filepath\n\n$prefix$suffix" } } \ No newline at end of file diff --git a/src/main/kotlin/ai/devchat/plugin/completion/editor/InlineCompletionService.kt b/src/main/kotlin/ai/devchat/plugin/completion/editor/InlineCompletionService.kt index bbfa7d8..da94cf6 100644 --- a/src/main/kotlin/ai/devchat/plugin/completion/editor/InlineCompletionService.kt +++ b/src/main/kotlin/ai/devchat/plugin/completion/editor/InlineCompletionService.kt @@ -20,6 +20,7 @@ import com.intellij.ui.JBColor import com.intellij.util.ui.UIUtil import ai.devchat.plugin.completion.agent.Agent import ai.devchat.plugin.completion.agent.AgentService +import com.intellij.openapi.fileEditor.FileDocumentManager import com.intellij.psi.PsiDocumentManager import com.intellij.psi.codeStyle.CodeStyleManager import kotlinx.coroutines.launch @@ -154,6 +155,7 @@ class InlineCompletionService { shownInlineCompletion = InlineCompletion(editor, offset, completion, inlays, markups, id, displayAt) val agentService = service() + val virtualFile = FileDocumentManager.getInstance().getFile(editor.document) agentService.scope.launch { agentService.postEvent( Agent.LogEventRequest( @@ -161,6 +163,8 @@ class InlineCompletionService { completionId = completion.id, lines = textLines.size, length = text.length, + ide = "intellij", + language = virtualFile?.extension ?: "" ) ) } @@ -217,13 +221,16 @@ class InlineCompletionService { currentCompletion.inlays.forEach(Disposer::dispose) } val agentService = service() + val virtualFile = FileDocumentManager.getInstance().getFile(currentCompletion.editor.document) agentService.scope.launch { agentService.postEvent( Agent.LogEventRequest( type = Agent.LogEventRequest.EventType.SELECT, completionId = currentCompletion.completion.id, lines = text.lines().size, - length = text.length + length = text.length, + ide = "intellij", + language = virtualFile?.extension ?: "" ) ) }