UI: handle back press on Model Selection screen

examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt

@@ -293,6 +293,9 @@ fun AppContent(
                 // Model Selection Screen
                 composable(AppDestinations.MODEL_SELECTION_ROUTE) {
                     ModelSelectionScreen(
+                        onNavigateBack = {
+                            navigationActions.navigateUp()
+                        },
                         onModelConfirmed = { modelInfo ->
                             navigationActions.navigateToModelLoading()
                         },

examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelSelectionScreen.kt

@@ -1,5 +1,6 @@
 package com.example.llama.revamp.ui.screens
 
+import androidx.activity.compose.BackHandler
 import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.animation.fadeIn
 import androidx.compose.animation.fadeOut
@@ -47,9 +48,9 @@ fun ModelSelectionScreen(
     val models by viewModel.availableModels.collectAsState()
     val preselectedModel by viewModel.preselectedModel.collectAsState()
 
-    val handleModelSelection = { model: ModelInfo ->
+    // Handle back button press
-        viewModel.confirmSelectedModel(model)
+    BackHandler(preselectedModel != null) {
-        onModelConfirmed(model)
+        viewModel.onBackPressed()
     }
 
     Box(modifier = Modifier.fillMaxSize()) {
@@ -89,7 +90,12 @@ fun ModelSelectionScreen(
             exit = scaleOut() + fadeOut()
         ) {
             FloatingActionButton(
-                onClick = { preselectedModel?.let { handleModelSelection(it) } },
+                onClick = {
+                    preselectedModel?.let {
+                        viewModel.confirmSelectedModel(it)
+                        onModelConfirmed(it)
+                    }
+                },
                 containerColor = MaterialTheme.colorScheme.primary
             ) {
                 Icon(

examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelSelectionViewModel.kt

@@ -55,6 +55,15 @@ class ModelSelectionViewModel @Inject constructor(
         _preselectedModel.value = null
     }
 
+    /**
+     * Handle back press from both back button and top bar
+     */
+    fun onBackPressed() {
+        if (_preselectedModel.value != null) {
+            resetSelection()
+        }
+    }
+
     companion object {
         private val TAG = ModelSelectionViewModel::class.java.simpleName