feat: update for js context

unit-mesh · Aug 1, 2023 · 4c2d883 · 4c2d883
1 parent aee8504
commit 4c2d883
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 32 deletions.
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/flow/provider/JvmIdeaContextPrompter.kt b/java/src/main/kotlin/cc/unitmesh/idea/flow/provider/JvmIdeaContextPrompter.kt
@@ -87,23 +87,6 @@ open class JvmIdeaContextPrompter : ContextPrompter() {
         }
     }
 
-    val chatContextCache: MutableMap<ChatCreationContext, String> = mutableMapOf()
-    private suspend fun collectionContext(creationContext: ChatCreationContext): String {
-        if (chatContextCache.containsKey(creationContext)) {
-            return chatContextCache[creationContext]!!
-        }
-
-        var chatContext = ""
-
-        val contextItems = ChatContextProvider.collectChatContextList(project!!, creationContext)
-        contextItems.forEach {
-            chatContext += it.text + "\n"
-        }
-
-        chatContextCache[creationContext] = chatContext
-        return chatContext
-    }
-
 
     private suspend fun createPrompt(selectedText: String): String {
         var prompt = action!!.instruction(lang)

diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/ContextPrompter.kt b/src/main/kotlin/cc/unitmesh/devti/provider/ContextPrompter.kt
@@ -1,6 +1,8 @@
 package cc.unitmesh.devti.provider
 
 import cc.unitmesh.devti.gui.chat.ChatActionType
+import cc.unitmesh.devti.provider.context.ChatContextProvider
+import cc.unitmesh.devti.provider.context.ChatCreationContext
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.project.Project
@@ -16,6 +18,23 @@ abstract class ContextPrompter : LazyExtensionInstance<ContextPrompter>() {
     protected var lang: String = ""
     protected var offset: Int = 0
 
+    val chatContextCache: MutableMap<ChatCreationContext, String> = mutableMapOf()
+    suspend fun collectionContext(creationContext: ChatCreationContext): String {
+        if (chatContextCache.containsKey(creationContext)) {
+            return chatContextCache[creationContext]!!
+        }
+
+        var chatContext = ""
+
+        val contextItems = ChatContextProvider.collectChatContextList(project!!, creationContext)
+        contextItems.forEach {
+            chatContext += it.text + "\n"
+        }
+
+        chatContextCache[creationContext] = chatContext
+        return chatContext
+    }
+
     open fun initContext(
         actionType: ChatActionType,
         selectedText: String,

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
@@ -1,43 +1,45 @@
 package cc.unitmesh.ide.webstorm.provider
 
+import cc.unitmesh.devti.gui.chat.ChatActionType
 import cc.unitmesh.devti.provider.ContextPrompter
 import cc.unitmesh.devti.provider.context.ChatContextProvider
 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.PsiFile
 import kotlinx.coroutines.runBlocking
 
 class JavaScriptContextPrompter : ContextPrompter() {
+    private lateinit var creationContext: ChatCreationContext
     private var additionContext = ""
 
     companion object {
         val log = logger<JavaScriptContextPrompter>()
     }
 
+    override fun initContext(
+        actionType: ChatActionType,
+        selectedText: String,
+        file: PsiFile?,
+        project: Project,
+        offset: Int
+    ) {
+        super.initContext(actionType, selectedText, file, project, offset)
+        creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file)
+    }
+
     override fun displayPrompt(): String {
-        val creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file)
-        additionContext = ""
 
         return runBlocking {
-            val contextItems = ChatContextProvider.collectChatContextList(project!!, creationContext)
-            contextItems.forEach {
-                additionContext += it.text + "\n"
-            }
-
+            additionContext = collectionContext(creationContext)
             return@runBlocking "${action!!.instruction(lang)}:\n```markdown\n$additionContext```\n```${lang}\n$selectedText\n```"
         }
     }
 
     override fun requestPrompt(): String {
-        val creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file)
-        additionContext = ""
-
         return runBlocking {
-            val contextItems = ChatContextProvider.collectChatContextList(project!!, creationContext)
-            contextItems.forEach {
-                additionContext += it.text + "\n"
-            }
-
+            additionContext = collectionContext(creationContext)
             return@runBlocking "${action!!.instruction(lang)}:\n$additionContext\n```${lang}\n$selectedText\n```"
         }
     }

diff --git a/webstorm/src/main/kotlin/cc/unitmesh/ide/webstorm/provider/JavaScriptContextProvider.kt b/webstorm/src/main/kotlin/cc/unitmesh/ide/webstorm/provider/JavaScriptContextProvider.kt
@@ -32,6 +32,7 @@ class JavaScriptContextProvider : ChatContextProvider {
         if (mostPopularPackagesContext != null) results.add(mostPopularPackagesContext)
 
         val techStack = prepareStack(snapshot)
+        logger<JavaScriptContextProvider>().warn("Tech stack: $techStack")
         if (techStack.coreFrameworks().isNotEmpty()) {
             val element = ChatContextItem(
                 JavaScriptContextProvider::class,