chore: Fix some code quality warnings

fmueller · Nov 20, 2024 · bf2c674 · bf2c674
1 parent 37f93ee
commit bf2c674
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 8 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
@@ -23,9 +23,9 @@ object OllamaService {
     // TODO add something about selected code has higher priority than the open files or references
     private val systemPrompt = """
                     You are Jarvis, an intelligent and helpful coding assistant on the level of an expert software developer. You assist users by providing code completions, debugging tips, explanations, and suggestions in various programming languages. Your responses are clear, concise, and directly address the user's needs.
-                    
+
                     All your responses should be formatted in Markdown. To improve readability, use paragraphs, lists, and headlines where appropriate.
-                    
+
                     - Follow coding best practices for readability, efficiency, and security.
                     - Use appropriate code blocks with syntax highlighting.
                     - Include comments and brief explanations where necessary.
@@ -134,7 +134,7 @@ object OllamaService {
         // TODO check if model is available
         // TODO if not, download model
 
-        var nextMessagePrompt =
+        val nextMessagePrompt =
             if (conversation.getLastUserMessage() != null) {
                 val lastUserMessage = conversation.getLastUserMessage()!!
                 """
@@ -144,7 +144,7 @@ object OllamaService {
                 |
                 |[Code Context]:
                 |
-                |${getCodeContextPrompt(lastUserMessage.codeContext, useCodeContext)}
+                |${getCodeContextPrompt(lastUserMessage.codeContext, !lastUserMessage.isHelpMessage() && useCodeContext)}
                 |
                 |[Assistant]: """.trimMargin()
             } else {
@@ -153,7 +153,7 @@ object OllamaService {
 
         val responseInFlight = StringBuilder()
         try {
-            suspendCancellableCoroutine<String> { continuation ->
+            suspendCancellableCoroutine { continuation ->
                 assistant
                     .chat(nextMessagePrompt)
                     .onNext { update ->

diff --git a/src/main/kotlin/com/github/fmueller/jarvis/ai/SimpleTokenizer.kt b/src/main/kotlin/com/github/fmueller/jarvis/ai/SimpleTokenizer.kt
@@ -2,7 +2,10 @@ package com.github.fmueller.jarvis.ai
 
 import dev.langchain4j.agent.tool.ToolExecutionRequest
 import dev.langchain4j.agent.tool.ToolSpecification
+import dev.langchain4j.data.message.AiMessage
 import dev.langchain4j.data.message.ChatMessage
+import dev.langchain4j.data.message.SystemMessage
+import dev.langchain4j.data.message.UserMessage
 import dev.langchain4j.model.Tokenizer
 import kotlin.math.roundToInt
 
@@ -16,7 +19,16 @@ class SimpleTokenizer : Tokenizer {
     }
 
     override fun estimateTokenCountInMessage(message: ChatMessage?): Int {
-        return estimateTokenCountInText(message?.text() ?: "") + 2
+        return estimateTokenCountInText(getText(message)) + 2
+    }
+
+    private fun getText(message: ChatMessage?): String {
+        return when (message) {
+            is AiMessage -> message.text()
+            is SystemMessage -> message.text()
+            is UserMessage -> if (message.hasSingleText()) message.singleText() else ""
+            else -> ""
+        }
     }
 
     override fun estimateTokenCountInMessages(messages: Iterable<ChatMessage?>?): Int {

diff --git a/src/main/kotlin/com/github/fmueller/jarvis/conversation/CodeContextHelper.kt b/src/main/kotlin/com/github/fmueller/jarvis/conversation/CodeContextHelper.kt
@@ -8,7 +8,7 @@ import com.intellij.psi.PsiManager
 
 object CodeContextHelper {
 
-    fun getCodeContext(project: Project): CodeContext? {
+    fun getCodeContext(project: Project): CodeContext {
         val editor = FileEditorManager.getInstance(project).selectedTextEditor
         if (editor != null) {
             val selection = editor.selectionModel

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,4 +1,4 @@
-package com.github.fmueller.jarvis.conversation;
+package com.github.fmueller.jarvis.conversation
 
 import com.github.fmueller.jarvis.ui.SyntaxHighlightedCodeHelper
 import com.intellij.openapi.Disposable