feat(core): add postExecute callback to code execution tasks #29

This commit introduces a postExecute callback to the code execution tasks. This callback is invoked after the task execution, passing the final output as a parameter. This change affects the CodeCompletionTask, EditorInteractionProvider, CodeCompletionRequest, FileGenerateTask, and LocationInteractionProvider classes. It also modifies the LocationInteractionContext to include a prompt string.
phodal · Jul 1, 2024 · af6567b · af6567b
1 parent 10cfc67
commit af6567b
Show file tree

Hide file tree

Showing 8 changed files with 45 additions and 27 deletions.
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
@@ -20,4 +20,9 @@ data class LocationInteractionContext(
     val editor: Editor?,
 
     val project: Project,
+
+    /**
+     * the [com.phodal.shirecore.llm.ChatMessage]
+     */
+    val prompt: String,
 )
diff --git a/core/src/main/kotlin/com/phodal/shirecore/provider/ide/LocationInteractionProvider.kt b/core/src/main/kotlin/com/phodal/shirecore/provider/ide/LocationInteractionProvider.kt
@@ -12,7 +12,7 @@ import com.intellij.openapi.extensions.ExtensionPointName
 interface LocationInteractionProvider {
     fun isApplicable(context: LocationInteractionContext): Boolean
 
-    fun execute(context: LocationInteractionContext) : String
+    fun execute(context: LocationInteractionContext, postExecute: (String) -> Unit)
 
     companion object {
         private val EP_NAME: ExtensionPointName<LocationInteractionProvider> =

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
@@ -43,10 +43,10 @@ class CodeCompletionTask(private val request: CodeCompletionRequest) :
             currentOffset.element = request.offset
 
             val project = request.project
-            val suggestion = StringBuilder()
+            val finalOutput = StringBuilder()
 
             flow.collect {
-                suggestion.append(it)
+                finalOutput.append(it)
                 invokeLater {
                     WriteCommandAction.runWriteCommandAction(project, codeMessage, writeActionGroupId, {
                         insertStringAndSaveChange(project, it, editor.document, currentOffset.element, false)
@@ -58,7 +58,8 @@ class CodeCompletionTask(private val request: CodeCompletionRequest) :
                 }
             }
 
-            logger.info("Suggestion: $suggestion")
+            logger.info("Suggestion: $finalOutput")
+            request.postExecute?.invoke(finalOutput.toString())
         }
     }
 

diff --git a/core/src/main/kotlin/com/phodal/shirecore/provider/impl/EditorInteractionProvider.kt b/core/src/main/kotlin/com/phodal/shirecore/provider/impl/EditorInteractionProvider.kt
@@ -1,7 +1,6 @@
 package com.phodal.shirecore.provider.impl
 
 import com.intellij.openapi.application.runReadAction
-import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator
 import com.phodal.shirecore.ShirelangNotifications
@@ -18,20 +17,22 @@ class EditorInteractionProvider : LocationInteractionProvider {
         return true
     }
 
-    override fun execute(context: LocationInteractionContext): String {
-        val msgs: List<ChatMessage> = listOf()
+    override fun execute(context: LocationInteractionContext, postExecute: (String) -> Unit) {
+        val msgs: List<ChatMessage> = listOf(
+            // todo: add system prompt
+            ChatMessage(ChatRole.User, context.prompt),
+        )
         val targetFile = context.editor?.virtualFile
-        val result: String = ""
 
         when (context.interactionType) {
             InteractionType.AppendCursor,
             InteractionType.AppendCursorStream,
             -> {
-                val task = createTask(context, msgs, isReplacement = false)
+                val task = createTask(context, msgs, isReplacement = false, postExecute = postExecute)
 
                 if (task == null) {
                     ShirelangNotifications.error(context.project, "Failed to create code completion task.")
-                    return ""
+                    return
                 }
 
                 ProgressManager.getInstance()
@@ -40,17 +41,17 @@ class EditorInteractionProvider : LocationInteractionProvider {
 
             InteractionType.OutputFile -> {
                 val fileName = targetFile?.name
-                val task = FileGenerateTask(context.project, msgs, fileName)
+                val task = FileGenerateTask(context.project, msgs, fileName, postExecute = postExecute)
                 ProgressManager.getInstance()
                     .runProcessWithProgressAsynchronously(task, BackgroundableProcessIndicator(task))
             }
 
             InteractionType.ReplaceSelection -> {
-                val task = createTask(context, msgs, true)
+                val task = createTask(context, msgs, true, postExecute)
 
                 if (task == null) {
                     ShirelangNotifications.error(context.project, "Failed to create code completion task.")
-                    return ""
+                    return
                 }
 
                 ProgressManager.getInstance()
@@ -59,7 +60,7 @@ class EditorInteractionProvider : LocationInteractionProvider {
 
             InteractionType.ReplaceCurrentFile -> {
                 val fileName = targetFile?.name
-                val task = FileGenerateTask(context.project, msgs, fileName)
+                val task = FileGenerateTask(context.project, msgs, fileName, postExecute = postExecute)
 
                 ProgressManager.getInstance()
                     .runProcessWithProgressAsynchronously(task, BackgroundableProcessIndicator(task))
@@ -69,14 +70,13 @@ class EditorInteractionProvider : LocationInteractionProvider {
                 TODO()
             }
         }
-
-        return result
     }
 
     private fun createTask(
         context: LocationInteractionContext,
         msgs: List<ChatMessage>,
         isReplacement: Boolean,
+        postExecute: (String) -> Unit,
     ): CodeCompletionTask? {
         if (context.editor == null) {
             ShirelangNotifications.error(context.project, "Editor is null, please open a file to continue.")
@@ -90,7 +90,15 @@ class EditorInteractionProvider : LocationInteractionProvider {
         val userPrompt = msgs.filter { it.role == ChatRole.User }.joinToString("\n") { it.content }
 
         val request = runReadAction {
-            CodeCompletionRequest.create(editor, offset, element, null, userPrompt, isReplacement = isReplacement)
+            CodeCompletionRequest.create(
+                editor,
+                offset,
+                element,
+                null,
+                userPrompt,
+                isReplacement = isReplacement,
+                postExecute = postExecute
+            )
         } ?: return null
 
         val task = CodeCompletionTask(request)

diff --git a/core/src/main/kotlin/com/phodal/shirecore/provider/impl/FileGenerateTask.kt b/core/src/main/kotlin/com/phodal/shirecore/provider/impl/FileGenerateTask.kt
@@ -25,7 +25,8 @@ open class FileGenerateTask(
     val messages: List<ChatMessage>,
     val fileName: String?,
     private val codeOnly: Boolean = false,
-    private val taskName: String = ShireCoreBundle.message("intentions.request.background.process.title")
+    private val taskName: String = ShireCoreBundle.message("intentions.request.background.process.title"),
+    val postExecute: ((String) -> Unit)?
 ) :
     Task.Backgroundable(project, taskName) {
     private val projectRoot = project.guessProjectDir()!!
@@ -66,6 +67,8 @@ open class FileGenerateTask(
             file.writeText(result)
             refreshAndOpenInEditor(Path(projectRoot.path), projectRoot)
         }
+
+        postExecute?.invoke(result)
     }
 
     private fun refreshAndOpenInEditor(file: Path, parentDir: VirtualFile) {

diff --git a/core/src/main/kotlin/com/phodal/shirecore/provider/impl/dto/CodeCompletionRequest.kt b/core/src/main/kotlin/com/phodal/shirecore/provider/impl/dto/CodeCompletionRequest.kt
@@ -21,6 +21,7 @@ class CodeCompletionRequest(
     val editor: Editor,
     val suffixText: String,
     val isReplacement: Boolean = false,
+    val postExecute: ((String) -> Unit)?,
 ) : Disposable {
     companion object {
         fun create(
@@ -29,7 +30,8 @@ class CodeCompletionRequest(
             element: PsiElement?,
             prefix: String?,
             suffix: String?,
-            isReplacement: Boolean = false
+            isReplacement: Boolean = false,
+            postExecute: (String) -> Unit,
         ): CodeCompletionRequest? {
             val project = editor.project ?: return null
             val document = editor.document
@@ -56,7 +58,8 @@ class CodeCompletionRequest(
                 element,
                 editor,
                 suffixText,
-                isReplacement
+                isReplacement,
+                postExecute = postExecute
             )
 
         }

diff --git a/shirelang/src/main/kotlin/com/phodal/shirelang/run/runner/ShireDefaultRunner.kt b/shirelang/src/main/kotlin/com/phodal/shirelang/run/runner/ShireDefaultRunner.kt
@@ -29,14 +29,16 @@ class ShireDefaultRunner(
                 interactionType = InteractionType.AppendCursorStream,
                 editor = context.editor,
                 project = context.myProject,
+                prompt = context.prompt
             )
 
             if (context.hole?.interaction != null) {
                 val interactionProvider = LocationInteractionProvider.provide(interactionContext)
                 if (interactionProvider != null) {
-                    val response = interactionProvider.execute(interactionContext)
-                    postFunction(response)
-                    context.processHandler.detachProcess()
+                    interactionProvider.execute(interactionContext) {
+                        postFunction(it)
+                        context.processHandler.detachProcess()
+                    }
                     return@invokeLater
                 }
             }

diff --git a/shirelang/src/main/kotlin/com/phodal/shirelang/run/runner/ShireRunner.kt b/shirelang/src/main/kotlin/com/phodal/shirelang/run/runner/ShireRunner.kt
@@ -4,10 +4,6 @@ import com.intellij.execution.process.ProcessHandler
 import com.intellij.execution.ui.ConsoleView
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
-import com.phodal.shirecore.action.ShireActionLocation
-import com.phodal.shirecore.agent.InteractionType
-import com.phodal.shirecore.provider.ide.LocationInteractionContext
-import com.phodal.shirecore.provider.ide.LocationInteractionProvider
 import com.phodal.shirelang.compiler.hobbit.HobbitHole
 import com.phodal.shirelang.run.ShireConfiguration