From 0d0d121c6dfc5204fa044ae53d1111f4a82c016d Mon Sep 17 00:00:00 2001
From: Phodal Huang <h@phodal.com>
Date: Thu, 21 Dec 2023 15:17:23 +0800
Subject: [PATCH] fix: fix simliar chunk length issue

---
 .../kotlin/cc/unitmesh/pick/builder/PickerOption.kt   |  3 ++-
 .../cc/unitmesh/pick/similar/JavaSimilarChunker.kt    | 11 ++++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/PickerOption.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/PickerOption.kt
index a6f7240f..77af397d 100644
--- a/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/PickerOption.kt
+++ b/unit-picker/src/main/kotlin/cc/unitmesh/pick/builder/PickerOption.kt
@@ -49,7 +49,8 @@ data class PickerOption(
      *
      */
     val codeContextStrategies: List<CodeContextStrategy> = listOf(
-        CodeContextStrategy.RELATED_CODE
+//        CodeContextStrategy.RELATED_CODE,
+        CodeContextStrategy.SIMILAR_CHUNKS,
     ),
     /**
      * The [CompletionBuilderType], which will according you IDE strategy to generate the type.
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/similar/JavaSimilarChunker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/similar/JavaSimilarChunker.kt
index e9a42ca3..4f434ecd 100644
--- a/unit-picker/src/main/kotlin/cc/unitmesh/pick/similar/JavaSimilarChunker.kt
+++ b/unit-picker/src/main/kotlin/cc/unitmesh/pick/similar/JavaSimilarChunker.kt
@@ -36,25 +36,26 @@ class JavaSimilarChunker(private val fileTree: HashMap<String, InstructionFileJo
             .take(maxRelevantFiles)
             .map { it.second }
 
-        logger.info("canonicalName: $canonicalName, relatedCodePath: $relatedCodePath")
+//        logger.info("canonicalName: $canonicalName, relatedCodePath: $relatedCodePath")
 
         val chunks = chunkedCode(beforeCursor)
         val allRelatedChunks = relatedCodePath
             .mapNotNull { fileTree[it] }
             .map { chunkedCode(it.code).joinToString("\n") }
 
-        val similarChunks = allRelatedChunks.map {
+        val similarChunks: List<Pair<Double, String>> = allRelatedChunks.map {
             val score = similarityScore(tokenize(it).toSet(), chunks.toSet())
             score to it
         }.sortedByDescending { it.first }
+            .filter { it.first > 0.0 }
             .take(maxRelevantFiles)
-            .map { it.second }
 
         // take the first 3 similar chunks or empty
         val similarChunksText = if (similarChunks.size > 3) {
-            similarChunks.take(3)
+            println("score: ${similarChunks.map { it.first }}")
+            similarChunks.take(3).map { it.second }
         } else {
-            similarChunks
+            similarChunks.map { it.second }
         }.map {
             // clean up the similar chunks, if start with multiple \n, remove to one
             it.replace(Regex("^\\n+"), "\n")