fix: fix lost element issue

unit-mesh · Aug 7, 2023 · 3dd043e · 3dd043e
1 parent d2e7bfe
commit 3dd043e
Show file tree

Hide file tree

Showing 9 changed files with 49 additions and 10 deletions.
diff --git a/clion/src/main/kotlin/cc/unitmesh/clion/provider/RustContextPrompter.kt b/clion/src/main/kotlin/cc/unitmesh/clion/provider/RustContextPrompter.kt
@@ -27,7 +27,7 @@ class RustContextPrompter : ContextPrompter() {
         element: PsiElement?
     ) {
         super.initContext(actionType, selectedText, file, project, offset, element)
-        creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file)
+        creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file, listOf(), element)
     }
 
     override fun displayPrompt(): String {

diff --git a/pycharm/src/main/kotlin/cc/unitmesh/pycharm/PythonFrameworkContextProvider.kt b/pycharm/src/main/kotlin/cc/unitmesh/pycharm/PythonFrameworkContextProvider.kt
@@ -0,0 +1,24 @@
+package cc.unitmesh.pycharm
+
+import cc.unitmesh.devti.provider.context.ChatContextItem
+import cc.unitmesh.devti.provider.context.ChatContextProvider
+import cc.unitmesh.devti.provider.context.ChatCreationContext
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.project.modules
+import com.jetbrains.python.PythonLanguage
+import com.jetbrains.python.sdk.configuration.PyProjectSdkConfigurationExtension
+
+class PythonFrameworkContextProvider: ChatContextProvider {
+    override fun isApplicable(project: Project, creationContext: ChatCreationContext): Boolean {
+        println(creationContext.element?.language)
+        return creationContext.element?.language is PythonLanguage
+    }
+
+    override suspend fun collect(project: Project, creationContext: ChatCreationContext): List<ChatContextItem> {
+        project.modules.asSequence()
+            .map { PyProjectSdkConfigurationExtension.findForModule(it) }
+            .forEach { println(it) }
+
+        return listOf()
+    }
+}
diff --git a/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonContextPrompter.kt b/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonContextPrompter.kt
@@ -3,12 +3,16 @@ package cc.unitmesh.pycharm.provider
 import com.intellij.temporary.similar.chunks.SimilarChunksWithPaths
 import cc.unitmesh.devti.gui.chat.ChatActionType
 import cc.unitmesh.devti.provider.ContextPrompter
+import cc.unitmesh.devti.provider.context.ChatCreationContext
+import cc.unitmesh.devti.provider.context.ChatOrigin
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import kotlinx.coroutines.runBlocking
 
 class PythonContextPrompter : ContextPrompter() {
+    private lateinit var creationContext: ChatCreationContext
     private var additionContext: String = ""
 
     companion object {
@@ -21,19 +25,26 @@ class PythonContextPrompter : ContextPrompter() {
         file: PsiFile?,
         project: Project,
         offset: Int,
-        element: PsiElement?
+        element: PsiElement?,
     ) {
         super.initContext(actionType, selectedText, file, project, offset, element)
         additionContext = SimilarChunksWithPaths.createQuery(file!!) ?: ""
+        creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file, listOf(), element)
     }
 
     override fun displayPrompt(): String {
-        return "$action\n```${lang}\n$selectedText\n```"
+        return runBlocking {
+            additionContext = collectionContext(creationContext)
+            return@runBlocking "$action\n```${lang}\n$selectedText\n```"
+        }
     }
 
     override fun requestPrompt(): String {
-        val prompt = "$action\n```${lang}\n$additionContext\n$selectedText\n```"
-        log.info("final prompt: $prompt")
-        return prompt
+        return runBlocking {
+            additionContext = collectionContext(creationContext)
+            val prompt = "$action\n```${lang}\n$additionContext\n$selectedText\n```"
+            log.info("final prompt: $prompt")
+            return@runBlocking prompt
+        }
     }
 }
diff --git a/pycharm/src/main/resources/cc.unitmesh.pycharm.xml b/pycharm/src/main/resources/cc.unitmesh.pycharm.xml
@@ -14,6 +14,8 @@
         <variableContextBuilder language="Python"
                                 implementationClass="cc.unitmesh.pycharm.context.PythonVariableContextBuilder"/>
 
+        <chatContextProvider implementation="cc.unitmesh.pycharm.PythonFrameworkContextProvider"/>
+
         <contextPrompter
                 language="Python"
                 implementation="cc.unitmesh.pycharm.provider.PythonContextPrompter"/>

diff --git a/src/main/kotlin/cc/unitmesh/devti/actions/chat/GenerateCodeInplaceAction.kt b/src/main/kotlin/cc/unitmesh/devti/actions/chat/GenerateCodeInplaceAction.kt
@@ -64,7 +64,7 @@ class GenerateCodeInplaceAction : AnAction() {
     }
 
     fun getChatCreationContext(list: List<ChatContextItem>, sourceFile: PsiFile?): ChatCreationContext {
-        return ChatCreationContext(ChatOrigin.Unknown, ChatActionType.CUSTOM_COMPLETE, sourceFile, list)
+        return ChatCreationContext(ChatOrigin.Unknown, ChatActionType.CUSTOM_COMPLETE, sourceFile, list, null)
     }
 
     private fun setupPrompt(inlay: Inlay<InlayComponent<AIInplacePrompt>>, sourceFile: PsiFile?) {

diff --git a/src/main/kotlin/cc/unitmesh/devti/intentions/task/TestCodeGenTask.kt b/src/main/kotlin/cc/unitmesh/devti/intentions/task/TestCodeGenTask.kt
@@ -63,7 +63,7 @@ class TestCodeGenTask(val request: TestCodeGenRequest) :
         indicator.text = AutoDevBundle.message("intentions.chat.code.test.step.collect-context")
         indicator.fraction = 0.3
 
-        val creationContext = ChatCreationContext(ChatOrigin.Intention, actionType, request.file)
+        val creationContext = ChatCreationContext(ChatOrigin.Intention, actionType, request.file, listOf(), element = request.element)
         val contextItems: List<ChatContextItem> = runBlocking {
             return@runBlocking ChatContextProvider.collectChatContextList(request.project, creationContext)
         }

diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/ContextPrompter.kt b/src/main/kotlin/cc/unitmesh/devti/provider/ContextPrompter.kt
@@ -13,6 +13,7 @@ import com.intellij.serviceContainer.LazyExtensionInstance
 import com.intellij.util.xmlb.annotations.Attribute
 
 abstract class ContextPrompter : LazyExtensionInstance<ContextPrompter>() {
+    private var element: PsiElement? = null
     protected var action: ChatActionType? = null
     protected var selectedText: String = ""
     protected var file: PsiFile? = null
@@ -57,6 +58,7 @@ abstract class ContextPrompter : LazyExtensionInstance<ContextPrompter>() {
         this.project = project
         this.lang = file?.language?.displayName ?: ""
         this.offset = offset
+        this.element = element
     }
 
     @Attribute("language")

diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/context/ChatCreationContext.kt b/src/main/kotlin/cc/unitmesh/devti/provider/context/ChatCreationContext.kt
@@ -9,5 +9,5 @@ data class ChatCreationContext(
     val action: ChatActionType,
     val sourceFile: PsiFile?,
     val extraItems: List<ChatContextItem> = emptyList(),
-    val element: PsiElement? = null
+    val element: PsiElement?
 )
diff --git a/webstorm/src/main/kotlin/cc/unitmesh/ide/webstorm/provider/JavaScriptContextPrompter.kt b/webstorm/src/main/kotlin/cc/unitmesh/ide/webstorm/provider/JavaScriptContextPrompter.kt
@@ -27,7 +27,7 @@ class JavaScriptContextPrompter : ContextPrompter() {
         element: PsiElement?
     ) {
         super.initContext(actionType, selectedText, file, project, offset, element)
-        creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file)
+        creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file, listOf(), element)
     }
 
     override fun displayPrompt(): String {