feat(unitest): make it works for instructions

unit-mesh · Dec 23, 2023 · 1d5fd62 · 1d5fd62
1 parent 46e9cad
commit 1d5fd62
Show file tree

Hide file tree

Showing 7 changed files with 35 additions and 13 deletions.
diff --git a/unit-core/src/main/kotlin/cc/unitmesh/core/completion/CompletionBuilder.kt b/unit-core/src/main/kotlin/cc/unitmesh/core/completion/CompletionBuilder.kt
@@ -4,7 +4,7 @@ import chapi.domain.core.CodeDataStruct
 import chapi.domain.core.CodeFunction
 
 interface CompletionBuilder {
-    fun build(dataStruct: CodeDataStruct): List<CodeCompletionIns> {
+    fun build(dataStruct: CodeDataStruct): List<TypedIns> {
         return listOf()
     }
 

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/FullFileCodeCompletionBuilder.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/FullFileCodeCompletionBuilder.kt
@@ -0,0 +1,24 @@
+package cc.unitmesh.pick.builder
+
+import cc.unitmesh.core.completion.CodeCompletionIns
+import cc.unitmesh.core.completion.CompletionBuilder
+import cc.unitmesh.core.completion.TypedIns
+import cc.unitmesh.pick.builder.unittest.lang.UnitTestService
+import cc.unitmesh.pick.worker.job.JobContext
+import chapi.domain.core.CodeDataStruct
+import chapi.domain.core.CodeFunction
+
+class FullFileCodeCompletionBuilder(val context: JobContext) : CompletionBuilder {
+    override fun build(dataStruct: CodeDataStruct): List<TypedIns> {
+        val testIns = UnitTestService.lookup(dataStruct, context).map {
+            it.build(dataStruct)
+        }.flatten()
+
+        return testIns
+    }
+
+    override fun build(function: CodeFunction): List<CodeCompletionIns> {
+        return listOf()
+    }
+
+}
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/InsBuilderUtil.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/InsBuilderUtil.kt
@@ -22,6 +22,7 @@ fun completionBuilder(completionBuilderType: CompletionBuilderType, context: Job
         CompletionBuilderType.INLINE_COMPLETION to InlineCodeCompletionBuilder(context),
         CompletionBuilderType.IN_BLOCK_COMPLETION to InBlockCodeCompletionBuilder(context),
         CompletionBuilderType.AFTER_BLOCK_COMPLETION to AfterBlockCodeCompletionBuilder(context),
+        CompletionBuilderType.FULL_FILE_COMPLETION to FullFileCodeCompletionBuilder(context),
     )[completionBuilderType] ?: throw SerializationException("Unknown message type: $completionBuilderType")
 }
 

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/ClassTestCodeBuilder.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/unittest/ClassTestCodeBuilder.kt
@@ -7,8 +7,6 @@ import cc.unitmesh.core.unittest.TypedTestIns
 import cc.unitmesh.core.unittest.TestCodeBuilder
 import cc.unitmesh.pick.worker.job.JobContext
 import chapi.domain.core.CodeDataStruct
-import kotlinx.serialization.encodeToString
-import kotlinx.serialization.json.Json
 
 class ClassTestCodeBuilder(private val context: JobContext) : TestCodeBuilder {
     /**
@@ -81,13 +79,10 @@ class ClassTestIns(
         input.append(underTestCode)
         input.append("\n```")
 
-        val instruction = Instruction(
+        return Instruction(
             instruction = "Write unit test for following code.",
             input = input.toString(),
             output = generatedCode,
         )
-
-        println("ClassTestIns: ${Json.encodeToString(instruction)}")
-        return instruction
     }
 }
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/option/InsPickerOption.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/option/InsPickerOption.kt
@@ -68,7 +68,8 @@ data class InsPickerOption(
     val completionTypes: List<CompletionBuilderType> = listOf(
         CompletionBuilderType.AFTER_BLOCK_COMPLETION,
         CompletionBuilderType.IN_BLOCK_COMPLETION,
-        CompletionBuilderType.INLINE_COMPLETION
+        CompletionBuilderType.INLINE_COMPLETION,
+        CompletionBuilderType.FULL_FILE_COMPLETION,
     ),
     /**
      * The [CodeQualityType], will be like a tree to hold the item.

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/strategy/BizCodeContextStrategy.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/strategy/BizCodeContextStrategy.kt
@@ -1,7 +1,6 @@
 package cc.unitmesh.pick.strategy
 
 import cc.unitmesh.pick.strategy.base.CodeStrategyBuilder
-import cc.unitmesh.core.completion.TypedIns
 import cc.unitmesh.pick.strategy.bizcode.RelatedCodeStrategyBuilder
 import cc.unitmesh.pick.strategy.bizcode.SimilarChunksStrategyBuilder
 import cc.unitmesh.pick.worker.job.JobContext

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/strategy/bizcode/RelatedCodeStrategyBuilder.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/strategy/bizcode/RelatedCodeStrategyBuilder.kt
@@ -1,5 +1,6 @@
 package cc.unitmesh.pick.strategy.bizcode
 
+import cc.unitmesh.core.completion.TypedIns
 import cc.unitmesh.pick.builder.completionBuilders
 import cc.unitmesh.pick.builder.ins.RelatedCodeIns
 import cc.unitmesh.pick.builder.unittest.lang.UnitTestService
@@ -10,7 +11,7 @@ import chapi.domain.core.CodeDataStruct
 
 class RelatedCodeStrategyBuilder(private val context: JobContext) : CodeStrategyBuilder {
 
-    override fun build(): List<RelatedCodeIns> {
+    override fun build(): List<TypedIns> {
         val language = context.job.fileSummary.language.lowercase()
         val container = context.job.container ?: return emptyList()
         val relatedCode = findRelatedCode(container)
@@ -27,10 +28,10 @@ class RelatedCodeStrategyBuilder(private val context: JobContext) : CodeStrategy
         val builders = completionBuilders(context.completionBuilderTypes, context)
 
         val codeCompletionIns = dataStructs.map { ds ->
-            UnitTestService.lookup(ds, context).map {
+            val blockIns = builders.asSequence().map {
                 it.build(ds)
-            }
-            ds.Functions.map { function ->
+            }.flatten()
+            val functionsIns = ds.Functions.map { function ->
                 builders.asSequence().map {
                     it.build(function)
                 }
@@ -49,6 +50,7 @@ class RelatedCodeStrategyBuilder(private val context: JobContext) : CodeStrategy
                         )
                     }.toList()
             }.flatten()
+            functionsIns + blockIns
         }.flatten()
 
         return codeCompletionIns