feat: add basic handle for cache

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

Hide file tree

Showing 4 changed files with 23 additions and 20 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
@@ -29,6 +29,7 @@ open class JvmIdeaContextPrompter : ContextPrompter() {
     private lateinit var mvcContextService: MvcContextService
     private var fileName = ""
     private lateinit var changeListManager: ChangeListManager
+    private lateinit var creationContext: ChatCreationContext
 
     override fun appendAdditionContext(context: String) {
         additionContext += context
@@ -47,6 +48,7 @@ open class JvmIdeaContextPrompter : ContextPrompter() {
 
         lang = file?.language?.displayName ?: ""
         fileName = file?.name ?: ""
+        creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file)
     }
 
     init {
@@ -55,17 +57,9 @@ open class JvmIdeaContextPrompter : ContextPrompter() {
     }
 
     override fun displayPrompt(): String {
-        val creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file)
-
         return runBlocking {
             val prompt = createPrompt(selectedText)
-            var chatContext = ""
-
-            val contextItems = ChatContextProvider.collectChatContextList(project!!, creationContext)
-            contextItems.forEach {
-                chatContext += it.text + "\n"
-            }
-
+            val chatContext = collectionContext(creationContext)
 
             val finalPrompt = if (additionContext.isNotEmpty()) {
                 "```\n$chatContext\n$additionContext\n```\n```$lang\n$selectedText\n```\n"
@@ -78,16 +72,10 @@ open class JvmIdeaContextPrompter : ContextPrompter() {
     }
 
     override fun requestPrompt(): String {
-        val creationContext = ChatCreationContext(ChatOrigin.ChatAction, action!!, file)
 
         return runBlocking {
             val prompt = createPrompt(selectedText)
-            var chatContext = ""
-
-            val contextItems = ChatContextProvider.collectChatContextList(project!!, creationContext)
-            contextItems.forEach {
-                chatContext += it.text + "\n"
-            }
+            val chatContext = collectionContext(creationContext)
 
             val finalPrompt = if (additionContext.isNotEmpty()) {
                 "\n$chatContext\n$additionContext\n```$lang\n$selectedText\n```\n"
@@ -99,6 +87,23 @@ 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/context/ChatCreationContext.kt b/src/main/kotlin/cc/unitmesh/devti/provider/context/ChatCreationContext.kt
@@ -3,7 +3,7 @@ package cc.unitmesh.devti.provider.context
 import cc.unitmesh.devti.gui.chat.ChatActionType
 import com.intellij.psi.PsiFile
 
-class ChatCreationContext(
+data class ChatCreationContext(
     val origin: ChatOrigin,
     val action: ChatActionType,
     val sourceFile: PsiFile?,

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
@@ -21,8 +21,7 @@ class JavaScriptContextPrompter : ContextPrompter() {
         return runBlocking {
             val contextItems = ChatContextProvider.collectChatContextList(project!!, creationContext)
             contextItems.forEach {
-
-                 += it.text + "\n"
+                additionContext += it.text + "\n"
             }
 
             return@runBlocking "${action!!.instruction(lang)}:\n```markdown\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,7 +32,6 @@ 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,