feat: add max completion in files support

unit-mesh · Dec 18, 2023 · 152648e · 152648e
1 parent 9e347b5
commit 152648e
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 3 deletions.
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/SimpleCodePicker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/SimpleCodePicker.kt
@@ -78,7 +78,8 @@ class SimpleCodePicker(private val config: PickerOption) : CodePicker {
             config.codeQualityTypes,
             config.builderConfig,
             pureDataFileName = config.pureDataFileName(),
-            config.completionTypes
+            config.completionTypes,
+            config.maxCompletionInOneFile
         ))
         val walkdirChannel = Channel<FileJob>()
         val summary = mutableListOf<Instruction>()

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/WorkerContext.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/WorkerContext.kt
@@ -11,4 +11,5 @@ data class WorkerContext(
     val builderConfig: BuilderConfig,
     val pureDataFileName: String,
     val completionTypes: List<CompletionBuilderType>,
+    val maxCompletionInOneFile: Int,
 )
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/worker/JavaWorker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/worker/JavaWorker.kt
@@ -83,14 +83,16 @@ class JavaWorker(private val context: WorkerContext) : LangWorker() {
             val jobContext =
                 JobContext(job, context.qualityTypes, fileTree, context.builderConfig, context.completionTypes)
 
-            context.codeContextStrategies.map { type ->
+            val instructions = context.codeContextStrategies.map { type ->
                 val instructionBuilder = type.builder(jobContext)
                 val list = instructionBuilder.build()
                 list.map {
                     file.appendText(it.toString() + "\n")
                 }
                 instructionBuilder.unique(list as List<Nothing>)
             }.flatten()
+
+            Instruction.takeStrategy(instructions, context.maxCompletionInOneFile)
         }.flatten()
 
         return@coroutineScope lists

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/worker/TypescriptWorker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/worker/TypescriptWorker.kt
@@ -43,14 +43,16 @@ class TypescriptWorker(private val context: WorkerContext) : LangWorker() {
         val lists = jobs.map { job ->
             val jobContext = JobContext(job, context.qualityTypes, fileTree, context.builderConfig, context.completionTypes)
 
-            context.codeContextStrategies.map { type ->
+            val instructions = context.codeContextStrategies.map { type ->
                 val instructionBuilder = type.builder(jobContext)
                 val list = instructionBuilder.build()
                 list.map {
                     file.appendText(it.toString() + "\n")
                 }
                 instructionBuilder.unique(list as List<Nothing>)
             }.flatten()
+
+            Instruction.takeStrategy(instructions, context.maxCompletionInOneFile)
         }.flatten()
 
         return@coroutineScope lists