fix: fix java gen doc return error format issue #99

Author: phodal

java/src/main/kotlin/cc/unitmesh/idea/provider/JavaLivingDocumentation.kt

@@ -10,6 +10,7 @@ import com.intellij.psi.*
 import com.intellij.psi.codeStyle.CodeStyleManager
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.util.IncorrectOperationException
+import org.jetbrains.annotations.NonNls
 
 class JavaLivingDocumentation : LivingDocumentation {
     override val parameterTagInstruction: String get() = "use @param tag"
@@ -38,8 +39,9 @@ class JavaLivingDocumentation : LivingDocumentation {
 
             when (type) {
                 LivingDocumentationType.COMMENT -> {
+                    val doc = Companion.preHandleDoc(newDoc)
                     val psiElementFactory = JavaPsiFacade.getElementFactory(project)
-                    val newDocComment = psiElementFactory.createDocCommentFromText(newDoc)
+                    val newDocComment = psiElementFactory.createDocCommentFromText(doc)
 
                     if (target is PsiDocCommentOwner) {
                         val oldDocComment = target.docComment
@@ -121,4 +123,17 @@ class JavaLivingDocumentation : LivingDocumentation {
         return methodsAndFieldsInRange
     }
 
+    companion object {
+        fun preHandleDoc(newDoc: String): @NonNls String {
+            // 1. remove ```java and ``` from the newDoc if it exists
+            val newDocWithoutCodeBlock = newDoc.removePrefix("```java")
+                .removePrefix("```")
+                .removeSuffix("```")
+
+            // 2. lookup for the first line of the newDoc
+            val fromSuggestion = LivingDocumentation.buildDocFromSuggestion(newDocWithoutCodeBlock, "/**", "*/")
+            return fromSuggestion
+        }
+    }
+
 }

java/src/test/kotlin/cc/unitmesh/idea/provider/JavaLivingDocumentationTest.kt

@@ -0,0 +1,32 @@
+package cc.unitmesh.idea.provider
+
+import com.intellij.testFramework.LightPlatformTestCase
+import org.junit.Test
+
+class JavaLivingDocumentationTest {
+    @Test
+    fun should_handle_when_llm_return_error_java_code() {
+        val code = """```Java
+/**
+    Converts a byte array to a short array.
+    @param bytes the input byte array.
+    @return the converted short array. If the input byte array is null, returns null.
+    */
+    public short[] bytesToShort(byte[] bytes) {
+    if (bytes == null) {
+    return null;
+    }
+    short[] shorts = new short[bytes.length / 2];
+    ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).asShortBuffer().get(shorts);
+    return shorts;
+    }
+"""
+
+        val result = JavaLivingDocumentation.preHandleDoc(code)
+        LightPlatformTestCase.assertEquals(result, """/**
+    Converts a byte array to a short array.
+    @param bytes the input byte array.
+    @return the converted short array. If the input byte array is null, returns null.
+    */""")
+    }
+}

kotlin/src/main/kotlin/cc/unitmesh/kotlin/provider/KotlinLivingDocumentation.kt

@@ -118,7 +118,7 @@ class KotlinLivingDocumentation : LivingDocumentation {
                     try {
                         doInsertComment(target, project, newDoc)
                     } catch (e: Exception) {
-                        val fromSuggestion = buildDocFromSuggestion(newDoc, "/**", "*/")
+                        val fromSuggestion = LivingDocumentation.buildDocFromSuggestion(newDoc, "/**", "*/")
                         if (fromSuggestion.isNotEmpty()) {
                             try {
                                 doInsertComment(target, project, fromSuggestion)

src/main/kotlin/cc/unitmesh/devti/provider/LivingDocumentation.kt

@@ -32,22 +32,6 @@ interface LivingDocumentation {
         return selectionModel.selectionStart <= element.textRange.startOffset && element.textRange.endOffset <= selectionModel.selectionEnd
     }
 
-    fun buildDocFromSuggestion(suggestion: String, commentStart: String, commentEnd: String): String {
-        val startIndex = suggestion.indexOf(commentStart)
-        if (startIndex < 0) {
-            return ""
-        }
-
-        val docComment = suggestion.substring(startIndex)
-        val endIndex = docComment.indexOf(commentEnd, commentStart.length)
-        if (endIndex < 0) {
-            return docComment + commentEnd
-        }
-
-        val substring = docComment.substring(0, endIndex + commentEnd.length)
-        return substring
-    }
-
     companion object {
         private val languageExtension: LanguageExtension<LivingDocumentation> =
             LanguageExtension("cc.unitmesh.livingDocumentation")
@@ -64,5 +48,21 @@ interface LivingDocumentation {
 
             return null
         }
+
+        fun buildDocFromSuggestion(suggestion: String, commentStart: String, commentEnd: String): String {
+            val startIndex = suggestion.indexOf(commentStart)
+            if (startIndex < 0) {
+                return ""
+            }
+
+            val docComment = suggestion.substring(startIndex)
+            val endIndex = docComment.indexOf(commentEnd, commentStart.length)
+            if (endIndex < 0) {
+                return docComment + commentEnd
+            }
+
+            val substring = docComment.substring(0, endIndex + commentEnd.length)
+            return substring
+        }
     }
 }