feat: add simple line count strategy

unit-mesh · Dec 21, 2023 · 744847e · 744847e
1 parent 5e48ab4
commit 744847e
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 3 deletions.
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/strategy/RelatedCodeCompletionIns.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/strategy/RelatedCodeCompletionIns.kt
@@ -2,6 +2,7 @@ package cc.unitmesh.pick.prompt.strategy
 
 import cc.unitmesh.pick.prompt.CompletionBuilderType
 import cc.unitmesh.pick.prompt.Instruction
+import cc.unitmesh.pick.threshold.QualityThreshold
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.Json
 
@@ -14,17 +15,38 @@ data class RelatedCodeCompletionIns(
     val output: String,
     override val type: CompletionBuilderType,
 ) : TypedCompletionIns {
-
     override fun toString(): String {
         return Json.encodeToString(serializer(), this)
     }
 
     override fun unique(): Instruction {
+        // Related code strategy
         val relatedCode = if (relatedCode.isNotBlank() && relatedCode.isNotEmpty()) {
-            "\n// Compare this snippets: \n ```${language}\n$relatedCode\n```"
+            // limit relatedCode to 30 lines
+            val relatedCodeLines = relatedCode.lines()
+            val maxLine = QualityThreshold.MAX_RELATED_CODE_LINE
+            if (relatedCodeLines.size > maxLine) {
+                relatedCodeLines.take(maxLine).joinToString("\n")
+            } else {
+                relatedCode
+            }
+            "\n// Compare this snippets: \n ```${language}\n${relatedCodeLines.joinToString("\n")}\n```"
         } else {
             ""
         }
+
+        // Count strategy
+        val maxLine = QualityThreshold.MAX_LINE_IN_CODE
+        val beforeCursorLine = beforeCursor.count { it == '\n' }
+        val afterCursorLine = output.count { it == '\n' }
+        // drop from the start of beforeCursor
+        val beforeCursor = if (beforeCursorLine + afterCursorLine > maxLine) {
+            val dropLine = beforeCursorLine + afterCursorLine - maxLine
+            beforeCursor.lines().drop(dropLine).joinToString("\n")
+        } else {
+            beforeCursor
+        }
+
         val input = "$relatedCode\n\nCode:\n```${language}\n$beforeCursor\n```"
         return Instruction(
             instruction = "Complete $language code, return rest code, no explaining",

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/strategy/SimilarChunkCompletionIns.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/strategy/SimilarChunkCompletionIns.kt
@@ -2,6 +2,7 @@ package cc.unitmesh.pick.prompt.strategy
 
 import cc.unitmesh.pick.prompt.CompletionBuilderType
 import cc.unitmesh.pick.prompt.Instruction
+import cc.unitmesh.pick.threshold.QualityThreshold
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.Json
 
@@ -19,7 +20,34 @@ data class SimilarChunkCompletionIns(
     }
 
     override fun unique(): Instruction {
-        val input = "\n${similarChunks}              \nCode:\n```${language}\n${beforeCursor}\n```"
+        // Similar chunk strategy
+        val similarChunks = if (similarChunks.isNotBlank() && similarChunks.isNotEmpty()) {
+            // limit similarChunks to 30 lines
+            val similarChunksLines = similarChunks.lines()
+            val maxLine = QualityThreshold.MAX_RELATED_CODE_LINE
+            if (similarChunksLines.size > maxLine) {
+                similarChunksLines.take(maxLine).joinToString("\n")
+            } else {
+                similarChunks
+            }
+            "\n// Similar chunk:\n ```${language}\n${similarChunksLines.joinToString("\n")}\n```"
+        } else {
+            ""
+        }
+
+        // Count strategy
+        val maxLine = QualityThreshold.MAX_LINE_IN_CODE
+        val beforeCursorLine = beforeCursor.count { it == '\n' }
+        val afterCursorLine = output.count { it == '\n' }
+        // drop from the start of beforeCursor
+        val beforeCursor = if (beforeCursorLine + afterCursorLine > maxLine) {
+            val dropLine = beforeCursorLine + afterCursorLine - maxLine
+            beforeCursor.lines().drop(dropLine).joinToString("\n")
+        } else {
+            beforeCursor
+        }
+
+        val input = "$similarChunks\n\nCode:\n```${language}\n$beforeCursor\n```"
 
         return Instruction(
             instruction = "Complete $language code, return rest code, no explaining",

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/strategy/TypedCompletionIns.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/strategy/TypedCompletionIns.kt
@@ -5,5 +5,6 @@ import cc.unitmesh.pick.prompt.Instruction
 
 interface TypedCompletionIns {
     val type: CompletionBuilderType
+
     fun unique(): Instruction
 }
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/threshold/QualityThreshold.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/threshold/QualityThreshold.kt
@@ -26,5 +26,6 @@ data class QualityThreshold(
         const val MAX_FILE_SIZE: Long = 1024 * 64
         const val MAX_LINE_IN_CODE: Int = 160
         const val MAX_CHAR_IN_CODE: Int = 1500
+        const val MAX_RELATED_CODE_LINE: Int = 30
     }
 }