feat: init for psi code

unit-mesh · Jul 14, 2023 · b459fa4 · b459fa4
1 parent 59132c0
commit b459fa4
Show file tree

Hide file tree

Showing 7 changed files with 73 additions and 15 deletions.
diff --git a/src/main/kotlin/cc/unitmesh/devti/connector/openai/PromptGenerator.kt b/src/main/kotlin/cc/unitmesh/devti/connector/openai/PromptGenerator.kt
@@ -3,6 +3,7 @@ package cc.unitmesh.devti.connector.openai
 import cc.unitmesh.devti.analysis.DtClass
 import cc.unitmesh.devti.flow.model.SimpleProjectInfo
 import cc.unitmesh.devti.prompting.PromptConfig
+import com.intellij.openapi.util.NlsSafe
 import java.io.InputStream
 
 class PromptGenerator {
@@ -48,7 +49,7 @@ class PromptGenerator {
             .replace("{controllers}", targetClazz.format())
             .replace("{storyDetail}", storyDetail)
             .replace("{models}", models.joinToString("\n", transform = DtClass::formatDto))
-            .replace("{services}", models.joinToString("\n", transform = DtClass::format))
+            .replace("{services}", services.joinToString("\n", transform = DtClass::format))
             .replace("{spec}", spec ?: "")
     }
 
@@ -65,4 +66,11 @@ class PromptGenerator {
         return promptTextString
             .replace("{code}", text)
     }
+
+    fun createServiceAndRepository(controller: @NlsSafe String): String {
+        val promptText: InputStream = getResource("create_service_and_repository")!!
+        val promptTextString = promptText.bufferedReader().use { it.readText() }
+        return promptTextString
+            .replace("{controllerCode}", controller)
+    }
 }
diff --git a/src/main/kotlin/cc/unitmesh/devti/flow/AutoDevFlow.kt b/src/main/kotlin/cc/unitmesh/devti/flow/AutoDevFlow.kt
@@ -11,8 +11,10 @@ import cc.unitmesh.devti.flow.model.TargetEndpoint
 import cc.unitmesh.devti.gui.chat.ChatCodingComponent
 import cc.unitmesh.devti.parser.parseCodeFromString
 import cc.unitmesh.devti.runconfig.AutoDevRunProfileState
+import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.diagnostic.logger
+import com.intellij.psi.PsiFile
 import kotlinx.coroutines.runBlocking
 
 class AutoDevFlow(
@@ -22,6 +24,7 @@ class AutoDevFlow(
     val ui: ChatCodingComponent,
 ) : DevtiFlowAction {
     private val promptGenerator = PromptGenerator()
+    var selectedControllerName = ""
 
     /**
      * Step 1: check story detail is valid, if not, fill story detail
@@ -86,9 +89,10 @@ class AutoDevFlow(
     }
 
     /**
-     * Step 3: update endpoint method
+     * Step 4: update endpoint method
      */
     fun updateEndpointMethod(target: TargetEndpoint, storyDetail: String) {
+        selectedControllerName = target.controller.name
         try {
             doExecuteUpdateEndpoint(target, storyDetail)
         } catch (e: Exception) {
@@ -97,6 +101,31 @@ class AutoDevFlow(
         }
     }
 
+    /**
+     * Step 5: create service and repository
+     */
+    fun createServiceAndRepository() {
+        // filter controllerName == selectedControllerName
+        val files: List<PsiFile> = processor?.getAllControllerFiles()?.filter { it.name == selectedControllerName }
+            ?: emptyList()
+        val controller = files.firstOrNull() ?: return
+
+        val controllerCode = runReadAction {
+            controller.text
+        }
+
+        val promptText = promptGenerator.createServiceAndRepository(controllerCode)
+
+        logger.warn("createServiceAndController prompt text: $promptText")
+        val result = executePrompt(promptText)
+
+        val services = parseCodeFromString(result)
+        services.forEach { service ->
+            processor?.let { createCodeByType(service, it, true) }
+        }
+    }
+
+
     private fun doExecuteUpdateEndpoint(target: TargetEndpoint, storyDetail: String) {
         val codes = fetchForEndpoint(target.endpoint, target.controller, storyDetail)
         if (codes.isEmpty()) {

diff --git a/src/main/kotlin/cc/unitmesh/devti/flow/JavaSpringBaseCrud.kt b/src/main/kotlin/cc/unitmesh/devti/flow/JavaSpringBaseCrud.kt
@@ -35,10 +35,10 @@ class JavaSpringBaseCrud(val project: Project) : SpringBaseCrud {
         return files.map(DtClass.Companion::fromJavaFile)
     }
 
-    private fun getAllControllerFiles(): List<PsiFile> = filterFilesByFunc(::controllerFilter)
-    private fun getAllEntityFiles(): List<PsiFile> = filterFilesByFunc(::entityFilter)
-    private fun getAllDtoFiles(): List<PsiFile> = filterFilesByFunc(::dtoFilter)
-    private fun getAllServiceFiles(): List<PsiFile> = filterFilesByFunc(::serviceFilter)
+    override fun getAllControllerFiles(): List<PsiFile> = filterFilesByFunc(::controllerFilter)
+    override fun getAllEntityFiles(): List<PsiFile> = filterFilesByFunc(::entityFilter)
+    override fun getAllDtoFiles(): List<PsiFile> = filterFilesByFunc(::dtoFilter)
+    override fun getAllServiceFiles(): List<PsiFile> = filterFilesByFunc(::serviceFilter)
 
     private fun filterFilesByFunc(filter: KFunction1<PsiClass, Boolean>): List<PsiFile> {
         return runReadAction {

diff --git a/src/main/kotlin/cc/unitmesh/devti/flow/SpringBaseCrud.kt b/src/main/kotlin/cc/unitmesh/devti/flow/SpringBaseCrud.kt
@@ -2,16 +2,23 @@ package cc.unitmesh.devti.flow
 
 import cc.unitmesh.devti.analysis.DtClass
 import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiFile
 
+// TODO: remove this interface
 interface SpringBaseCrud {
     fun controllerList(): List<DtClass>
+    fun entityList(): List<DtClass>
     fun serviceList(): List<DtClass>
-
     /**
      * return all entity class + dto class
      */
     fun modelList(): List<DtClass>
-    fun entityList(): List<DtClass>
+    fun getAllControllerFiles(): List<PsiFile>
+    fun getAllEntityFiles(): List<PsiFile>
+    fun getAllDtoFiles(): List<PsiFile>
+    fun getAllServiceFiles(): List<PsiFile>
+
+
     fun createControllerOrUpdateMethod(targetController: String, code: String, isControllerExist: Boolean)
     fun createController(endpoint: String, code: String): DtClass?
     fun createEntity(code: String): DtClass?

diff --git a/src/main/kotlin/cc/unitmesh/devti/runconfig/AutoDevRunProfileState.kt b/src/main/kotlin/cc/unitmesh/devti/runconfig/AutoDevRunProfileState.kt
@@ -71,19 +71,26 @@ class AutoDevRunProfileState(
                     val storyId = options.storyId()
                     val storyDetail = autoDevFlow.fillStoryDetail(storyId)
 
-                    indicator.fraction = 0.3
+                    indicator.fraction = 0.2
 
                     indicator.text = AutoDevBundle.message("devti.generatingDtoAndEntity")
                     autoDevFlow.generateDtoAndEntity(storyDetail)
 
+                    indicator.fraction = 0.4
+
                     indicator.text = AutoDevBundle.message("devti.progress.fetchingSuggestEndpoint")
                     val target = autoDevFlow.fetchSuggestEndpoint(storyDetail)
 
-                    indicator.fraction = 0.7
+                    indicator.fraction = 0.6
 
                     indicator.text = AutoDevBundle.message("devti.progress.updatingEndpointMethod")
                     autoDevFlow.updateEndpointMethod(target, storyDetail)
 
+                    indicator.fraction = 0.8
+
+                    indicator.text = AutoDevBundle.message("devti.progress.creatingServiceAndRepository")
+                    autoDevFlow.createServiceAndRepository()
+
                     indicator.fraction = 1.0
                 }
             }

diff --git a/src/main/resources/messages/AutoDevBundle.properties b/src/main/resources/messages/AutoDevBundle.properties
@@ -2,10 +2,12 @@ name=AutoDev
 projectService=Project service: {0}
 configGithubToken=GitHub token:
 configOpenAIToken=OpenAI token:
-devti.progress.creatingStory=Processing story detail (1/4)
-devti.generatingDtoAndEntity=Create DTO and Entity (2/4)
-devti.progress.fetchingSuggestEndpoint=Fetching suggest endpoint (Controller) (3/4)
-devti.progress.updatingEndpointMethod=Updating endpoint (Controller) method (4/4)
+devti.progress.creatingStory=Processing story detail (1/5)
+devti.generatingDtoAndEntity=Create DTO and Entity (2/5)
+devti.progress.fetchingSuggestEndpoint=Fetching suggest endpoint (Controller) (3/5)
+devti.progress.updatingEndpointMethod=Updating endpoint (Controller) method (4/5)
+devti.progress.creatingServiceAndRepository=Create Service and Repository (5/5)
+
 # Config
 autodev.flow=AutoDev Flow
 # Chat

diff --git a/src/main/resources/prompts/openai/create_service_rerpository.txt b/src/main/resources/prompts/openai/create_service_rerpository.txt
@@ -26,4 +26,9 @@ public class {xxxRepository} extends JpaRepository<{xxx}, Long> {
 
 }
 ```
-###
+###
+
+// Controller 代码
+```java
+{controllerCode}
+```