fix: add lost messages for Azure

unit-mesh · Dec 31, 2023 · 468a5ce · 468a5ce
1 parent 2fcab1c
commit 468a5ce
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/main/kotlin/cc/unitmesh/devti/llms/azure/AzureOpenAIProvider.kt b/src/main/kotlin/cc/unitmesh/devti/llms/azure/AzureOpenAIProvider.kt
@@ -173,6 +173,7 @@ class AzureOpenAIProvider(val project: Project) : LLMProvider {
                         ObjectMapper().readValue(sse!!.data, ChatCompletionResult::class.java)
                     val completion = result.choices[0].message
                     if (completion != null && completion.content != null) {
+                        output += completion.content
                         trySend(completion.content)
                     }
                 }

diff --git a/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomLLMProvider.kt b/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomLLMProvider.kt
@@ -52,8 +52,7 @@ class CustomLLMProvider(val project: Project) : LLMProvider {
 
     private var client = OkHttpClient()
     private val timeout = Duration.ofSeconds(600)
-    private val messages: MutableList<Message> = ArrayList()
-
+    private val messages: MutableList<Message> = mutableListOf()
     private val logger = logger<CustomLLMProvider>()
 
     override fun clearMessage() {
@@ -139,6 +138,7 @@ class CustomLLMProvider(val project: Project) : LLMProvider {
                                 }
                             }
                         }
+
                     close()
                 }
                 awaitClose()