Refactor Gemini Chatbot sample to use UI State and handle errors (#54)

ksemenova · web-flow · commit e092b9b38ebc · 2025-07-09T15:30:38.000-07:00
Co-authored-by: ksemenova &lt;10581063+ksemenova@users.noreply.github.com&gt;
diff --git a/ai-catalog/gradle/libs.versions.toml b/ai-catalog/gradle/libs.versions.toml
@@ -2,6 +2,7 @@
 agp = "8.8.2"
 coilCompose = "3.1.0"
 firebaseBom = "33.14.0"
+lifecycleRuntimeCompose = "2.9.1"
 mlkitGenAi = "1.0.0-beta1"
 kotlin = "2.1.0"
 coreKtx = "1.15.0"
@@ -31,6 +32,7 @@ uiTooling = "1.8.3"
 
 [libraries]
 androidx-core-ktx = { group = "androidx.core", name = "core-ktx", version.ref = "coreKtx" }
+androidx-lifecycle-runtime-compose = { module = "androidx.lifecycle:lifecycle-runtime-compose", version.ref = "lifecycleRuntimeCompose" }
 coil-compose = { module = "io.coil-kt.coil3:coil-compose", version.ref = "coilCompose" }
 firebase-bom = { module = "com.google.firebase:firebase-bom", version.ref = "firebaseBom" }
 firebase-ai = { group = "com.google.firebase", name = "firebase-ai" }
diff --git a/ai-catalog/samples/gemini-chatbot/build.gradle.kts b/ai-catalog/samples/gemini-chatbot/build.gradle.kts
@@ -65,6 +65,7 @@ dependencies {
     implementation(libs.hilt.android)
     implementation(libs.hilt.navigation.compose)
     implementation(libs.androidx.runtime.livedata)
+    implementation(libs.androidx.lifecycle.runtime.compose)
     implementation(libs.androidx.material3.android)
     ksp(libs.hilt.compiler)
 
diff --git a/ai-catalog/samples/gemini-chatbot/src/main/java/com/android/ai/samples/geminichatbot/GeminiChatbotScreen.kt b/ai-catalog/samples/gemini-chatbot/src/main/java/com/android/ai/samples/geminichatbot/GeminiChatbotScreen.kt
@@ -16,7 +16,6 @@
 package com.android.ai.samples.geminichatbot
 
 import android.content.Intent
-import android.net.Uri
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.PaddingValues
@@ -32,7 +31,9 @@ import androidx.compose.foundation.lazy.LazyColumn
 import androidx.compose.foundation.lazy.items
 import androidx.compose.material.icons.Icons
 import androidx.compose.material.icons.filled.Code
+import androidx.compose.material3.AlertDialog
 import androidx.compose.material3.Button
+import androidx.compose.material3.CircularProgressIndicator
 import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.material3.Icon
 import androidx.compose.material3.MaterialTheme
@@ -44,7 +45,6 @@ import androidx.compose.material3.TopAppBarDefaults
 import androidx.compose.material3.TopAppBarDefaults.topAppBarColors
 import androidx.compose.material3.rememberTopAppBarState
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.saveable.rememberSaveable
@@ -59,14 +59,16 @@ import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.unit.sp
+import androidx.core.net.toUri
 import androidx.hilt.navigation.compose.hiltViewModel
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
 
 @OptIn(ExperimentalMaterial3Api::class)
 @Composable
 fun GeminiChatbotScreen(viewModel: GeminiChatbotViewModel = hiltViewModel()) {
     val topAppBarState = rememberTopAppBarState()
     val scrollBehavior = TopAppBarDefaults.pinnedScrollBehavior(topAppBarState)
-    val messages by viewModel.messageList.collectAsState()
+    val uiState by viewModel.uiState.collectAsStateWithLifecycle()
     var message by rememberSaveable { mutableStateOf("") }
 
     Scaffold(
@@ -88,15 +90,50 @@ fun GeminiChatbotScreen(viewModel: GeminiChatbotViewModel = hiltViewModel()) {
             )
         },
     ) { innerPadding ->
+
         Column {
             val layoutDirection = LocalLayoutDirection.current
+
+            val messages = when (val state = uiState) {
+                is GeminiChatbotUiState.Initial -> emptyList()
+                is GeminiChatbotUiState.Generating -> state.messages
+                is GeminiChatbotUiState.Success -> state.messages
+                is GeminiChatbotUiState.Error -> state.messages
+            }
+
             MessageList(
                 modifier = Modifier
                     .fillMaxWidth()
                     .weight(1f),
-                messages = messages.sortedBy { - it.timestamp },
+                messages = messages.sortedByDescending { it.timestamp },
                 contentPadding = innerPadding,
             )
+
+            when (val state = uiState) {
+                is GeminiChatbotUiState.Generating -> {
+                    CircularProgressIndicator(
+                        modifier = Modifier
+                            .padding(vertical = 8.dp)
+                            .align(Alignment.CenterHorizontally),
+                    )
+                }
+
+                is GeminiChatbotUiState.Error -> {
+                    AlertDialog(
+                        onDismissRequest = { viewModel.dismissError() },
+                        title = { Text(text = stringResource(R.string.error)) },
+                        text = { Text(text = state.errorMessage) },
+                        confirmButton = {
+                            Button(onClick = { viewModel.dismissError() }) {
+                                Text(text = stringResource(R.string.dismiss_button))
+                            }
+                        },
+                    )
+                }
+
+                else -> { /* No additional UI for Initial or Success states */ }
+            }
+
             InputBar(
                 value = message,
                 placeholder = stringResource(R.string.geminichatbot_input_placeholder),
@@ -108,7 +145,7 @@ fun GeminiChatbotScreen(viewModel: GeminiChatbotViewModel = hiltViewModel()) {
                     message = ""
                 },
                 contentPadding = innerPadding.copy(layoutDirection, top = 0.dp),
-                sendEnabled = true,
+                sendEnabled = uiState !is GeminiChatbotUiState.Generating,
             )
         }
     }
@@ -176,7 +213,7 @@ fun SeeCodeButton() {
 
     Button(
         onClick = {
-            val intent = Intent(Intent.ACTION_VIEW, Uri.parse(githubLink))
+            val intent = Intent(Intent.ACTION_VIEW, githubLink.toUri())
             context.startActivity(intent)
         },
         modifier = Modifier.padding(end = 8.dp),
diff --git a/ai-catalog/samples/gemini-chatbot/src/main/java/com/android/ai/samples/geminichatbot/GeminiChatbotViewModel.kt b/ai-catalog/samples/gemini-chatbot/src/main/java/com/android/ai/samples/geminichatbot/GeminiChatbotViewModel.kt
@@ -28,12 +28,23 @@ import com.google.firebase.ai.type.generationConfig
 import javax.inject.Inject
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.launch
 
+sealed class GeminiChatbotUiState {
+    data object Initial : GeminiChatbotUiState()
+    data class Generating(val messages: List<ChatMessage>) : GeminiChatbotUiState()
+    data class Success(val messages: List<ChatMessage>) : GeminiChatbotUiState()
+    data class Error(
+        val errorMessage: String,
+        val messages: List<ChatMessage>,
+    ) : GeminiChatbotUiState()
+}
+
 class GeminiChatbotViewModel @Inject constructor() : ViewModel() {
 
-    private val _messageList = MutableStateFlow(mutableListOf<ChatMessage>())
-    val messageList: StateFlow<List<ChatMessage>> = _messageList
+    private val _uiState = MutableStateFlow<GeminiChatbotUiState>(GeminiChatbotUiState.Initial)
+    val uiState: StateFlow<GeminiChatbotUiState> = _uiState.asStateFlow()
 
     private val generativeModel by lazy {
         Firebase.ai(backend = GenerativeBackend.googleAI()).generativeModel(
@@ -59,21 +70,24 @@ class GeminiChatbotViewModel @Inject constructor() : ViewModel() {
     private val chat = generativeModel.startChat()
 
     fun sendMessage(message: String) {
-        viewModelScope.launch {
-            _messageList.value.add(
-                ChatMessage(
-                    message,
-                    System.currentTimeMillis(),
-                    false,
-                    null,
-                ),
-            )
+        val currentMessages =
+            when (val state = _uiState.value) {
+                is GeminiChatbotUiState.Success -> state.messages
+                is GeminiChatbotUiState.Error -> state.messages
+                else -> emptyList()
+            }
 
-            val response = chat.sendMessage(message)
+        val newMessages =
+            currentMessages.toMutableList().apply {
+                add(ChatMessage(message, System.currentTimeMillis(), false, null))
+            }
+        viewModelScope.launch {
+            try {
+                _uiState.value = GeminiChatbotUiState.Generating(newMessages)
 
-            response.text?.let {
-                _messageList.value = _messageList.value.toMutableList().apply {
-                    add(
+                val response = chat.sendMessage(message)
+                response.text?.let {
+                    newMessages.add(
                         ChatMessage(
                             it.trim(),
                             System.currentTimeMillis(),
@@ -82,7 +96,19 @@ class GeminiChatbotViewModel @Inject constructor() : ViewModel() {
                         ),
                     )
                 }
+                _uiState.value = GeminiChatbotUiState.Success(newMessages)
+            } catch (e: Exception) {
+                _uiState.value =
+                    GeminiChatbotUiState.Error(
+                        e.localizedMessage ?: "Something went wrong, try again",
+                        newMessages,
+                    )
             }
         }
     }
+
+    fun dismissError() {
+        val errorState = _uiState.value as? GeminiChatbotUiState.Error ?: return
+        _uiState.value = GeminiChatbotUiState.Success(errorState.messages)
+    }
 }
diff --git a/ai-catalog/samples/gemini-chatbot/src/main/res/values/strings.xml b/ai-catalog/samples/gemini-chatbot/src/main/res/values/strings.xml
@@ -1,6 +1,8 @@
 <?xml version="1.0" encoding="utf-8"?>
 <resources>
-    <string name="geminichatbot_title_bar">Gemini chatbot</string>
+    <string name="geminichatbot_title_bar">Gemini Chatbot</string>
     <string name="geminichatbot_input_placeholder">Type your message</string>
     <string name="see_code">See code</string>
+    <string name="dismiss_button">Dismiss</string>
+    <string name="error">Error</string>
 </resources>
