feat: add offset for handle prompt text

unit-mesh · Jul 20, 2023 · 5bdf188 · 5bdf188
1 parent 174afea
commit 5bdf188
Show file tree

Hide file tree

Showing 7 changed files with 11 additions and 13 deletions.
diff --git a/idea/src/main/kotlin/cc/unitmesh/idea/kotlin/KotlinContextPrompter.kt b/idea/src/main/kotlin/cc/unitmesh/idea/kotlin/KotlinContextPrompter.kt
@@ -14,9 +14,9 @@ class KotlinContextPrompter : ContextPrompter() {
     private var project: Project? = null
     private var lang: String = ""
 
-    override fun initContext(actionType: ChatBotActionType, prefixText: String, file: PsiFile?, project: Project) {
+    override fun initContext(actionType: ChatBotActionType, text: String, file: PsiFile?, project: Project, offset: Int) {
         this.action = actionType
-        this.selectedText = prefixText
+        this.selectedText = text
         this.file = file
         this.project = project
         this.lang = file?.language?.displayName ?: ""

diff --git a/idea/src/main/kotlin/cc/unitmesh/idea/provider/JavaContextPrompter.kt b/idea/src/main/kotlin/cc/unitmesh/idea/provider/JavaContextPrompter.kt
@@ -33,9 +33,9 @@ class JavaContextPrompter : ContextPrompter() {
     private val isService = fileName.endsWith("Service.java") || fileName.endsWith("ServiceImpl.java")
     private lateinit var changeListManager: ChangeListManager
 
-    override fun initContext(actionType: ChatBotActionType, prefixText: String, file: PsiFile?, project: Project) {
+    override fun initContext(actionType: ChatBotActionType, text: String, file: PsiFile?, project: Project, offset: Int) {
         this.action = actionType
-        this.prefixText = prefixText
+        this.prefixText = text
         this.file = file
         this.project = project
         changeListManager = ChangeListManagerImpl.getInstance(project)

diff --git a/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonContextPrompter.kt b/pycharm/src/main/kotlin/cc/unitmesh/pycharm/provider/PythonContextPrompter.kt
@@ -14,9 +14,9 @@ class PythonContextPrompter : ContextPrompter() {
     private var project: Project? = null
     private var lang: String = ""
 
-    override fun initContext(actionType: ChatBotActionType, prefixText: String, file: PsiFile?, project: Project) {
+    override fun initContext(actionType: ChatBotActionType, text: String, file: PsiFile?, project: Project, offset: Int) {
         this.action = actionType
-        this.selectedText = prefixText
+        this.selectedText = text
         this.file = file
         this.project = project
         this.lang = file?.language?.displayName ?: ""

diff --git a/src/main/kotlin/cc/unitmesh/devti/actions/chat/ChatBaseAction.kt b/src/main/kotlin/cc/unitmesh/devti/actions/chat/ChatBaseAction.kt
@@ -53,7 +53,7 @@ abstract class ChatBaseAction : AnAction() {
 
         val prompter = ContextPrompter.prompter(file?.language?.displayName ?: "")
         logger.info("use prompter: ${prompter?.javaClass}")
-        prompter?.initContext(actionType, prefixText, file, project)
+        prompter?.initContext(actionType, prefixText, file, project, caretModel?.offset ?: 0)
 
         toolWindowManager?.activate {
             val chatContext = ChatContext(

diff --git a/src/main/kotlin/cc/unitmesh/devti/intentions/editor/AbstractChatIntention.kt b/src/main/kotlin/cc/unitmesh/devti/intentions/editor/AbstractChatIntention.kt
@@ -52,7 +52,7 @@ abstract class AbstractChatIntention : IntentionAction {
 
 
         val prompter = ContextPrompter.prompter(file.language.displayName)
-        prompter?.initContext(actionType, selectedText, file, project)
+        prompter?.initContext(actionType, selectedText, file, project, editor.caretModel.offset)
         sendToChat(project, actionType, prompter!!)
     }
 

diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/ContextPrompter.kt b/src/main/kotlin/cc/unitmesh/devti/provider/ContextPrompter.kt
@@ -1,8 +1,6 @@
 package cc.unitmesh.devti.provider
 
 import cc.unitmesh.devti.gui.chat.ChatBotActionType
-import cc.unitmesh.devti.prompting.model.PromptConfig
-import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiFile
@@ -22,7 +20,7 @@ abstract class ContextPrompter : LazyExtensionInstance<ContextPrompter>() {
 
     open fun getUIPrompt(): String = ""
     open fun getRequestPrompt(): String = ""
-    open fun initContext(actionType: ChatBotActionType, prefixText: String, file: PsiFile?, project: Project) {}
+    open fun initContext(actionType: ChatBotActionType, text: String, file: PsiFile?, project: Project, offset: Int) {}
 
     companion object {
         private val EP_NAME: ExtensionPointName<ContextPrompter> =

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
@@ -12,9 +12,9 @@ class JavaScriptContextPrompter : ContextPrompter() {
     private var project: Project? = null
 
     private val lang: String = file?.language?.displayName ?: ""
-    override fun initContext(actionType: ChatBotActionType, prefixText: String, file: PsiFile?, project: Project) {
+    override fun initContext(actionType: ChatBotActionType, text: String, file: PsiFile?, project: Project, offset: Int) {
         this.action = actionType
-        this.selectedText = prefixText
+        this.selectedText = text
         this.file = file
         this.project = project
     }