feat: Add proper memory management

- Dispose editors when project is unloaded - Make OllamaService an object - Couple other simplifications of lifecycle management
fmueller · Jun 4, 2024 · 5dbce9a · 5dbce9a
1 parent e94d6a1
commit 5dbce9a
Show file tree

Hide file tree

Showing 8 changed files with 58 additions and 26 deletions.
diff --git a/src/main/kotlin/com/github/fmueller/jarvis/ai/OllamaService.kt b/src/main/kotlin/com/github/fmueller/jarvis/ai/OllamaService.kt
@@ -1,7 +1,6 @@
 package com.github.fmueller.jarvis.ai
 
 import com.github.fmueller.jarvis.conversation.Message
-import com.intellij.openapi.components.Service
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.withContext
 import kotlinx.serialization.Serializable
@@ -26,8 +25,7 @@ private data class ChatRequest(
 @Serializable
 private data class ChatResponse(val message: ChatMessage)
 
-@Service(Service.Level.PROJECT)
-class OllamaService {
+object OllamaService {
 
     suspend fun chat(messages: List<Message>): String = withContext(Dispatchers.IO) {
         // TODO check if model is available

diff --git a/src/main/kotlin/com/github/fmueller/jarvis/commands/ChatCommand.kt b/src/main/kotlin/com/github/fmueller/jarvis/commands/ChatCommand.kt
@@ -5,10 +5,10 @@ import com.github.fmueller.jarvis.conversation.Conversation
 import com.github.fmueller.jarvis.conversation.Message
 import com.github.fmueller.jarvis.conversation.Role
 
-class ChatCommand(private val ollamaService: OllamaService) : SlashCommand {
+class ChatCommand : SlashCommand {
 
     override suspend fun run(conversation: Conversation): Conversation {
-        if (!ollamaService.isAvailable()) {
+        if (!OllamaService.isAvailable()) {
             conversation.addMessage(
                 Message(
                     Role.ASSISTANT,
@@ -18,7 +18,7 @@ class ChatCommand(private val ollamaService: OllamaService) : SlashCommand {
             return conversation
         }
 
-        val response = ollamaService.chat(conversation.messages).trim()
+        val response = OllamaService.chat(conversation.messages).trim()
         conversation.addMessage(Message(Role.ASSISTANT, response))
         return conversation
     }

diff --git a/src/main/kotlin/com/github/fmueller/jarvis/commands/SlashCommandParser.kt b/src/main/kotlin/com/github/fmueller/jarvis/commands/SlashCommandParser.kt
@@ -1,8 +1,6 @@
 package com.github.fmueller.jarvis.commands
 
-import com.github.fmueller.jarvis.ai.OllamaService
-
-class SlashCommandParser(private val ollamaService: OllamaService) {
+object SlashCommandParser {
 
     fun parse(message: String): SlashCommand {
         val trimmedMessage = message.trim().lowercase()
@@ -14,6 +12,6 @@ class SlashCommandParser(private val ollamaService: OllamaService) {
             return NewConversationCommand()
         }
 
-        return ChatCommand(ollamaService)
+        return ChatCommand()
     }
 }
diff --git a/src/main/kotlin/com/github/fmueller/jarvis/conversation/Conversation.kt b/src/main/kotlin/com/github/fmueller/jarvis/conversation/Conversation.kt
@@ -1,7 +1,8 @@
 package com.github.fmueller.jarvis.conversation
 
-import com.github.fmueller.jarvis.ai.OllamaService
 import com.github.fmueller.jarvis.commands.SlashCommandParser
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.Service
 import java.beans.PropertyChangeListener
 import java.beans.PropertyChangeSupport
 import java.time.LocalDateTime
@@ -16,19 +17,20 @@ enum class Role {
 
 data class Message(val role: Role, val content: String, val createdAt: LocalDateTime = LocalDateTime.now())
 
-class Conversation(ollamaService: OllamaService) {
+// as long as we don't have conversation history persistence,
+// we can keep the conversation in memory
+// and declare it as a project-level service
+@Service(Service.Level.PROJECT)
+class Conversation : Disposable {
 
     private var _messages = mutableListOf<Message>()
     val messages get() = _messages.toList()
 
-    private val commandParser = SlashCommandParser(ollamaService)
     private val propertyChangeSupport = PropertyChangeSupport(this)
 
     suspend fun chat(message: String): Conversation {
         addMessage(Message(Role.USER, message.trim()))
-
-        val command = commandParser.parse(message)
-        return command.run(this)
+        return SlashCommandParser.parse(message).run(this)
     }
 
     fun addMessage(message: Message) {
@@ -47,4 +49,10 @@ class Conversation(ollamaService: OllamaService) {
     fun addPropertyChangeListener(listener: PropertyChangeListener) {
         propertyChangeSupport.addPropertyChangeListener(listener)
     }
+
+    override fun dispose() {
+        propertyChangeSupport.getPropertyChangeListeners("messages").forEach {
+            propertyChangeSupport.removePropertyChangeListener(it)
+        }
+    }
 }
diff --git a/src/main/kotlin/com/github/fmueller/jarvis/conversation/ConversationPanel.kt b/src/main/kotlin/com/github/fmueller/jarvis/conversation/ConversationPanel.kt
@@ -1,14 +1,16 @@
 package com.github.fmueller.jarvis.conversation
 
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
 import com.intellij.ui.components.JBScrollPane
 import org.jdesktop.swingx.VerticalLayout
 import java.awt.event.MouseAdapter
 import java.awt.event.MouseEvent
 import javax.swing.JPanel
 import javax.swing.SwingUtilities
 
-class ConversationPanel(conversation: Conversation, private val project: Project) {
+class ConversationPanel(conversation: Conversation, private val project: Project) : Disposable {
 
     private val panel = JPanel().apply {
         layout = VerticalLayout(1)
@@ -22,6 +24,8 @@ class ConversationPanel(conversation: Conversation, private val project: Project
     }
 
     init {
+        Disposer.register(conversation, this)
+
         var isUserScrolling = false
 
         scrollableContainer.verticalScrollBar.addAdjustmentListener { e ->
@@ -53,11 +57,23 @@ class ConversationPanel(conversation: Conversation, private val project: Project
     }
 
     private fun update(messages: List<Message>) {
-        panel.removeAll()
-        messages.forEach { message ->
-            panel.add(MessagePanel(message, project))
+        if (messages.size <= 1) {
+            panel.components.filter { it is Disposable }.map { it as Disposable }.forEach { it.dispose() }
+            panel.removeAll()
+        }
+
+        if (messages.isNotEmpty()) {
+            val messagePanel = MessagePanel(messages.last(), project)
+            Disposer.register(this, messagePanel)
+            panel.add(messagePanel)
         }
+
         panel.revalidate()
         panel.repaint()
     }
+
+    override fun dispose() {
+        // nothing to dispose here, all message panels are disposed when the panel is cleared
+        // or automatically when the conversation is disposed
+    }
 }
diff --git a/src/main/kotlin/com/github/fmueller/jarvis/conversation/MessagePanel.kt b/src/main/kotlin/com/github/fmueller/jarvis/conversation/MessagePanel.kt
@@ -1,6 +1,7 @@
 package com.github.fmueller.jarvis.conversation;
 
 import com.github.fmueller.jarvis.ui.SyntaxHighlightedCodeHelper
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.editor.colors.EditorColorsManager
 import com.intellij.openapi.editor.colors.TextAttributesKey
 import com.intellij.openapi.project.Project
@@ -22,15 +23,15 @@ import javax.swing.BorderFactory
 import javax.swing.JEditorPane
 import javax.swing.JPanel
 
-class MessagePanel(private val message: Message, project: Project) : JPanel() {
+class MessagePanel(private val message: Message, project: Project) : JPanel(), Disposable {
 
     private companion object {
         private val codeBlockPattern = Pattern.compile("```(\\w+)?\\n(.*?)\\n```", Pattern.DOTALL)
         private val assistantBgColor = { UIUtil.getPanelBackground() }
         private val userBgColor = { UIUtil.getTextFieldBackground() }
     }
 
-    private val syntaxHelper = SyntaxHighlightedCodeHelper(project)
+    private val highlightedCodeHelper = SyntaxHighlightedCodeHelper(project)
 
     init {
         buildPanel()
@@ -41,12 +42,17 @@ class MessagePanel(private val message: Message, project: Project) : JPanel() {
         buildPanel()
     }
 
+    override fun dispose() {
+        highlightedCodeHelper.disposeAllEditors()
+    }
+
     private fun buildPanel() {
         if (message == null) {
             return
         }
 
         removeAll()
+        highlightedCodeHelper.disposeAllEditors()
 
         layout = VerticalLayout(5)
         background = if (message.role == Role.ASSISTANT) assistantBgColor() else userBgColor()
@@ -126,7 +132,7 @@ class MessagePanel(private val message: Message, project: Project) : JPanel() {
     }
 
     private fun addHighlightedCode(languageId: String, code: String) {
-        val editor = syntaxHelper.getHighlightedEditor(languageId, code)
+        val editor = highlightedCodeHelper.getHighlightedEditor(languageId, code)
         if (editor != null) {
             editor.contentComponent.border = BorderFactory.createEmptyBorder(5, 5, 5, 5)
             val outerPanelBackground = background

diff --git a/src/main/kotlin/com/github/fmueller/jarvis/toolWindow/ConversationWindowFactory.kt b/src/main/kotlin/com/github/fmueller/jarvis/toolWindow/ConversationWindowFactory.kt
@@ -1,6 +1,5 @@
 package com.github.fmueller.jarvis.toolWindow
 
-import com.github.fmueller.jarvis.ai.OllamaService
 import com.github.fmueller.jarvis.conversation.Conversation
 import com.github.fmueller.jarvis.conversation.ConversationPanel
 import com.github.fmueller.jarvis.conversation.InputArea
@@ -33,7 +32,7 @@ class ConversationWindowFactory : ToolWindowFactory {
 
     class ConversationWindow(toolWindow: ToolWindow) {
 
-        private val conversation = Conversation(toolWindow.project.service<OllamaService>())
+        private val conversation = toolWindow.project.service<Conversation>()
         private val conversationPanel = ConversationPanel(conversation, toolWindow.project)
 
         @OptIn(DelicateCoroutinesApi::class)

diff --git a/src/main/kotlin/com/github/fmueller/jarvis/ui/SyntaxHighlightedCodeHelper.kt b/src/main/kotlin/com/github/fmueller/jarvis/ui/SyntaxHighlightedCodeHelper.kt
@@ -11,11 +11,12 @@ import java.time.format.DateTimeFormatter
 
 class SyntaxHighlightedCodeHelper(private val project: Project) {
 
+    private val createdEditors = mutableListOf<Editor>()
+
     companion object {
         val formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd_HH-mm-ss.SSS")
     }
 
-    // TODO add disposal of created editors
     fun getHighlightedEditor(languageId: String, code: String): Editor? {
         // TODO test with some Kotlin code and see if the language is correctly detected
         val language = Language.findLanguageByID(languageId) ?: IgnoreLanguage.INSTANCE
@@ -29,6 +30,7 @@ class SyntaxHighlightedCodeHelper(private val project: Project) {
             )
 
         val editor = EditorFactory.getInstance().createEditor(file.viewProvider.document, project, fileType, true)
+        createdEditors.add(editor)
 
         editor.settings.apply {
             isLineNumbersShown = false
@@ -48,4 +50,9 @@ class SyntaxHighlightedCodeHelper(private val project: Project) {
         }
         return editor
     }
+
+    fun disposeAllEditors() {
+        createdEditors.forEach { EditorFactory.getInstance().releaseEditor(it) }
+        createdEditors.clear()
+    }
 }