feat: make code can be re-try

unit-mesh · Apr 16, 2023 · d9d695c · d9d695c
1 parent 4fad27e
commit d9d695c
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 13 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -38,7 +38,8 @@ dependencies {
 //    implementation(libs.annotations)
     implementation("org.kohsuke:github-api:1.314")
     implementation("com.aallam.openai:openai-client:3.2.0")
-    implementation("org.jetbrains:markdown:0.4.1")
+//    implementation("org.jetbrains:markdown:0.4.1")
+    implementation("org.commonmark:commonmark:0.21.0")
 
     implementation("org.slf4j:slf4j-api:2.0.7")
     // for local testting

diff --git a/src/main/kotlin/cc/unitmesh/devti/DevtiFlow.kt b/src/main/kotlin/cc/unitmesh/devti/DevtiFlow.kt
@@ -8,6 +8,9 @@ import cc.unitmesh.devti.prompt.DevtiFlowAction
 import cc.unitmesh.devti.runconfig.DtRunState
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.diagnostic.logger
+import org.commonmark.node.*
+import org.commonmark.parser.Parser
+
 
 class DevtiFlow(
     private val kanban: Kanban,
@@ -49,9 +52,35 @@ class DevtiFlow(
         }
 
         // 3. update endpoint method
-        val code = flowAction.needUpdateMethodForController(targetEndpoint, targetController)
+        val code = fetchCode(targetEndpoint, targetController, storyDetail)
+        try {
+            analyser?.updateMethod(targetController.name, code)
+        } catch (e: Exception) {
+            logger.warn("update method failed: $e")
+            logger.warn("try to fill update method 2nd")
+
+            val code = fetchCode(targetEndpoint, targetController, storyDetail)
+            analyser?.updateMethod(targetController.name, code)
+        }
+    }
+
+    private fun fetchCode(
+        targetEndpoint: String,
+        targetController: DtClass,
+        storyDetail: String
+    ): String {
+        val content = flowAction.needUpdateMethodForController(targetEndpoint, targetController, storyDetail)
+        val code = parseCodeFromString(content)
         logger.warn("update method code: $code")
-        analyser?.updateMethod(targetController.name, code)
+        return code
+    }
+
+    private fun parseCodeFromString(markdown: String): String {
+        val parser: Parser = Parser.builder().build()
+        val node: Node = parser.parse(markdown)
+        val visitor = CodeVisitor()
+        node.accept(visitor)
+        return visitor.code
     }
 
     private fun getController(targetEndpoint: String): String? {
@@ -64,3 +93,16 @@ class DevtiFlow(
         private val logger: Logger = logger<DtRunState>()
     }
 }
+
+
+internal class CodeVisitor : AbstractVisitor() {
+    var code = ""
+
+    override fun visit(fencedCodeBlock: FencedCodeBlock?) {
+        this.code = fencedCodeBlock?.literal ?: ""
+    }
+
+    override fun visit(indentedCodeBlock: IndentedCodeBlock?) {
+        super.visit(indentedCodeBlock)
+    }
+}
diff --git a/src/main/kotlin/cc/unitmesh/devti/prompt/DevtiFlowAction.kt b/src/main/kotlin/cc/unitmesh/devti/prompt/DevtiFlowAction.kt
@@ -6,5 +6,5 @@ import cc.unitmesh.devti.kanban.SimpleProjectInfo
 interface DevtiFlowAction {
     fun fillStoryDetail(project: SimpleProjectInfo, story: String): String
     fun analysisEndpoint(storyDetail: String, classes: List<DtClass>): String
-    fun needUpdateMethodForController(targetEndpoint: String, clazz: DtClass): String
+    fun needUpdateMethodForController(targetEndpoint: String, clazz: DtClass, storyDetail: String): String
 }
diff --git a/src/main/kotlin/cc/unitmesh/devti/prompt/openai/GptPromptText.kt b/src/main/kotlin/cc/unitmesh/devti/prompt/openai/GptPromptText.kt
@@ -40,12 +40,12 @@ class GptPromptText() {
             .replace("{storyDetail}", storyDetail)
     }
 
-    fun fillUpdateMethod(targetEndpoint: String, targetClazz: DtClass): String {
+    fun fillUpdateMethod(targetEndpoint: String, targetClazz: DtClass, storyDetail: String): String {
         val promptText: InputStream = this::class.java.classLoader.getResourceAsStream("prompts/update_controller_method.txt")!!
         val promptTextString = promptText.bufferedReader().use { it.readText() }
         return promptTextString
             .replace("{controllerName}", targetClazz.name)
             .replace("{controllers}", targetClazz.format())
-            .replace("{targetEndpoint}", targetEndpoint)
+            .replace("{storyDetail}", storyDetail)
     }
 }
diff --git a/src/main/kotlin/cc/unitmesh/devti/prompt/openai/OpenAIAction.kt b/src/main/kotlin/cc/unitmesh/devti/prompt/openai/OpenAIAction.kt
@@ -4,7 +4,6 @@ import cc.unitmesh.devti.analysis.DtClass
 import cc.unitmesh.devti.kanban.SimpleProjectInfo
 import cc.unitmesh.devti.prompt.AiAction
 import cc.unitmesh.devti.prompt.DevtiFlowAction
-import cc.unitmesh.devti.runconfig.DtRunState
 import com.aallam.openai.api.BetaOpenAI
 import com.aallam.openai.api.chat.ChatCompletion
 import com.aallam.openai.api.chat.ChatCompletionRequest
@@ -52,9 +51,9 @@ class OpenAIAction(val openAIKey: String, val version: String) : AiAction, Devti
         }
     }
 
-    override fun needUpdateMethodForController(targetEndpoint: String, clazz: DtClass): String {
-        val promptText = gptPromptText.fillUpdateMethod(targetEndpoint, clazz)
-        logger.info("needUpdateMethodForController prompt text: $promptText")
+    override fun needUpdateMethodForController(targetEndpoint: String, clazz: DtClass, storyDetail: String): String {
+        val promptText = gptPromptText.fillUpdateMethod(targetEndpoint, clazz, storyDetail)
+        logger.warn("needUpdateMethodForController prompt text: $promptText")
         return runBlocking {
             val prompt = prompt(promptText)
             return@runBlocking prompt

diff --git a/src/main/resources/prompts/update_controller_method.txt b/src/main/resources/prompts/update_controller_method.txt
@@ -1,11 +1,15 @@
-请根据下面的用户故事以及 {controllerName} 的相关信息，修改代码。要求如下：
+请根据下面的用户故事以及 {controllerName} 的相关信息，编写代码。要求如下：
 
 1. 使用 Java + Spring 编写，只返回 Controller 对应方法的代码。
 2. 不做解释，只返回修改完的函数代码。
 
-```
+{controllerName} 相关信息如下：
+
+###
 {controllers}
-```
+###
+
+用户故事如下：
 
 ###
 {storyDetail}