feat(doc): add basdic prompt

phodal · phodal · commit 4fddd6d28777 · 2023-08-09T11:31:57.000+08:00
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/promting/JavaContextPrompter.kt b/java/src/main/kotlin/cc/unitmesh/idea/promting/JavaContextPrompter.kt
@@ -93,25 +93,12 @@ open class JavaContextPrompter : ContextPrompter() {
 
         when (action!!) {
             ChatActionType.EXPLAIN -> {
-                val autoComment = customPromptConfig?.autoComment
-                if (autoComment?.instruction?.isNotEmpty() == true) {
-                    prompt = autoComment.instruction
-                }
             }
 
             ChatActionType.REFACTOR -> {
-                val refactor = customPromptConfig?.refactor
-                if (refactor?.instruction?.isNotEmpty() == true) {
-                    prompt = refactor.instruction
-                }
             }
 
             ChatActionType.CODE_COMPLETE -> {
-                val codeComplete = customPromptConfig?.autoComplete
-                if (codeComplete?.instruction?.isNotEmpty() == true) {
-                    prompt = codeComplete.instruction
-                }
-
                 when {
                     MvcUtil.isController(fileName, lang) -> {
                         val spec = CustomPromptConfig.load().spec["controller"]
@@ -135,12 +122,7 @@ open class JavaContextPrompter : ContextPrompter() {
                 }
             }
 
-            ChatActionType.GENERATE_TEST -> {
-                val writeTest = customPromptConfig?.writeTest
-                if (writeTest?.instruction?.isNotEmpty() == true) {
-                    prompt = writeTest.instruction
-                }
-            }
+            ChatActionType.GENERATE_TEST -> {}
 
             ChatActionType.FIX_ISSUE -> {
                 addFixIssueContext(selectedText)
@@ -158,6 +140,7 @@ open class JavaContextPrompter : ContextPrompter() {
 
             ChatActionType.CUSTOM_COMPLETE -> {
             }
+
             ChatActionType.CUSTOM_ACTION -> TODO()
         }
 
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/provider/JavaLivingDocumentation.kt b/java/src/main/kotlin/cc/unitmesh/idea/provider/JavaLivingDocumentation.kt
@@ -2,12 +2,11 @@ package cc.unitmesh.idea.provider
 
 import cc.unitmesh.devti.provider.LivingDocumentation
 import cc.unitmesh.devti.provider.LivingDocumentationType
+import com.intellij.codeInsight.daemon.impl.CollectHighlightsUtil
 import com.intellij.openapi.editor.SelectionModel
-import com.intellij.psi.PsiClass
-import com.intellij.psi.PsiElement
-import com.intellij.psi.PsiMethod
-import com.intellij.psi.PsiNameIdentifierOwner
+import com.intellij.psi.*
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.IncorrectOperationException
 
 class JavaLivingDocumentation : LivingDocumentation {
     override fun startEndString(type: LivingDocumentationType): Pair<String, String> {
@@ -19,9 +18,23 @@ class JavaLivingDocumentation : LivingDocumentation {
     }
 
     override fun updateDoc(psiElement: PsiElement, str: String) {
-        TODO("Not yet implemented")
+        val project = psiElement.project
+        val psiElementFactory = JavaPsiFacade.getElementFactory(project)
+        val newDocComment = psiElementFactory.createDocCommentFromText(str)
+
+        if (psiElement is PsiDocCommentOwner) {
+            try {
+                psiElement.docComment?.replace(newDocComment)
+            } catch (e: IncorrectOperationException) {
+                val firstChild = psiElement.firstChild
+                psiElement.addBefore(newDocComment, firstChild)
+            }
+        } else {
+            throw IncorrectOperationException("Unable to update documentation")
+        }
     }
 
+
     override fun findExampleDoc(psiNameIdentifierOwner: PsiNameIdentifierOwner): String {
         return ""
     }
@@ -38,11 +51,51 @@ class JavaLivingDocumentation : LivingDocumentation {
 
     }
 
+    fun containsElement(selectionModel: SelectionModel, element: PsiElement): Boolean {
+        return selectionModel.selectionStart <= element.textRange.startOffset && element.textRange.endOffset <= selectionModel.selectionEnd
+    }
+
     override fun findDocTargetsInSelection(
         psiElement: PsiElement,
         selectionModel: SelectionModel,
     ): List<PsiNameIdentifierOwner> {
-        TODO("Not yet implemented")
+        val findCommonParent = CollectHighlightsUtil.findCommonParent(
+            psiElement,
+            selectionModel.selectionStart,
+            selectionModel.selectionEnd
+        ) ?: return emptyList()
+
+        if (findCommonParent is PsiJavaFile) {
+            val classAndFieldMethods = mutableListOf<PsiNameIdentifierOwner>()
+            val classes = findCommonParent.classes
+            for (psiClass in classes) {
+                if (containsElement(selectionModel, psiClass)) {
+                    classAndFieldMethods.add(psiClass)
+                }
+            }
+
+            return classAndFieldMethods
+        }
+
+        val target = findNearestDocumentationTarget(findCommonParent) ?: return emptyList()
+
+        if (target !is PsiClass || containsElement(selectionModel, target)) {
+            return listOf(target)
+        }
+
+        val methodsAndFieldsInRange = mutableListOf<PsiNameIdentifierOwner>()
+        for (psiField in target.fields) {
+            if (containsElement(selectionModel, psiField)) {
+                methodsAndFieldsInRange.add(psiField)
+            }
+        }
+        for (psiMethod in target.methods) {
+            if (containsElement(selectionModel, psiMethod)) {
+                methodsAndFieldsInRange.add(psiMethod)
+            }
+        }
+
+        return methodsAndFieldsInRange
     }
 
 }
diff --git a/src/main/kotlin/cc/unitmesh/devti/context/ClassContextProvider.kt b/src/main/kotlin/cc/unitmesh/devti/context/ClassContextProvider.kt
@@ -1,12 +1,13 @@
 package cc.unitmesh.devti.context
 
+import cc.unitmesh.devti.context.base.LLMQueryContextProvider
 import cc.unitmesh.devti.context.builder.ClassContextBuilder
 import com.intellij.lang.Language
 import com.intellij.lang.LanguageExtension
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.psi.PsiElement
 
-class ClassContextProvider(private val gatherUsages: Boolean) {
+class ClassContextProvider(private val gatherUsages: Boolean) : LLMQueryContextProvider<PsiElement> {
     private val languageExtension = LanguageExtension<ClassContextBuilder>("cc.unitmesh.classContextBuilder")
     private val providers: List<ClassContextBuilder>
 
@@ -19,7 +20,7 @@ class ClassContextProvider(private val gatherUsages: Boolean) {
         val logger = logger<ClassContextProvider>()
     }
 
-    fun from(psiElement: PsiElement): ClassContext {
+    override fun from(psiElement: PsiElement): ClassContext {
         for (provider in providers) {
             provider.getClassContext(psiElement, gatherUsages)?.let {
                 return it
diff --git a/src/main/kotlin/cc/unitmesh/devti/context/FileContextProvider.kt b/src/main/kotlin/cc/unitmesh/devti/context/FileContextProvider.kt
@@ -1,11 +1,12 @@
 package cc.unitmesh.devti.context
 
+import cc.unitmesh.devti.context.base.LLMQueryContextProvider
 import cc.unitmesh.devti.context.builder.FileContextBuilder
 import com.intellij.lang.Language
 import com.intellij.lang.LanguageExtension
 import com.intellij.psi.PsiFile
 
-class FileContextProvider {
+class FileContextProvider: LLMQueryContextProvider<PsiFile> {
     private val languageExtension: LanguageExtension<FileContextBuilder> =
         LanguageExtension("cc.unitmesh.fileContextBuilder")
 
@@ -16,14 +17,14 @@ class FileContextProvider {
         providers = registeredLanguages.mapNotNull { languageExtension.forLanguage(it) }
     }
 
-    fun from(psiFile: PsiFile): FileContext {
+    override fun from(psiElement: PsiFile): FileContext {
         for (provider in providers) {
-            val fileContext = provider.getFileContext(psiFile)
+            val fileContext = provider.getFileContext(psiElement)
             if (fileContext != null) {
                 return fileContext
             }
         }
 
-        return FileContext(psiFile, psiFile.name, psiFile.virtualFile?.path!!)
+        return FileContext(psiElement, psiElement.name, psiElement.virtualFile?.path!!)
     }
 }
diff --git a/src/main/kotlin/cc/unitmesh/devti/context/MethodContextProvider.kt b/src/main/kotlin/cc/unitmesh/devti/context/MethodContextProvider.kt
@@ -1,12 +1,14 @@
 package cc.unitmesh.devti.context
 
+import cc.unitmesh.devti.context.base.LLMQueryContextProvider
 import cc.unitmesh.devti.context.builder.MethodContextBuilder
 import com.intellij.lang.Language
 import com.intellij.lang.LanguageExtension
 import com.intellij.psi.PsiElement
 import org.jetbrains.annotations.NotNull
 
-class MethodContextProvider(private val includeClassContext: Boolean, private val gatherUsages: Boolean) {
+class MethodContextProvider(private val includeClassContext: Boolean, private val gatherUsages: Boolean):
+    LLMQueryContextProvider<PsiElement> {
     @NotNull
     private val languageExtension: LanguageExtension<MethodContextBuilder> =
         LanguageExtension("cc.unitmesh.methodContextBuilder")
@@ -20,7 +22,7 @@ class MethodContextProvider(private val includeClassContext: Boolean, private va
     }
 
     @NotNull
-    fun from(@NotNull psiElement: PsiElement): MethodContext {
+    override fun from(@NotNull psiElement: PsiElement): MethodContext {
         val iterator = providers.iterator()
         while (iterator.hasNext()) {
             val provider = iterator.next()
diff --git a/src/main/kotlin/cc/unitmesh/devti/context/VariableContextProvider.kt b/src/main/kotlin/cc/unitmesh/devti/context/VariableContextProvider.kt
@@ -1,5 +1,6 @@
 package cc.unitmesh.devti.context
 
+import cc.unitmesh.devti.context.base.LLMQueryContextProvider
 import cc.unitmesh.devti.context.builder.VariableContextBuilder
 import com.intellij.lang.Language
 import com.intellij.lang.LanguageExtension
@@ -9,7 +10,7 @@ class VariableContextProvider(
     private val includeMethodContext: Boolean,
     private val includeClassContext: Boolean,
     private val gatherUsages: Boolean
-) {
+): LLMQueryContextProvider<PsiElement> {
     private val languageExtension: LanguageExtension<VariableContextBuilder> =
         LanguageExtension("cc.unitmesh.variableContextBuilder")
 
@@ -20,7 +21,7 @@ class VariableContextProvider(
         providers = registeredLanguages.mapNotNull(languageExtension::forLanguage)
     }
 
-    fun from(psiElement: PsiElement): VariableContext {
+    override fun from(psiElement: PsiElement): VariableContext {
         for (provider in providers) {
             val variableContext =
                 provider.getVariableContext(psiElement, includeMethodContext, includeClassContext, gatherUsages)
diff --git a/src/main/kotlin/cc/unitmesh/devti/context/base/LLMQueryContextProvider.kt b/src/main/kotlin/cc/unitmesh/devti/context/base/LLMQueryContextProvider.kt
@@ -0,0 +1,8 @@
+package cc.unitmesh.devti.context.base;
+
+import com.intellij.psi.PsiElement
+
+interface LLMQueryContextProvider<T : PsiElement?> {
+    fun from(psiElement: T): LLMQueryContext
+}
+
diff --git a/src/main/kotlin/cc/unitmesh/devti/intentions/action/LivingDocumentationIntention.kt b/src/main/kotlin/cc/unitmesh/devti/intentions/action/LivingDocumentationIntention.kt
@@ -10,6 +10,7 @@ import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.progress.Task
 import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
 import com.intellij.psi.PsiNameIdentifierOwner
@@ -39,13 +40,30 @@ class LivingDocumentationIntention : AbstractChatIntention() {
         if (selectedText != null) {
             val owners: List<PsiNameIdentifierOwner> = findSelectedElementToDocument(editor, project, selectionModel)
             for (identifierOwner in owners) {
-                val task: Task.Backgroundable = LivingDocumentationTask(editor, identifierOwner)
-
-                ProgressManager.getInstance()
-                    .runProcessWithProgressAsynchronously(task, BackgroundableProcessIndicator(task))
+                writeForDocument(editor, identifierOwner)
             }
         }
 
+        val closestToCaretNamedElement: PsiNameIdentifierOwner? = getClosestToCaretNamedElement(editor)
+        if (closestToCaretNamedElement != null) {
+            writeForDocument(editor, closestToCaretNamedElement)
+        }
+    }
+
+    private fun writeForDocument(editor: Editor, element: PsiNameIdentifierOwner) {
+        val task: Task.Backgroundable = LivingDocumentationTask(editor, element)
+        ProgressManager.getInstance()
+            .runProcessWithProgressAsynchronously(task, BackgroundableProcessIndicator(task))
+    }
+
+    private fun getClosestToCaretNamedElement(editor: Editor): PsiNameIdentifierOwner? {
+        val element = PsiUtilBase.getElementAtCaret(editor) ?: return null
+        return getClosestNamedElement(element)
+    }
+
+    private fun getClosestNamedElement(element: PsiElement): PsiNameIdentifierOwner? {
+        val support = LivingDocumentation.forLanguage(element.language) ?: return null
+        return support.findNearestDocumentationTarget(element)
     }
 
     private fun findSelectedElementToDocument(
diff --git a/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/LivingDocumentationTask.kt b/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/LivingDocumentationTask.kt

Original file line number	Diff line number	Diff line change
`@@ -1,11 +1,12 @@`
`1`	`1`	`package cc.unitmesh.devti.context`
`2`	`2`
	`3`	`+import cc.unitmesh.devti.context.base.LLMQueryContextProvider`
`3`	`4`	`import cc.unitmesh.devti.context.builder.FileContextBuilder`
`4`	`5`	`import com.intellij.lang.Language`
`5`	`6`	`import com.intellij.lang.LanguageExtension`
`6`	`7`	`import com.intellij.psi.PsiFile`
`7`	`8`
`8`		`-class FileContextProvider {`
	`9`	`+class FileContextProvider: LLMQueryContextProvider<PsiFile> {`
`9`	`10`	`private val languageExtension: LanguageExtension<FileContextBuilder> =`
`10`	`11`	`LanguageExtension("cc.unitmesh.fileContextBuilder")`
`11`	`12`
`@@ -16,14 +17,14 @@ class FileContextProvider {`
`16`	`17`	`providers = registeredLanguages.mapNotNull { languageExtension.forLanguage(it) }`
`17`	`18`	`}`
`18`	`19`
`19`		`- fun from(psiFile: PsiFile): FileContext {`
	`20`	`+ override fun from(psiElement: PsiFile): FileContext {`
`20`	`21`	`for (provider in providers) {`
`21`		`- val fileContext = provider.getFileContext(psiFile)`
	`22`	`+ val fileContext = provider.getFileContext(psiElement)`
`22`	`23`	`if (fileContext != null) {`
`23`	`24`	`return fileContext`
`24`	`25`	`}`
`25`	`26`	`}`
`26`	`27`
`27`		`- return FileContext(psiFile, psiFile.name, psiFile.virtualFile?.path!!)`
	`28`	`+ return FileContext(psiElement, psiElement.name, psiElement.virtualFile?.path!!)`
`28`	`29`	`}`
`29`	`30`	`}`