feat: add worker context to text builder

unit-mesh · Dec 7, 2023 · ac6eef7 · ac6eef7
1 parent 9a542e9
commit ac6eef7
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 12 deletions.
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/picker/CodePicker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/picker/CodePicker.kt
@@ -1,6 +1,8 @@
 package cc.unitmesh.pick.picker
 
+import cc.unitmesh.pick.prompt.Instruction
 import cc.unitmesh.pick.prompt.InstructionBuilder
+import cc.unitmesh.pick.worker.WorkerContext
 import cc.unitmesh.pick.worker.WorkerManager
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.Channel
@@ -25,9 +27,9 @@ class CodePicker(private val config: PickerConfig) {
             logger.info("start picker")
 
             val languageWorker = LanguageWorker()
-            val workerManager = WorkerManager()
+            val workerManager = WorkerManager(WorkerContext(config.builderTypes))
             val walkdirChannel = Channel<FileJob>()
-            val summary = mutableListOf<InstructionBuilder>()
+            val summary = mutableListOf<Instruction>()
 
             launch {
                 launch {
@@ -43,11 +45,11 @@ class CodePicker(private val config: PickerConfig) {
 
                     summary.addAll(workerManager.runAll())
                 }
+
+                // todo: add summary to jsonl
             }.join()
 
             logger.info("stop picker")
-
-
             // 3. generate tree to jsonl
         }
     }

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/picker/PickerConfig.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/picker/PickerConfig.kt
@@ -1,8 +1,13 @@
 package cc.unitmesh.pick.picker
 
+import cc.unitmesh.pick.prompt.InstructionType
+
 data class PickerConfig(
     val url: String,
     val branch: String = "master",
-    val baseDir: String = "."
+    val baseDir: String = ".",
+    val builderTypes: List<InstructionType> = listOf(
+        InstructionType.RELATED_CODE_COMPLETION
+    ),
 ) {
 }
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/LangWorker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/LangWorker.kt
@@ -2,8 +2,9 @@ package cc.unitmesh.pick.worker
 
 import cc.unitmesh.pick.prompt.InstructionBuilder
 import cc.unitmesh.pick.picker.InstructionJob
+import cc.unitmesh.pick.prompt.Instruction
 
 abstract class LangWorker {
     abstract fun addJob(job: InstructionJob)
-    abstract suspend fun start() : List<InstructionBuilder>
+    abstract suspend fun start() : Collection<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
@@ -0,0 +1,7 @@
+package cc.unitmesh.pick.worker
+
+import cc.unitmesh.pick.prompt.InstructionType
+
+data class WorkerContext(
+    val builderType: List<InstructionType>,
+)
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/WorkerManager.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/WorkerManager.kt
@@ -2,11 +2,13 @@ package cc.unitmesh.pick.worker
 
 import cc.unitmesh.pick.prompt.InstructionBuilder
 import cc.unitmesh.pick.picker.InstructionJob
+import cc.unitmesh.pick.prompt.Instruction
+import cc.unitmesh.pick.worker.worker.JavaLangWorker
 import org.archguard.rule.common.Language
 
-class WorkerManager {
+class WorkerManager(workerContext: WorkerContext) {
     private val workers: Map<Language, LangWorker> = mapOf(
-        Language.JAVA to JavaLangWorker(),
+        Language.JAVA to JavaLangWorker(workerContext),
     )
 
     fun addJob(job: InstructionJob) {
@@ -15,7 +17,7 @@ class WorkerManager {
         worker?.addJob(job)
     }
 
-    suspend fun runAll() : List<InstructionBuilder> {
+    suspend fun runAll() : List<Instruction> {
         return workers.map { (_, worker) ->
             worker.start()
         }.flatten()

diff --git a/...cc/unitmesh/pick/worker/JavaLangWorker.kt → ...mesh/pick/worker/worker/JavaLangWorker.kt b/...cc/unitmesh/pick/worker/JavaLangWorker.kt → ...mesh/pick/worker/worker/JavaLangWorker.kt
@@ -1,7 +1,10 @@
-package cc.unitmesh.pick.worker
+package cc.unitmesh.pick.worker.worker
 
 import cc.unitmesh.pick.prompt.InstructionBuilder
 import cc.unitmesh.pick.picker.InstructionJob
+import cc.unitmesh.pick.prompt.Instruction
+import cc.unitmesh.pick.worker.LangWorker
+import cc.unitmesh.pick.worker.WorkerContext
 import chapi.ast.javaast.JavaAnalyser
 import kotlinx.coroutines.coroutineScope
 import org.archguard.scanner.analyser.count.FileJob
@@ -20,7 +23,7 @@ import org.archguard.scanner.analyser.count.FileJob
  * - by Horizontal (with Import File):
  * - by Vertical (with History Change):
  */
-class JavaLangWorker : LangWorker() {
+class JavaLangWorker(val workerContext: WorkerContext) : LangWorker() {
     private val jobs: MutableList<InstructionJob> = mutableListOf()
     private val packageTree: MutableMap<String, InstructionJob> = mutableMapOf()
 
@@ -44,7 +47,7 @@ class JavaLangWorker : LangWorker() {
         }
     }
 
-    override suspend fun start(): List<InstructionBuilder> = coroutineScope {
+    override suspend fun start(): Collection<Instruction> = coroutineScope {
         jobs.map {
 
         }