feat(openAi): retrieve token async

Author: Blarc

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/AICommitsUtils.kt

@@ -5,7 +5,7 @@ import com.github.blarc.ai.commits.intellij.plugin.notifications.sendNotificatio
 import com.github.blarc.ai.commits.intellij.plugin.settings.AppSettings2
 import com.github.blarc.ai.commits.intellij.plugin.settings.ProjectSettings
 import com.intellij.credentialStore.CredentialAttributes
-import com.intellij.credentialStore.Credentials
+import com.intellij.credentialStore.OneTimeString
 import com.intellij.ide.passwordSafe.PasswordSafe
 import com.intellij.openapi.components.service
 import com.intellij.openapi.diff.impl.patch.IdeaTextPatchBuilder
@@ -15,6 +15,8 @@ import com.intellij.openapi.vcs.changes.Change
 import com.knuddels.jtokkit.Encodings
 import com.knuddels.jtokkit.api.ModelType
 import git4idea.repo.GitRepositoryManager
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.withContext
 import java.io.StringWriter
 import java.nio.file.FileSystems
 
@@ -138,9 +140,21 @@ object AICommitsUtils {
         return encoding.countTokens(prompt) > modelType.maxContextLength
     }
 
-    fun retrieveToken(title: String): String? {
-        val credentials: Credentials? = PasswordSafe.instance.get(getCredentialAttributes(title))
-        return credentials?.getPasswordAsString()
+    // TODO @Blarc: Slow operations are prohibited on EDT
+    fun saveToken(title: String, token: String) {
+        try {
+            PasswordSafe.instance.setPassword(getCredentialAttributes(title), token)
+        } catch (e: Exception) {
+            sendNotification(Notification.unableToSaveToken(e.message))
+        }
+    }
+
+    suspend fun retrieveToken(title: String): OneTimeString? {
+        val credentialAttributes = getCredentialAttributes(title)
+        val credentials = withContext(Dispatchers.IO) {
+            PasswordSafe.instance.get(credentialAttributes)
+        }
+        return credentials?.password
     }
 
     fun getCredentialAttributes(title: String): CredentialAttributes {

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/AppSettings2.kt

@@ -1,13 +1,16 @@
 package com.github.blarc.ai.commits.intellij.plugin.settings
 
 import com.github.blarc.ai.commits.intellij.plugin.AICommitsUtils
+import com.github.blarc.ai.commits.intellij.plugin.AICommitsUtils.getCredentialAttributes
+import com.github.blarc.ai.commits.intellij.plugin.AICommitsUtils.saveToken
 import com.github.blarc.ai.commits.intellij.plugin.notifications.Notification
 import com.github.blarc.ai.commits.intellij.plugin.notifications.sendNotification
 import com.github.blarc.ai.commits.intellij.plugin.settings.clients.LLMClientConfiguration
 import com.github.blarc.ai.commits.intellij.plugin.settings.clients.ollama.OllamaClientConfiguration
 import com.github.blarc.ai.commits.intellij.plugin.settings.clients.openAi.OpenAiClientConfiguration
 import com.github.blarc.ai.commits.intellij.plugin.settings.clients.openAi.OpenAiClientSharedState
 import com.github.blarc.ai.commits.intellij.plugin.settings.prompts.DefaultPrompts
+import com.intellij.ide.passwordSafe.PasswordSafe
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.PersistentStateComponent
 import com.intellij.openapi.components.Service
@@ -72,7 +75,7 @@ class AppSettings2 : PersistentStateComponent<AppSettings2> {
         val appSettings = AppSettings.instance
         migrateSettingsFromVersion1(appSettings)
         val openAiLlmClient = llmClientConfigurations.find { it.displayName == "OpenAI" }
-        migrateOpenAiClientFromVersion1(openAiLlmClient, appSettings)
+        migrateOpenAiClientFromVersion1(openAiLlmClient as OpenAiClientConfiguration, appSettings)
     }
 
     private fun migrateSettingsFromVersion1(appSettings: AppSettings) {
@@ -85,14 +88,22 @@ class AppSettings2 : PersistentStateComponent<AppSettings2> {
         appExclusions = appSettings.appExclusions
     }
 
-    private fun migrateOpenAiClientFromVersion1(openAiLlmClientConfiguration: LLMClientConfiguration?, appSettings: AppSettings) {
+    private fun migrateOpenAiClientFromVersion1(openAiLlmClientConfiguration: OpenAiClientConfiguration?, appSettings: AppSettings) {
         openAiLlmClientConfiguration?.apply {
             host = appSettings.openAIHost
             appSettings.openAISocketTimeout.toIntOrNull()?.let { timeout = it }
             proxyUrl = appSettings.proxyUrl
             modelId = appSettings.openAIModelId
             temperature = appSettings.openAITemperature
-            AICommitsUtils.retrieveToken(appSettings.openAITokenTitle)?.let { token = it }
+
+            val credentialAttributes = getCredentialAttributes(appSettings.openAITokenTitle)
+            PasswordSafe.instance.getAsync(credentialAttributes)
+                .onSuccess {
+                    it?.password?.let { token ->
+                        saveToken(displayName, token.toString(false))
+                        tokenIsStored = true
+                    }
+                }
         }
 
         OpenAiClientSharedState.getInstance().hosts.addAll(appSettings.openAIHosts)

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/LLMClientConfiguration.kt

@@ -1,14 +1,8 @@
 package com.github.blarc.ai.commits.intellij.plugin.settings.clients
 
-import com.github.blarc.ai.commits.intellij.plugin.AICommitsUtils.getCredentialAttributes
-import com.github.blarc.ai.commits.intellij.plugin.AICommitsUtils.retrieveToken
-import com.github.blarc.ai.commits.intellij.plugin.notifications.Notification
-import com.github.blarc.ai.commits.intellij.plugin.notifications.sendNotification
-import com.intellij.ide.passwordSafe.PasswordSafe
 import com.intellij.openapi.ui.ComboBox
 import com.intellij.openapi.vcs.ui.CommitMessage
 import com.intellij.util.xmlb.annotations.Attribute
-import com.intellij.util.xmlb.annotations.Transient
 import javax.swing.Icon
 
 abstract class LLMClientConfiguration(
@@ -19,10 +13,6 @@ abstract class LLMClientConfiguration(
     @Attribute var modelId: String,
     @Attribute var temperature: String,
 ) : Cloneable, Comparable<LLMClientConfiguration> {
-    @get:Transient
-    var token: String
-        get() = retrieveToken(displayName) ?: ""
-        set(token) = saveToken(token)
 
     abstract fun getIcon(): Icon
 
@@ -52,14 +42,6 @@ abstract class LLMClientConfiguration(
 
     abstract fun panel(): LLMClientPanel
 
-    private fun saveToken(token: String) {
-        try {
-            PasswordSafe.instance.setPassword(getCredentialAttributes(displayName), token)
-        } catch (e: Exception) {
-            sendNotification(Notification.unableToSaveToken(e.message))
-        }
-    }
-
     override fun compareTo(other: LLMClientConfiguration): Int {
         return displayName.compareTo(other.displayName)
     }

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/LLMClientService.kt

@@ -15,46 +15,51 @@ import kotlinx.coroutines.withContext
 
 abstract class LLMClientService<T : LLMClientConfiguration>(private val cs: CoroutineScope) {
 
-    abstract fun buildChatModel(client: T): ChatLanguageModel
+    abstract suspend fun buildChatModel(client: T): ChatLanguageModel
 
     fun generateCommitMessage(client: T, prompt: String, commitMessage: CommitMessage) {
-        val model = buildChatModel(client)
-        sendRequest(model, prompt, onSuccess = {
-            commitMessage.setCommitMessage(it)
-            AppSettings2.instance.recordHit()
-        }, onError = {
-            commitMessage.setCommitMessage(it)
-        })
+        cs.launch(Dispatchers.Default) {
+            sendRequest(client, prompt, onSuccess = {
+                commitMessage.setCommitMessage(it)
+                AppSettings2.instance.recordHit()
+            }, onError = {
+                commitMessage.setCommitMessage(it)
+            })
+        }
     }
 
     fun verifyConfiguration(client: T, label: JBLabel) {
-        val model = buildChatModel(client)
-        sendRequest(model, "test", onSuccess = {
-            label.text = message("settings.verify.valid")
-            label.icon = AllIcons.General.InspectionsOK
-        }, onError = {
-            label.text = it.wrap(80)
-            label.icon = AllIcons.General.InspectionsError
-        })
+        // TODO @Blarc: Can you make this better?
+        label.text = "Verifying configuration..."
+        cs.launch(Dispatchers.Default) {
+            sendRequest(client, "test", onSuccess = {
+                label.text = message("settings.verify.valid")
+                label.icon = AllIcons.General.InspectionsOK
+            }, onError = {
+                label.text = it.wrap(80)
+                label.icon = AllIcons.General.InspectionsError
+            })
+        }
     }
 
-    private fun sendRequest(model: ChatLanguageModel, text: String, onSuccess: suspend (r: String) -> Unit, onError: suspend (r: String) -> Unit) {
-        cs.launch(Dispatchers.Default) {
-            try {
-                val response = withContext(Dispatchers.IO) {
-                    model.generate(
-                        listOf(
-                            UserMessage.from(
-                                "user",
-                                text
-                            )
+    private suspend fun sendRequest(client: T, text: String, onSuccess: suspend (r: String) -> Unit, onError: suspend (r: String) -> Unit) {
+        try {
+            val model = buildChatModel(client)
+            val response = withContext(Dispatchers.IO) {
+                model.generate(
+                    listOf(
+                        UserMessage.from(
+                            "user",
+                            text
                         )
-                    ).content().text()
-                }
-                onSuccess(response)
-            } catch (e: Exception) {
-                onError(e.message ?: "Unknown error.")
+                    )
+                ).content().text()
             }
+            onSuccess(response)
+        } catch (e: IllegalArgumentException) {
+            onError("Invalid configuration: ${e.message ?: "unknown"}.")
+        } catch (e: Exception) {
+            onError(e.message ?: "Unknown error.")
         }
     }
 }

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/ollama/OllamaClientService.kt

@@ -45,7 +45,7 @@ class OllamaClientService(private val cs: CoroutineScope) : LLMClientService<Oll
         }
     }
 
-    override fun buildChatModel(client: OllamaClientConfiguration): ChatLanguageModel {
+    override suspend fun buildChatModel(client: OllamaClientConfiguration): ChatLanguageModel {
         return OllamaChatModel.builder()
             .modelName(client.modelId)
             .temperature(client.temperature.toDouble())

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/openAi/OpenAiClientConfiguration.kt

@@ -3,11 +3,11 @@ package com.github.blarc.ai.commits.intellij.plugin.settings.clients.openAi
 import com.github.blarc.ai.commits.intellij.plugin.Icons
 import com.github.blarc.ai.commits.intellij.plugin.settings.clients.LLMClientConfiguration
 import com.intellij.openapi.vcs.ui.CommitMessage
+import com.intellij.util.xmlb.annotations.Transient
 import javax.swing.Icon
 
 class OpenAiClientConfiguration(
-    displayName: String = "OpenAI",
-    var organizationId: String? = null
+    displayName: String = "OpenAI"
 ) : LLMClientConfiguration(
     displayName,
     "https://api.openai.com/v1",
@@ -16,6 +16,11 @@ class OpenAiClientConfiguration(
     "gpt-3.5-turbo",
     "0.7"
 ) {
+    var organizationId: String? = null
+    var tokenIsStored: Boolean = false
+    @Transient
+    var token: String? = null
+
     override fun getIcon(): Icon {
         return Icons.OPEN_AI
     }
@@ -39,6 +44,7 @@ class OpenAiClientConfiguration(
         copy.modelId = modelId
         copy.organizationId = organizationId
         copy.temperature = temperature
+        copy.tokenIsStored = tokenIsStored
         return copy
     }
 

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/openAi/OpenAiClientPanel.kt

@@ -1,6 +1,7 @@
 package com.github.blarc.ai.commits.intellij.plugin.settings.clients.openAi
 
 import com.github.blarc.ai.commits.intellij.plugin.AICommitsBundle.message
+import com.github.blarc.ai.commits.intellij.plugin.AICommitsUtils.saveToken
 import com.github.blarc.ai.commits.intellij.plugin.emptyText
 import com.github.blarc.ai.commits.intellij.plugin.settings.clients.LLMClientPanel
 import com.intellij.ui.components.JBPasswordField
@@ -22,8 +23,10 @@ class OpenAiClientPanel(private val clientConfiguration: OpenAiClientConfigurati
             label(message("settings.llmClient.token"))
                 .widthGroup("label")
             cell(tokenPasswordField)
-                .bindText(clientConfiguration::token)
-                .emptyText(message("settings.openAI.token.example"))
+                .bindText(getter = {""}, setter = {
+                    saveToken(clientConfiguration.displayName, it)
+                })
+                .emptyText(if (clientConfiguration.tokenIsStored) message("settings.openAI.token.stored") else message("settings.openAI.token.example"))
                 .resizableColumn()
                 .focused()
                 .widthGroup("input")

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/openAi/OpenAiClientService.kt

@@ -1,8 +1,10 @@
 package com.github.blarc.ai.commits.intellij.plugin.settings.clients.openAi
 
+import com.github.blarc.ai.commits.intellij.plugin.AICommitsUtils.retrieveToken
 import com.github.blarc.ai.commits.intellij.plugin.settings.clients.LLMClientService
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
+import com.intellij.util.text.nullize
 import dev.langchain4j.model.chat.ChatLanguageModel
 import dev.langchain4j.model.openai.OpenAiChatModel
 import kotlinx.coroutines.CoroutineScope
@@ -19,9 +21,10 @@ class OpenAiClientService(cs: CoroutineScope) : LLMClientService<OpenAiClientCon
         fun getInstance(): OpenAiClientService = service()
     }
 
-    override fun buildChatModel(client: OpenAiClientConfiguration): ChatLanguageModel {
+    override suspend fun buildChatModel(client: OpenAiClientConfiguration): ChatLanguageModel {
+        val token = client.token.nullize(true) ?: retrieveToken(client.displayName)?.toString(true)
         val builder = OpenAiChatModel.builder()
-            .apiKey(client.token)
+            .apiKey(token ?: "")
             .modelName(client.modelId)
             .temperature(client.temperature.toDouble())
             .timeout(Duration.ofSeconds(client.timeout.toLong()))

src/main/resources/messages/MyBundle.properties

@@ -72,6 +72,7 @@ settings.llmClient.temperature.comment=What sampling temperature to use, between
   the output more random,while lower values like 0.2 will make it more focused and deterministic.
 
 settings.openAI.token.example=sk-ABCdefgHIjKlxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+settings.openAI.token.stored=<hidden>
 settings.openAi.token.comment=\
   <p>You can get your token <a href="https://platform.openai.com/account/api-keys">here.</a/></p>
 settings.openAi.organizationId=Organization ID