diff --git a/clion/src/main/kotlin/cc/unitmesh/clion/provider/RustContextPrompter.kt b/clion/src/main/kotlin/cc/unitmesh/clion/provider/RustContextPrompter.kt index 6763dd96e5..939bab050a 100644 --- 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 new file mode 100644 index 0000000000..f0385f7f7c --- /dev/null +++ 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 { + project.modules.asSequence() + .map { PyProjectSdkConfigurationExtension.findForModule(it) } + .forEach { println(it) } + + return listOf() + } +} \ No newline at end of file diff --git a/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonContextPrompter.kt b/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonContextPrompter.kt index 59834189c6..758f4803b8 100644 --- 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 index 2e025e864c..d899c21cc0 100644 --- a/pycharm/src/main/resources/cc.unitmesh.pycharm.xml +++ b/pycharm/src/main/resources/cc.unitmesh.pycharm.xml @@ -14,6 +14,8 @@ + + diff --git a/src/main/kotlin/cc/unitmesh/devti/actions/chat/GenerateCodeInplaceAction.kt b/src/main/kotlin/cc/unitmesh/devti/actions/chat/GenerateCodeInplaceAction.kt index 688cb5c941..81aadad506 100644 --- 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, 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>, 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 index dcd5d97af2..9663cf6b49 100644 --- 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 = 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 index 1400e4a3c9..511cde79ef 100644 --- 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() { + 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() { 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 index 77ae4e0372..b4a03e011b 100644 --- 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 = emptyList(), - val element: PsiElement? = null + val element: PsiElement? ) \ No newline at end of file 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 index 862ed75d81..9f90312872 100644 --- 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 {