feat: impl output bound data #52

unit-mesh · Nov 20, 2023 · 1a2e01e · 1a2e01e
1 parent da6f461
commit 1a2e01e
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 17 deletions.
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/context/JavaContextCollection.kt b/java/src/main/kotlin/cc/unitmesh/idea/context/JavaContextCollection.kt
@@ -41,7 +41,7 @@ object JavaContextCollection {
      *```
      */
     fun dataStructure(clazz: PsiClass): SimpleClassStructure {
-        return createSimpleStructure(clazz)
+        return simpleStructure(clazz)
     }
 
     /**
@@ -54,7 +54,7 @@ object JavaContextCollection {
      * If the field type is a custom class, the method recursively creates a SimpleClassStructure object for that class.
      * If the field type cannot be resolved, it is skipped.
      */
-    private fun createSimpleStructure(clazz: PsiClass): SimpleClassStructure {
+    fun simpleStructure(clazz: PsiClass): SimpleClassStructure {
         val fields = clazz.fields
         val children = fields.mapNotNull { field ->
             when {
@@ -70,7 +70,7 @@ object JavaContextCollection {
 
                 else -> {
                     val classStructure =
-                        (field.type as PsiClassType).resolve()?.let { createSimpleStructure(it) }
+                        (field.type as PsiClassType).resolve()?.let { simpleStructure(it) }
                             ?: return@mapNotNull null
                     classStructure.builtIn = false
                     classStructure

diff --git a/java/src/main/kotlin/cc/unitmesh/idea/prompting/JavaContextPrompter.kt b/java/src/main/kotlin/cc/unitmesh/idea/prompting/JavaContextPrompter.kt
@@ -9,6 +9,7 @@ import cc.unitmesh.devti.provider.context.ChatOrigin
 import cc.unitmesh.devti.settings.AutoDevSettingsState
 import cc.unitmesh.idea.MvcUtil
 import cc.unitmesh.idea.flow.MvcContextService
+import cc.unitmesh.idea.provider.JavaTestDataBuilder
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.project.Project
@@ -132,15 +133,13 @@ open class JavaContextPrompter : ContextPrompter() {
     private fun prepareDataStructure(creationContext: ChatCreationContext) {
         val element = creationContext.element ?: return logger.error("element is null")
 
-        val methodContext = MethodContextProvider(false, false).from(element)
-        selectedText = methodContext.text
-
-        val datastructures = methodContext.inputOutputString()
-        additionContext += """
-// input and output's Class/DataStructures:
-
-$datastructures
-"""
+        val testDataBuilder = JavaTestDataBuilder()
+        testDataBuilder.inBoundData(element).forEach { (key, value) ->
+            additionContext += "input Class\n$value\n"
+        }
+        testDataBuilder.outBoundData(element).forEach { (key, value) ->
+            additionContext += "output Class\n$value\n"
+        }
     }
 
     private fun addFixIssueContext(selectedText: String) {

diff --git a/java/src/main/kotlin/cc/unitmesh/idea/provider/JavaTestDataBuilder.kt b/java/src/main/kotlin/cc/unitmesh/idea/provider/JavaTestDataBuilder.kt
@@ -1,13 +1,26 @@
 package cc.unitmesh.idea.provider
 
 import cc.unitmesh.devti.provider.TestDataBuilder
+import cc.unitmesh.idea.context.JavaContextCollection
+import com.intellij.psi.PsiClassType
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
 
 class JavaTestDataBuilder : TestDataBuilder {
-    override fun inBoundData(element: PsiElement): Pair<String, String>? {
-        if (element !is PsiMethod) return null
+    override fun outBoundData(element: PsiElement): Map<String, String> {
+        if (element !is PsiMethod) return emptyMap()
 
-        return null
+        val returnType = element.returnType ?: return emptyMap()
+        when  {
+            returnType is PsiClassType -> {
+                returnType.resolve()?.let {
+                    val qualifiedName = it.qualifiedName!!
+                    val simpleClassStructure = JavaContextCollection.dataStructure(it)
+                    return mapOf(qualifiedName to simpleClassStructure.toString())
+                }
+            }
+        }
+
+        return emptyMap()
     }
 }
diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/TestDataBuilder.kt b/src/main/kotlin/cc/unitmesh/devti/provider/TestDataBuilder.kt
@@ -5,9 +5,9 @@ import com.intellij.lang.LanguageExtension
 import com.intellij.psi.PsiElement
 
 interface TestDataBuilder {
-    fun inBoundData(element: PsiElement): Pair<String, String>? = null
+    fun inBoundData(element: PsiElement): Map<String, String> = mapOf()
 
-    fun outBoundData(element: PsiElement): Pair<String, String>? = null
+    fun outBoundData(element: PsiElement): Map<String, String> = mapOf()
 
     companion object {
         private val languageExtension: LanguageExtension<TestDataBuilder> =