From 8b219ab17ef6377334ff93c4c357bdf92a3a48d0 Mon Sep 17 00:00:00 2001
From: twiceYuan <twiceyuan@gmail.com>
Date: Wed, 17 Jan 2024 17:07:51 +0800
Subject: [PATCH] Replace response handling for chat in the sample with
 streaming content generation

---
 .../ai/sample/feature/chat/ChatUiState.kt     |  7 ++++
 .../ai/sample/feature/chat/ChatViewModel.kt   | 37 ++++++++++++++-----
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/generativeai-android-sample/app/src/main/kotlin/com/google/ai/sample/feature/chat/ChatUiState.kt b/generativeai-android-sample/app/src/main/kotlin/com/google/ai/sample/feature/chat/ChatUiState.kt
index 48dea90e..2ac5f0ff 100644
--- a/generativeai-android-sample/app/src/main/kotlin/com/google/ai/sample/feature/chat/ChatUiState.kt
+++ b/generativeai-android-sample/app/src/main/kotlin/com/google/ai/sample/feature/chat/ChatUiState.kt
@@ -36,4 +36,11 @@ class ChatUiState(
             _messages.add(newMessage)
         }
     }
+
+    fun updateMessage(id: String, newMessage: ChatMessage) {
+        val index = _messages.indexOfFirst { it.id == id }
+        if (index != -1) {
+            _messages[index] = newMessage
+        }
+    }
 }
diff --git a/generativeai-android-sample/app/src/main/kotlin/com/google/ai/sample/feature/chat/ChatViewModel.kt b/generativeai-android-sample/app/src/main/kotlin/com/google/ai/sample/feature/chat/ChatViewModel.kt
index f102678d..d56d5e49 100644
--- a/generativeai-android-sample/app/src/main/kotlin/com/google/ai/sample/feature/chat/ChatViewModel.kt
+++ b/generativeai-android-sample/app/src/main/kotlin/com/google/ai/sample/feature/chat/ChatViewModel.kt
@@ -24,6 +24,7 @@ import com.google.ai.client.generativeai.type.content
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.launch
 
 class ChatViewModel(
@@ -61,18 +62,36 @@ class ChatViewModel(
 
         viewModelScope.launch {
             try {
-                val response = chat.sendMessage(userMessage)
+                val responseFlow = chat.sendMessageStream(userMessage)
 
-                _uiState.value.replaceLastPendingMessage()
+                var outputContent = ""
+                var messageId: String? = null
 
-                response.text?.let { modelResponse ->
-                    _uiState.value.addMessage(
-                        ChatMessage(
-                            text = modelResponse,
-                            participant = Participant.MODEL,
-                            isPending = false
+                responseFlow.collect { response ->
+                    outputContent += response.text
+                    val id = messageId
+                    if (id == null) {
+                        _uiState.value.replaceLastPendingMessage()
+                        _uiState.value.addMessage(
+                            ChatMessage(
+                                text = outputContent,
+                                participant = Participant.MODEL,
+                                isPending = false
+                            ).also {
+                                messageId = it.id
+                            }
                         )
-                    )
+                    } else {
+                        _uiState.value.updateMessage(
+                            id = id,
+                            newMessage = ChatMessage(
+                                id = id,
+                                text = outputContent,
+                                participant = Participant.MODEL,
+                                isPending = false
+                            )
+                        )
+                    }
                 }
             } catch (e: Exception) {
                 _uiState.value.replaceLastPendingMessage()