chat: add button to stop response generation

shubham0204 · Dec 2, 2024 · 4da5ce5 · 4da5ce5
1 parent c4da126
commit 4da5ce5
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1 +1,4 @@
-This is the first alpha release of SmolChat.
+This release:
+
+- adds a button to stop response generation
+- disables chat options when no chat is selected 
diff --git a/app/src/main/java/io/shubham0204/smollmandroid/ui/screens/chat/ChatActivity.kt b/app/src/main/java/io/shubham0204/smollmandroid/ui/screens/chat/ChatActivity.kt
@@ -53,6 +53,7 @@ import androidx.compose.material.icons.automirrored.filled.ArrowForward
 import androidx.compose.material.icons.filled.Android
 import androidx.compose.material.icons.filled.Menu
 import androidx.compose.material.icons.filled.MoreVert
+import androidx.compose.material.icons.filled.Stop
 import androidx.compose.material3.CircularProgressIndicator
 import androidx.compose.material3.DrawerValue
 import androidx.compose.material3.ExperimentalMaterial3Api
@@ -448,7 +449,14 @@ private fun MessageInput(viewModel: ChatScreenViewModel) {
             )
             Spacer(modifier = Modifier.width(8.dp))
             if (isGeneratingResponse || isInitializingModel) {
-                CircularProgressIndicator(color = AppAccentColor)
+                Box(contentAlignment = Alignment.Center) {
+                    CircularProgressIndicator(color = AppAccentColor)
+                    if (isGeneratingResponse) {
+                        IconButton(onClick = { viewModel.stopGeneration() }) {
+                            Icon(Icons.Default.Stop, contentDescription = "Stop", tint = Color.DarkGray)
+                        }
+                    }
+                }
             } else {
                 IconButton(
                     enabled = questionText.isNotEmpty(),

diff --git a/app/src/main/java/io/shubham0204/smollmandroid/ui/screens/chat/ChatScreenViewModel.kt b/app/src/main/java/io/shubham0204/smollmandroid/ui/screens/chat/ChatScreenViewModel.kt
@@ -41,6 +41,7 @@ import io.shubham0204.smollmandroid.llm.ModelsRepository
 import io.shubham0204.smollmandroid.prism4j.PrismGrammarLocator
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
@@ -70,6 +71,7 @@ class ChatScreenViewModel(
     val showTaskListBottomListState = mutableStateOf(false)
 
     val isInitializingModel = mutableStateOf(false)
+    var responseGenerationJob: Job? = null
 
     val markwon: Markwon
 
@@ -134,20 +136,32 @@ class ChatScreenViewModel(
             }
             messagesDB.addUserMessage(chat.id, query)
             isGeneratingResponse.value = true
-            CoroutineScope(Dispatchers.Default).launch {
-                partialResponse.value = ""
-                smolLM.getResponse(query).collect { partialResponse.value += it }
-                messagesDB.addAssistantMessage(chat.id, partialResponse.value)
-                withContext(Dispatchers.Main) { isGeneratingResponse.value = false }
+            responseGenerationJob =
+                CoroutineScope(Dispatchers.Default).launch {
+                    partialResponse.value = ""
+                    smolLM.getResponse(query).collect { partialResponse.value += it }
+                    messagesDB.addAssistantMessage(chat.id, partialResponse.value)
+                    withContext(Dispatchers.Main) { isGeneratingResponse.value = false }
+                }
+        }
+    }
+
+    fun stopGeneration() {
+        isGeneratingResponse.value = false
+        responseGenerationJob?.let { job ->
+            if (job.isActive) {
+                job.cancel()
             }
         }
     }
 
     fun switchChat(chat: Chat) {
+        stopGeneration()
         currChatState.value = chat
     }
 
     fun deleteChat(chat: Chat) {
+        stopGeneration()
         chatsDB.deleteChat(chat)
         messagesDB.deleteMessages(chat.id)
         currChatState.value = null