feat(clients): support streaming response

Blarc · Blarc · commit f54b74e65a3d · 2024-10-08T20:29:43.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 ### Added
 
+- Support streaming response.
 - Support for Hugging Face.
 
 ## [2.5.0] - 2024-09-22
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/AppSettings2.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/AppSettings2.kt
@@ -76,6 +76,9 @@ class AppSettings2 : PersistentStateComponent<AppSettings2> {
 
     var appExclusions: Set<String> = setOf()
 
+    @Attribute
+    var useStreamingResponse: Boolean = true
+
     override fun getState() = this
 
     override fun loadState(state: AppSettings2) {
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/AppSettingsConfigurable.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/AppSettingsConfigurable.kt
@@ -42,6 +42,10 @@ class AppSettingsConfigurable(val project: Project, cs: CoroutineScope) : BoundC
                 .bindSelected(project.service<ProjectSettings>()::isProjectSpecificLLMClient)
             contextHelp(message("settings.llmClient.projectSpecific.contextHelp"))
                 .align(AlignX.LEFT)
+            checkBox(message("settings.llmClient.streamingResponse"))
+                .bindSelected(AppSettings2.instance::useStreamingResponse)
+            contextHelp(message("settings.llmClient.streamingResponse.contextHelp"))
+                .align(AlignX.LEFT)
         }
         row {
             llmClientToolbarDecorator = ToolbarDecorator.createDecorator(llmClientTable.table)
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/LLMClientService.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/LLMClientService.kt
@@ -21,8 +21,11 @@ import com.intellij.platform.ide.progress.withBackgroundProgress
 import com.intellij.ui.components.JBLabel
 import com.intellij.vcs.commit.AbstractCommitWorkflowHandler
 import com.intellij.vcs.commit.isAmendCommitMode
+import dev.langchain4j.data.message.AiMessage
 import dev.langchain4j.data.message.UserMessage
+import dev.langchain4j.model.StreamingResponseHandler
 import dev.langchain4j.model.chat.ChatLanguageModel
+import dev.langchain4j.model.chat.StreamingChatLanguageModel
 import git4idea.GitCommit
 import git4idea.history.GitHistoryUtils
 import git4idea.repo.GitRepositoryManager
@@ -35,6 +38,8 @@ abstract class LLMClientService<C : LLMClientConfiguration>(private val cs: Coro
 
     abstract suspend fun buildChatModel(client: C): ChatLanguageModel
 
+    abstract suspend fun buildStreamingChatModel(client: C): StreamingChatLanguageModel?
+
     fun generateCommitMessage(clientConfiguration: C, commitWorkflowHandler: AbstractCommitWorkflowHandler<*, *>, commitMessage: CommitMessage, project: Project) {
 
         val commitContext = commitWorkflowHandler.workflow.commitContext
@@ -58,7 +63,7 @@ abstract class LLMClientService<C : LLMClientConfiguration>(private val cs: Coro
                 val branch = commonBranch(includedChanges, project)
                 val prompt = constructPrompt(project.service<ProjectSettings>().activePrompt.content, diff, branch, commitMessage.text, project)
 
-                sendRequest(clientConfiguration, prompt, onSuccess = {
+                sendStreamingRequest(clientConfiguration, prompt, onSuccess = {
                     withContext(Dispatchers.EDT) {
                         commitMessage.setCommitMessage(it)
                     }
@@ -72,6 +77,7 @@ abstract class LLMClientService<C : LLMClientConfiguration>(private val cs: Coro
         }
     }
 
+
     fun verifyConfiguration(client: C, label: JBLabel) {
         label.text = message("settings.verify.running")
         cs.launch(ModalityState.current().asContextElement()) {
@@ -89,20 +95,15 @@ abstract class LLMClientService<C : LLMClientConfiguration>(private val cs: Coro
         }
     }
 
-    private suspend fun sendRequest(client: C, text: String, onSuccess: suspend (r: String) -> Unit, onError: suspend (r: String) -> Unit) {
+    private suspend fun sendStreamingRequest(client: C, 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()
+            if (AppSettings2.instance.useStreamingResponse) {
+                buildStreamingChatModel(client)?.let { streamingChatModel ->
+                    sendStreamingRequest(streamingChatModel, text, onSuccess, onError)
+                    return
+                }
             }
-            onSuccess(response)
+            sendRequest(client, text, onSuccess, onError)
         } catch (e: IllegalArgumentException) {
             onError(message("settings.verify.invalid", e.message ?: message("unknown-error")))
         } catch (e: Exception) {
@@ -112,6 +113,50 @@ abstract class LLMClientService<C : LLMClientConfiguration>(private val cs: Coro
         }
     }
 
+    private suspend fun sendStreamingRequest(streamingModel: StreamingChatLanguageModel, text: String, onSuccess: suspend (r: String) -> Unit, onError: suspend (r: String) -> Unit) {
+        var response = ""
+        withContext(Dispatchers.IO) {
+            streamingModel.generate(
+                listOf(
+                    UserMessage.from(
+                        "user",
+                        text
+                    )
+                ),
+                object : StreamingResponseHandler<AiMessage> {
+                    override fun onNext(token: String?) {
+                        response += token
+                        cs.launch {
+                            onSuccess(response)
+                        }
+                    }
+
+                    override fun onError(error: Throwable?) {
+                        response = error?.message.toString()
+                        cs.launch {
+                            onError(response)
+                        }
+                    }
+                }
+            )
+        }
+    }
+
+    private suspend fun sendRequest(client: C, text: String, onSuccess: suspend (r: String) -> Unit, onError: suspend (r: String) -> Unit) {
+        val model = buildChatModel(client)
+        val response = withContext(Dispatchers.IO) {
+            model.generate(
+                listOf(
+                    UserMessage.from(
+                        "user",
+                        text
+                    )
+                )
+            ).content().text()
+        }
+        onSuccess(response)
+    }
+
     private suspend fun getLastCommitChanges(project: Project): List<Change> {
         return withContext(Dispatchers.IO) {
             GitRepositoryManager.getInstance(project).repositories.map { repo ->
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/anthropic/AnthropicClientService.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/anthropic/AnthropicClientService.kt
@@ -10,7 +10,9 @@ import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.util.text.nullize
 import dev.langchain4j.model.anthropic.AnthropicChatModel
+import dev.langchain4j.model.anthropic.AnthropicStreamingChatModel
 import dev.langchain4j.model.chat.ChatLanguageModel
+import dev.langchain4j.model.chat.StreamingChatLanguageModel
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -45,6 +47,22 @@ class AnthropicClientService(private val cs: CoroutineScope) : LLMClientService<
 
     }
 
+    override suspend fun buildStreamingChatModel(client: AnthropicClientConfiguration): StreamingChatLanguageModel {
+        val token = client.token.nullize(true) ?: retrieveToken(client.id)?.toString(true)
+        val builder = AnthropicStreamingChatModel.builder()
+            .modelName(client.modelId)
+            .temperature(client.temperature.toDouble())
+            .apiKey(token ?: "")
+            .baseUrl(client.host)
+            .timeout(Duration.ofSeconds(client.timeout.toLong()))
+
+        client.version?.takeIf { it.isNotBlank() }?.let {
+            builder.version(it)
+        }
+
+        return builder.build()
+    }
+
     fun saveToken(client: AnthropicClientConfiguration, token: String) {
         cs.launch(Dispatchers.Default) {
             try {
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/azureOpenAi/AzureOpenAiClientService.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/azureOpenAi/AzureOpenAiClientService.kt
@@ -10,7 +10,9 @@ import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.util.text.nullize
 import dev.langchain4j.model.azure.AzureOpenAiChatModel
+import dev.langchain4j.model.azure.AzureOpenAiStreamingChatModel
 import dev.langchain4j.model.chat.ChatLanguageModel
+import dev.langchain4j.model.chat.StreamingChatLanguageModel
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -36,6 +38,17 @@ class AzureOpenAiClientService(private val cs: CoroutineScope) : LLMClientServic
             .build()
     }
 
+    override suspend fun buildStreamingChatModel(client: AzureOpenAiClientConfiguration): StreamingChatLanguageModel {
+        val token = client.token.nullize(true) ?: retrieveToken(client.id)?.toString(true)
+        return AzureOpenAiStreamingChatModel.builder()
+            .deploymentName(client.modelId)
+            .temperature(client.temperature.toDouble())
+            .timeout(Duration.ofSeconds(client.timeout.toLong()))
+            .endpoint(client.host)
+            .apiKey(token ?: "")
+            .build()
+    }
+
     fun saveToken(client: AzureOpenAiClientConfiguration, token: String) {
         cs.launch(Dispatchers.Default) {
             try {
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/gemini/GeminiClientService.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/gemini/GeminiClientService.kt
@@ -4,7 +4,9 @@ import com.github.blarc.ai.commits.intellij.plugin.settings.clients.LLMClientSer
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import dev.langchain4j.model.chat.ChatLanguageModel
+import dev.langchain4j.model.chat.StreamingChatLanguageModel
 import dev.langchain4j.model.vertexai.VertexAiGeminiChatModel
+import dev.langchain4j.model.vertexai.VertexAiGeminiStreamingChatModel
 import kotlinx.coroutines.CoroutineScope
 
 @Service(Service.Level.APP)
@@ -24,4 +26,12 @@ class GeminiClientService(private val cs: CoroutineScope): LLMClientService<Gemi
             .build()
     }
 
+    override suspend fun buildStreamingChatModel(client: GeminiClientConfiguration): StreamingChatLanguageModel {
+        return VertexAiGeminiStreamingChatModel.builder()
+            .project(client.projectId)
+            .location(client.location)
+            .modelName(client.modelId)
+            .temperature(client.temperature.toFloat())
+            .build()
+    }
 }
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/huggingface/HuggingFaceClientService.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/huggingface/HuggingFaceClientService.kt
@@ -37,6 +37,8 @@ class HuggingFaceClientService(private val cs: CoroutineScope) : LLMClientServic
             .build()
     }
 
+    override suspend fun buildStreamingChatModel(client: HuggingFaceClientConfiguration) = null
+
     fun saveToken(client: HuggingFaceClientConfiguration, token: String) {
         cs.launch(Dispatchers.Default) {
             try {
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/ollama/OllamaClientService.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/ollama/OllamaClientService.kt
@@ -6,8 +6,10 @@ import com.intellij.openapi.components.service
 import com.intellij.openapi.ui.ComboBox
 import com.intellij.openapi.ui.naturalSorted
 import dev.langchain4j.model.chat.ChatLanguageModel
+import dev.langchain4j.model.chat.StreamingChatLanguageModel
 import dev.langchain4j.model.ollama.OllamaChatModel
 import dev.langchain4j.model.ollama.OllamaModels
+import dev.langchain4j.model.ollama.OllamaStreamingChatModel
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -53,4 +55,13 @@ class OllamaClientService(private val cs: CoroutineScope) : LLMClientService<Oll
             .baseUrl(client.host)
             .build()
     }
+
+    override suspend fun buildStreamingChatModel(client: OllamaClientConfiguration): StreamingChatLanguageModel {
+        return OllamaStreamingChatModel.builder()
+            .modelName(client.modelId)
+            .temperature(client.temperature.toDouble())
+            .timeout(Duration.ofSeconds(client.timeout.toLong()))
+            .baseUrl(client.host)
+            .build()
+    }
 }
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/openAi/OpenAiClientService.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/openAi/OpenAiClientService.kt
@@ -10,7 +10,9 @@ 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.chat.StreamingChatLanguageModel
 import dev.langchain4j.model.openai.OpenAiChatModel
+import dev.langchain4j.model.openai.OpenAiStreamingChatModel
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -48,6 +50,27 @@ class OpenAiClientService(private val cs: CoroutineScope) : LLMClientService<Ope
         return builder.build()
     }
 
+    override suspend fun buildStreamingChatModel(client: OpenAiClientConfiguration): StreamingChatLanguageModel {
+        val token = client.token.nullize(true) ?: retrieveToken(client.id)?.toString(true)
+        val builder = OpenAiStreamingChatModel.builder()
+            .apiKey(token ?: "")
+            .modelName(client.modelId)
+            .temperature(client.temperature.toDouble())
+            .timeout(Duration.ofSeconds(client.timeout.toLong()))
+            .baseUrl(client.host)
+
+        client.proxyUrl?.takeIf { it.isNotBlank() }?.let {
+            val uri = URI(it)
+            builder.proxy(Proxy(Proxy.Type.HTTP, InetSocketAddress(uri.host, uri.port)))
+        }
+
+        client.organizationId?.takeIf { it.isNotBlank() }?.let {
+            builder.organizationId(it)
+        }
+
+        return builder.build()
+    }
+
     fun saveToken(client: OpenAiClientConfiguration, token: String) {
         cs.launch(Dispatchers.Default) {
             try {
diff --git a/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/qianfan/QianfanClientService.kt b/src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/clients/qianfan/QianfanClientService.kt
@@ -10,8 +10,10 @@ 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.chat.StreamingChatLanguageModel
 import dev.langchain4j.model.qianfan.QianfanChatModel
 import dev.langchain4j.model.qianfan.QianfanChatModelNameEnum
+import dev.langchain4j.model.qianfan.QianfanStreamingChatModel
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -42,6 +44,24 @@ class QianfanClientService(private val cs: CoroutineScope) : LLMClientService<Qi
         return builder.build()
     }
 
+    override suspend fun buildStreamingChatModel(client: QianfanClientConfiguration): StreamingChatLanguageModel {
+        val apiKey = client.apiKey.nullize(true) ?: retrieveToken(client.id + "apiKey")?.toString(true)
+        val secretKey = client.secretKey.nullize(true) ?: retrieveToken(client.id + "secretKey")?.toString(true)
+
+        val builder = QianfanStreamingChatModel.builder()
+            .baseUrl(client.host)
+            .apiKey(apiKey)
+            .secretKey(secretKey)
+            .modelName(client.modelId)
+            .temperature(client.temperature.toDouble())
+        // Fix https://github.com/langchain4j/langchain4j/pull/1426. Remove this 'if' statement when langchain4j releases a new version that resolves this issue.
+        if (client.modelId == QianfanChatModelNameEnum.ERNIE_SPEED_128K.modelName) {
+            builder.endpoint("ernie-speed-128k")
+        }
+
+        return builder.build()
+    }
+
     fun saveApiKey(client: QianfanClientConfiguration, key: String) {
         cs.launch(Dispatchers.Default) {
             try {
diff --git a/src/main/resources/messages/AiCommitsBundle.properties b/src/main/resources/messages/AiCommitsBundle.properties
@@ -76,6 +76,9 @@ settings.llmClient.timeout=Timeout
 settings.llmClient.temperature=Temperature
 settings.llmClient.temperature.comment=What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make \
 the output more random,while lower values like 0.2 will make it more focused and deterministic.
+settings.llmClient.streamingResponse=Streaming response
+settings.llmClient.streamingResponse.contextHelp=Some models do not support streaming response and will fall back to normal response.
+
 
 settings.openAI.token.example=sk-ABCdefgHIjKlxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
 settings.openAi.token.comment=You can get your token <a href="https://platform.openai.com/account/api-keys">here.</a>

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,8 @@ class HuggingFaceClientService(private val cs: CoroutineScope) : LLMClientServic`
`37`	`37`	`.build()`
`38`	`38`	`}`
`39`	`39`
	`40`	`+ override suspend fun buildStreamingChatModel(client: HuggingFaceClientConfiguration) = null`
	`41`	`+`
`40`	`42`	`fun saveToken(client: HuggingFaceClientConfiguration, token: String) {`
`41`	`43`	`cs.launch(Dispatchers.Default) {`
`42`	`44`	`try {`