diff --git a/core/src/main/kotlin/com/phodal/shirecore/middleware/select/SelectElementStrategy.kt b/core/src/main/kotlin/com/phodal/shirecore/middleware/select/SelectElementStrategy.kt index e94f0c4df..901386281 100644 --- a/core/src/main/kotlin/com/phodal/shirecore/middleware/select/SelectElementStrategy.kt +++ b/core/src/main/kotlin/com/phodal/shirecore/middleware/select/SelectElementStrategy.kt @@ -1,8 +1,6 @@ package com.phodal.shirecore.middleware.select import com.intellij.openapi.application.runInEdt -import com.intellij.openapi.application.runReadAction -import com.intellij.openapi.application.runWriteAction import com.intellij.openapi.diagnostic.logger import com.intellij.openapi.editor.Editor import com.intellij.openapi.project.Project @@ -131,6 +129,12 @@ sealed class SelectElementStrategy { } fun resolvePsiElement(myProject: Project, editor: Editor): PsiElement? { + val elementToAction = DefaultPsiElementStrategy().getElementToAction(myProject, editor) + + if (elementToAction != null) { + return elementToAction + } + val element: PsiElement? = try { editor.caretModel.currentCaret.offset.let { val psiFile = PsiUtilBase.getPsiFileInEditor(editor, myProject) ?: return@let null diff --git a/core/src/main/kotlin/com/phodal/shirecore/provider/ide/LocationInteractionContext.kt b/core/src/main/kotlin/com/phodal/shirecore/provider/ide/LocationInteractionContext.kt index a25700f69..4d2a8858b 100644 --- a/core/src/main/kotlin/com/phodal/shirecore/provider/ide/LocationInteractionContext.kt +++ b/core/src/main/kotlin/com/phodal/shirecore/provider/ide/LocationInteractionContext.kt @@ -2,6 +2,7 @@ package com.phodal.shirecore.provider.ide import com.intellij.openapi.editor.Editor import com.intellij.openapi.project.Project +import com.intellij.psi.PsiElement import com.phodal.shirecore.action.ShireActionLocation import com.phodal.shirecore.agent.InteractionType import kotlinx.coroutines.flow.Flow @@ -25,4 +26,6 @@ data class LocationInteractionContext( * the [com.phodal.shirecore.llm.ChatMessage] */ val prompt: String, + + val selectElement: PsiElement? = null, ) \ No newline at end of file diff --git a/core/src/main/kotlin/com/phodal/shirecore/provider/impl/CodeCompletionTask.kt b/core/src/main/kotlin/com/phodal/shirecore/provider/impl/CodeCompletionTask.kt index 460c6cee4..74fe6400d 100644 --- a/core/src/main/kotlin/com/phodal/shirecore/provider/impl/CodeCompletionTask.kt +++ b/core/src/main/kotlin/com/phodal/shirecore/provider/impl/CodeCompletionTask.kt @@ -2,10 +2,8 @@ package com.phodal.shirecore.provider.impl import com.intellij.openapi.application.invokeLater -import com.intellij.openapi.command.WriteCommandAction import com.intellij.openapi.diagnostic.logger import com.intellij.openapi.editor.Document -import com.intellij.openapi.editor.ScrollType import com.intellij.openapi.progress.ProgressIndicator import com.intellij.openapi.progress.Task import com.intellij.openapi.project.Project @@ -15,8 +13,11 @@ import com.intellij.psi.codeStyle.CodeStyleManager import com.phodal.shirecore.ShireCoreBundle import com.phodal.shirecore.ShireCoroutineScope import com.phodal.shirecore.llm.LlmProvider +import com.phodal.shirecore.markdown.Code import com.phodal.shirecore.provider.impl.dto.CodeCompletionRequest +import kotlinx.coroutines.cancel import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.flow.cancellable import kotlinx.coroutines.launch import kotlin.jvm.internal.Ref @@ -27,9 +28,10 @@ class CodeCompletionTask(private val request: CodeCompletionRequest) : private val writeActionGroupId = "code.complete.intention.write.action" private val codeMessage = ShireCoreBundle.message("intentions.chat.code.complete.name") -// private val chunksString = request.element?.let { SimilarChunksWithPaths.createQuery(it, 60) } + // private val chunksString = request.element?.let { SimilarChunksWithPaths.createQuery(it, 60) } // private val commenter = request.element?.let { LanguageCommenters.INSTANCE.forLanguage(it.language) } // private val commentPrefix = commenter?.lineCommentPrefix + private var isCanceled: Boolean = false override fun run(indicator: ProgressIndicator) { val prompt = promptText() @@ -39,27 +41,34 @@ class CodeCompletionTask(private val request: CodeCompletionRequest) : val editor = request.editor ShireCoroutineScope.scope(request.project).launch { - val currentOffset = Ref.IntRef() - currentOffset.element = request.offset + var currentOffset = request.offset val project = request.project - val finalOutput = StringBuilder() + val suggestion = StringBuilder() - flow.collect { - finalOutput.append(it) - invokeLater { - WriteCommandAction.runWriteCommandAction(project, codeMessage, writeActionGroupId, { - insertStringAndSaveChange(project, it, editor.document, currentOffset.element, false) - }) + flow.cancellable().collect { char -> + if (isCanceled) { + cancel() + return@collect + } + + val parsedContent = Code.parse(char).text; - currentOffset.element += it.length - editor.caretModel.moveToOffset(currentOffset.element) - editor.scrollingModel.scrollToCaret(ScrollType.MAKE_VISIBLE) + suggestion.append(parsedContent) + invokeLater { + if (!isCanceled && !request.isReplacement) { + InsertUtil.insertStreamingToDoc(project, parsedContent, editor, currentOffset) + currentOffset += char.length + } } } - logger.info("Suggestion: $finalOutput") - request.postExecute?.invoke(finalOutput.toString()) + if (request.isReplacement) { + InsertUtil.replaceText(project, editor, request.element, suggestion.toString()) + } + + logger.info("Suggestion: $suggestion") + request.postExecute?.invoke(suggestion.toString()) } } @@ -86,7 +95,7 @@ class CodeCompletionTask(private val request: CodeCompletionRequest) : suggestion: String, document: Document, startOffset: Int, - withReformat: Boolean + withReformat: Boolean, ) { document.insertString(startOffset, suggestion) PsiDocumentManager.getInstance(project).commitDocument(document) diff --git a/core/src/main/kotlin/com/phodal/shirecore/provider/impl/InsertUtil.kt b/core/src/main/kotlin/com/phodal/shirecore/provider/impl/InsertUtil.kt new file mode 100644 index 000000000..adff86914 --- /dev/null +++ b/core/src/main/kotlin/com/phodal/shirecore/provider/impl/InsertUtil.kt @@ -0,0 +1,58 @@ +package com.phodal.shirecore.provider.impl + +import com.intellij.openapi.application.runReadAction +import com.intellij.openapi.command.WriteCommandAction +import com.intellij.openapi.editor.Document +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.editor.ScrollType +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.TextRange +import com.intellij.psi.PsiDocumentManager +import com.intellij.psi.PsiElement +import com.intellij.psi.codeStyle.CodeStyleManager +import com.phodal.shirecore.ShireCoreBundle + +object InsertUtil { + fun insertStringAndSaveChange( + project: Project, + content: String, + document: Document, + startOffset: Int, + withReformat: Boolean, + ) { + document.insertString(startOffset, content) + PsiDocumentManager.getInstance(project).commitDocument(document) + + if (!withReformat) return + + val psiFile = PsiDocumentManager.getInstance(project).getPsiFile(document) + psiFile?.let { file -> + val reformatRange = TextRange(startOffset, startOffset + content.length) + CodeStyleManager.getInstance(project).reformatText(file, listOf(reformatRange)) + } + } + + fun insertStreamingToDoc(project: Project, char: String, editor: Editor, currentOffset: Int) { + WriteCommandAction.runWriteCommandAction( + project, + ShireCoreBundle.message("intentions.chat.code.complete.name"), + "intentions.write.action", + { + insertStringAndSaveChange(project, char, editor.document, currentOffset, false) + }) + + editor.caretModel.moveToOffset(currentOffset + char.length) + editor.scrollingModel.scrollToCaret(ScrollType.MAKE_VISIBLE) + } + + fun replaceText(project: Project, editor: Editor, element: PsiElement?, output: String) { + val primaryCaret = editor.caretModel.primaryCaret; + val start = runReadAction { primaryCaret.selectionStart } + val end = runReadAction { primaryCaret.selectionEnd } + val document = runReadAction { editor.document } + + WriteCommandAction.runWriteCommandAction(project) { + document.replaceString(start, end, output) + } + } +} \ No newline at end of file diff --git a/core/src/main/resources/messages/ShireCoreBundle.properties b/core/src/main/resources/messages/ShireCoreBundle.properties index aaef23a8a..bcb5db4f8 100644 --- a/core/src/main/resources/messages/ShireCoreBundle.properties +++ b/core/src/main/resources/messages/ShireCoreBundle.properties @@ -3,4 +3,4 @@ progress.run.task=Running task shire.ref.loading=Loading schema.custom-agent.json.display.name=Shire CustomAgent Schema intentions.request.background.process.title=Your LLM handle generate file -intentions.chat.code.complete.name=Code complete +intentions.chat.code.complete.name=Code Complete