UI: polish conversation screen

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

@@ -86,7 +86,8 @@ class InferenceEngine {
         // This would be replaced with actual token generation logic
         return flow {
             try {
-                delay(500) // Simulate processing time
+                // Simulate longer processing time (1.5 seconds)
+                delay(1500)
 
                 _state.value = State.Generating
 
@@ -96,7 +97,8 @@ class InferenceEngine {
 
                 for (word in words) {
                     emit(word + " ")
-                    delay(50) // Simulate token generation delay
+                    // Slower token generation (200ms per token instead of 50ms)
+                    delay(200)
                 }
 
                 _state.value = State.AwaitingUserPrompt

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

@@ -1,6 +1,16 @@
 package com.example.llama.revamp.ui.screens
 
 import androidx.compose.animation.AnimatedVisibility
+import androidx.compose.animation.core.LinearEasing
+import androidx.compose.animation.core.RepeatMode
+import androidx.compose.animation.core.animateFloat
+import androidx.compose.animation.core.infiniteRepeatable
+import androidx.compose.animation.core.rememberInfiniteTransition
+import androidx.compose.animation.core.tween
+import androidx.compose.animation.expandVertically
+import androidx.compose.animation.fadeIn
+import androidx.compose.animation.fadeOut
+import androidx.compose.animation.shrinkVertically
 import androidx.compose.foundation.background
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
@@ -13,48 +23,61 @@ import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.size
 import androidx.compose.foundation.layout.width
 import androidx.compose.foundation.lazy.LazyColumn
+import androidx.compose.foundation.lazy.LazyListState
 import androidx.compose.foundation.lazy.items
 import androidx.compose.foundation.lazy.rememberLazyListState
 import androidx.compose.foundation.shape.CircleShape
 import androidx.compose.foundation.shape.RoundedCornerShape
 import androidx.compose.material.icons.Icons
-import androidx.compose.material.icons.filled.ExpandLess
-import androidx.compose.material.icons.filled.ExpandMore
 import androidx.compose.material.icons.filled.Send
 import androidx.compose.material3.Card
+import androidx.compose.material3.CardDefaults
 import androidx.compose.material3.CircularProgressIndicator
 import androidx.compose.material3.DrawerState
-import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.material3.Icon
 import androidx.compose.material3.IconButton
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.OutlinedTextField
+import androidx.compose.material3.Surface
 import androidx.compose.material3.Text
+import androidx.compose.material3.TextFieldDefaults
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.collectAsState
+import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.clip
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.graphics.StrokeCap
+import androidx.compose.ui.platform.LocalLifecycleOwner
+import androidx.compose.ui.text.style.TextOverflow
 import androidx.compose.ui.unit.dp
-import androidx.lifecycle.viewmodel.compose.viewModel
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleEventObserver
+import androidx.lifecycle.repeatOnLifecycle
 import com.example.llama.revamp.engine.InferenceEngine
 import com.example.llama.revamp.navigation.NavigationActions
 import com.example.llama.revamp.ui.components.AppScaffold
 import com.example.llama.revamp.viewmodel.MainViewModel
 import com.example.llama.revamp.viewmodel.Message
+import kotlinx.coroutines.launch
 
-@OptIn(ExperimentalMaterial3Api::class)
+/**
+ * Screen for LLM conversation with user.
+ */
 @Composable
 fun ConversationScreen(
     onBackPressed: () -> Unit,
     drawerState: DrawerState,
     navigationActions: NavigationActions,
-    viewModel: MainViewModel = viewModel()
+    viewModel: MainViewModel
 ) {
     val engineState by viewModel.engineState.collectAsState()
     val messages by viewModel.messages.collectAsState()
@@ -64,13 +87,38 @@ fun ConversationScreen(
     val isProcessing = engineState is InferenceEngine.State.ProcessingUserPrompt
     val isGenerating = engineState is InferenceEngine.State.Generating
 
-    val lazyListState = rememberLazyListState()
+    val listState = rememberLazyListState()
     var inputText by remember { mutableStateOf("") }
+    val coroutineScope = rememberCoroutineScope()
+    val lifecycleOwner = LocalLifecycleOwner.current
 
-    // Auto-scroll to bottom when messages change
+    // Auto-scroll to bottom when messages change or when typing
-    LaunchedEffect(messages.size) {
+    val shouldScrollToBottom by remember(messages.size, isGenerating) {
+        derivedStateOf { true }
+    }
+
+    LaunchedEffect(shouldScrollToBottom, messages.size) {
         if (messages.isNotEmpty()) {
-            lazyListState.animateScrollToItem(messages.size - 1)
+            listState.animateScrollToItem(messages.size - 1)
+        }
+    }
+
+    // Set up lifecycle-aware message monitoring
+    DisposableEffect(lifecycleOwner) {
+        val observer = LifecycleEventObserver { _, event ->
+            if (event == Lifecycle.Event.ON_RESUME) {
+                // Scroll to bottom when returning to the screen
+                if (messages.isNotEmpty()) {
+                    coroutineScope.launch {
+                        listState.scrollToItem(messages.size - 1)
+                    }
+                }
+            }
+        }
+
+        lifecycleOwner.lifecycle.addObserver(observer)
+        onDispose {
+            lifecycleOwner.lifecycle.removeObserver(observer)
         }
     }
 
@@ -86,215 +134,187 @@ fun ConversationScreen(
                 .padding(paddingValues)
         ) {
             // System prompt display (collapsible)
-            systemPrompt?.let { prompt ->
+            AnimatedSystemPrompt(systemPrompt)
-                var expanded by remember { mutableStateOf(false) }
 
-                Card(
+            // Messages list
-                    modifier = Modifier
+            Box(
-                        .fillMaxWidth()
-                        .padding(horizontal = 16.dp, vertical = 8.dp)
-                ) {
-                    Column(
-                        modifier = Modifier.padding(8.dp)
-                    ) {
-                        Row(
-                            modifier = Modifier
-                                .fillMaxWidth()
-                                .padding(8.dp),
-                            verticalAlignment = Alignment.CenterVertically
-                        ) {
-                            Text(
-                                text = "System Prompt",
-                                style = MaterialTheme.typography.titleMedium
-                            )
-
-                            Spacer(modifier = Modifier.weight(1f))
-
-                            IconButton(onClick = { expanded = !expanded }) {
-                                Icon(
-                                    imageVector = if (expanded) Icons.Filled.ExpandLess else Icons.Filled.ExpandMore,
-                                    contentDescription = if (expanded) "Collapse" else "Expand"
-                                )
-                            }
-                        }
-
-                        AnimatedVisibility(visible = expanded) {
-                            Text(
-                                text = prompt,
-                                style = MaterialTheme.typography.bodySmall,
-                                modifier = Modifier.padding(8.dp)
-                            )
-                        }
-                    }
-                }
-            }
-
-            // Messages
-            LazyColumn(
-                state = lazyListState,
                 modifier = Modifier
                     .weight(1f)
                     .fillMaxWidth()
-                    .padding(horizontal = 16.dp)
             ) {
-                items(messages) { message ->
+                ConversationMessageList(
-                    MessageBubble(message = message)
+                    messages = messages,
-                }
+                    listState = listState,
-
-                // Show thinking indicator when processing
-                item {
-                    if (isProcessing) {
-                        Row(
-                            modifier = Modifier
-                                .fillMaxWidth()
-                                .padding(vertical = 8.dp),
-                            verticalAlignment = Alignment.CenterVertically
-                        ) {
-                            // Assistant avatar
-                            Box(
-                                modifier = Modifier
-                                    .size(40.dp)
-                                    .clip(CircleShape)
-                                    .background(MaterialTheme.colorScheme.primary),
-                                contentAlignment = Alignment.Center
-                            ) {
-                                Text(
-                                    text = "AI",
-                                    color = MaterialTheme.colorScheme.onPrimary
                 )
             }
 
-                            Spacer(modifier = Modifier.width(8.dp))
-
-                            CircularProgressIndicator(
-                                modifier = Modifier.size(16.dp),
-                                strokeWidth = 2.dp
-                            )
-
-                            Spacer(modifier = Modifier.width(8.dp))
-
-                            Text(
-                                text = "Thinking...",
-                                style = MaterialTheme.typography.bodyMedium,
-                                color = MaterialTheme.colorScheme.onSurfaceVariant
-                            )
-                        }
-                    }
-                }
-            }
-
             // Input area
-            Card(
+            ConversationInputField(
-                modifier = Modifier
-                    .fillMaxWidth()
-                    .padding(16.dp)
-            ) {
-                Row(
-                    modifier = Modifier
-                        .fillMaxWidth()
-                        .padding(8.dp),
-                    verticalAlignment = Alignment.CenterVertically
-                ) {
-                    OutlinedTextField(
                 value = inputText,
                 onValueChange = { inputText = it },
-                        modifier = Modifier.weight(1f),
+                onSendClick = {
-                        placeholder = { Text("Type your message...") },
-                        singleLine = false,
-                        maxLines = 5,
-                        enabled = !isProcessing && !isGenerating
-                    )
-
-                    Spacer(modifier = Modifier.width(8.dp))
-
-                    IconButton(
-                        onClick = {
                     if (inputText.isNotBlank()) {
                         viewModel.sendMessage(inputText)
                         inputText = ""
                     }
                 },
-                        enabled = inputText.isNotBlank() && !isProcessing && !isGenerating
+                isEnabled = !isProcessing && !isGenerating
-                    ) {
-                        Icon(
-                            imageVector = Icons.Default.Send,
-                            contentDescription = "Send",
-                            tint = if (inputText.isNotBlank() && !isProcessing && !isGenerating)
-                                MaterialTheme.colorScheme.primary
-                            else
-                                MaterialTheme.colorScheme.onSurface.copy(alpha = 0.38f)
             )
         }
     }
+}
+
+@Composable
+fun AnimatedSystemPrompt(systemPrompt: String?) {
+    var expanded by remember { mutableStateOf(false) }
+
+    if (!systemPrompt.isNullOrBlank()) {
+        Card(
+            modifier = Modifier
+                .fillMaxWidth()
+                .padding(horizontal = 16.dp, vertical = 8.dp),
+            onClick = { expanded = !expanded }
+        ) {
+            Column(
+                modifier = Modifier.padding(16.dp)
+            ) {
+                Row(
+                    modifier = Modifier.fillMaxWidth(),
+                    verticalAlignment = Alignment.CenterVertically
+                ) {
+                    Text(
+                        text = "System Prompt",
+                        style = MaterialTheme.typography.titleMedium,
+                        modifier = Modifier.weight(1f)
+                    )
+
+                    Text(
+                        text = if (expanded) "Hide" else "Show",
+                        style = MaterialTheme.typography.labelMedium,
+                        color = MaterialTheme.colorScheme.primary
+                    )
+                }
+
+                AnimatedVisibility(
+                    visible = expanded,
+                    enter = fadeIn() + expandVertically(),
+                    exit = fadeOut() + shrinkVertically()
+                ) {
+                    Text(
+                        text = systemPrompt,
+                        style = MaterialTheme.typography.bodySmall,
+                        color = MaterialTheme.colorScheme.onSurfaceVariant,
+                        modifier = Modifier
+                            .fillMaxWidth()
+                            .padding(top = 8.dp)
+                    )
+                }
             }
         }
     }
 }
 
+@Composable
+fun ConversationMessageList(
+    messages: List<Message>,
+    listState: LazyListState,
+) {
+    LazyColumn(
+        state = listState,
+        modifier = Modifier
+            .fillMaxSize()
+            .padding(horizontal = 16.dp),
+        reverseLayout = false
+    ) {
+        items(
+            items = messages,
+            key = { "${it::class.simpleName}_${it.timestamp}" }
+        ) { message ->
+            MessageBubble(message = message)
+        }
+
+        // Add extra space at the bottom for better UX
+        item { Spacer(modifier = Modifier.height(16.dp)) }
+    }
+}
+
 @Composable
 fun MessageBubble(message: Message) {
-    Row(
+    when (message) {
+        is Message.User -> UserMessageBubble(
+            content = message.content,
+            formattedTime = message.formattedTime
+        )
+
+        is Message.Assistant -> AssistantMessageBubble(
+            content = message.content,
+            formattedTime = message.formattedTime,
+            isComplete = message.isComplete,
+            isThinking = !message.isComplete && message.content.isBlank(),
+            metrics = if (message.isComplete && message.content.isNotBlank()) {
+                // TODO-han.yin: Generate some example metrics for now
+                // This would come from the actual LLM engine in a real implementation
+                val tokenCount = message.content.split("\\s+".toRegex()).size
+                val ttft = (200 + (Math.random() * 80)).toInt()
+                val tps = 8.5 + (Math.random() * 1.5)
+                "Tokens: $tokenCount, TTFT: ${ttft}ms, TPS: ${"%.1f".format(tps)}"
+            } else null
+        )
+    }
+}
+
+@Composable
+fun UserMessageBubble(content: String, formattedTime: String) {
+    Column(
         modifier = Modifier
             .fillMaxWidth()
             .padding(vertical = 8.dp),
-        verticalAlignment = Alignment.Top
+        horizontalAlignment = Alignment.End
     ) {
-        when (message) {
+        // Timestamp above bubble
-            is Message.User -> {
+        Text(
+            text = formattedTime,
+            style = MaterialTheme.typography.labelSmall,
+            color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
+            modifier = Modifier.padding(bottom = 4.dp)
+        )
+
+        Row {
             Spacer(modifier = Modifier.weight(1f))
 
-                Box(
+            Card(
-                    modifier = Modifier
+                shape = RoundedCornerShape(16.dp, 4.dp, 16.dp, 16.dp),
-                        .width(40.dp)
+                colors = CardDefaults.cardColors(
-                        .height(40.dp)
+                    containerColor = MaterialTheme.colorScheme.primaryContainer
-                        .clip(CircleShape)
+                ),
-                        .background(MaterialTheme.colorScheme.secondaryContainer),
+                elevation = CardDefaults.cardElevation(defaultElevation = 1.dp)
-                    contentAlignment = Alignment.Center
             ) {
                 Text(
-                        text = "You",
+                    text = content,
-                        color = MaterialTheme.colorScheme.onSecondaryContainer,
-                        style = MaterialTheme.typography.labelMedium
-                    )
-                }
-
-                Spacer(modifier = Modifier.width(8.dp))
-
-                Box(
-                    modifier = Modifier
-                        .weight(5f)
-                        .clip(
-                            RoundedCornerShape(
-                                topStart = 16.dp,
-                                topEnd = 16.dp,
-                                bottomStart = 16.dp,
-                                bottomEnd = 4.dp
-                            )
-                        )
-                        .background(MaterialTheme.colorScheme.primaryContainer)
-                        .padding(12.dp)
-                ) {
-                    Column {
-                        Text(
-                            text = message.content,
                     style = MaterialTheme.typography.bodyMedium,
-                            color = MaterialTheme.colorScheme.onPrimaryContainer
+                    color = MaterialTheme.colorScheme.onPrimaryContainer,
-                        )
+                    modifier = Modifier.padding(12.dp)
-
-                        Text(
-                            text = message.formattedTime,
-                            style = MaterialTheme.typography.labelSmall,
-                            color = MaterialTheme.colorScheme.onPrimaryContainer.copy(alpha = 0.7f),
-                            modifier = Modifier.align(Alignment.End)
                 )
             }
         }
     }
+}
 
-            is Message.Assistant -> {
+@Composable
+fun AssistantMessageBubble(
+    content: String,
+    formattedTime: String,
+    isComplete: Boolean,
+    isThinking: Boolean,
+    metrics: String? = null
+) {
+    Row(
+        verticalAlignment = Alignment.Top,
+    ) {
+        // Assistant avatar
         Box(
             modifier = Modifier
-                        .size(40.dp)
+                .size(36.dp)
                 .clip(CircleShape)
                 .background(MaterialTheme.colorScheme.primary),
             contentAlignment = Alignment.Center
@@ -308,47 +328,163 @@ fun MessageBubble(message: Message) {
 
         Spacer(modifier = Modifier.width(8.dp))
 
-                Box(
+        Column(
             modifier = Modifier
-                        .weight(5f)
+                .fillMaxWidth()
-                        .clip(
-                            RoundedCornerShape(
-                                topStart = 16.dp,
-                                topEnd = 16.dp,
-                                bottomStart = 4.dp,
-                                bottomEnd = 16.dp
-                            )
-                        )
-                        .background(MaterialTheme.colorScheme.surfaceVariant)
-                        .padding(12.dp)
         ) {
-                    Column {
+            // Timestamp above bubble
+            if (formattedTime.isNotBlank()) {
                 Text(
-                            text = message.content,
+                    text = formattedTime,
+                    style = MaterialTheme.typography.labelSmall,
+                    color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
+                    modifier = Modifier.padding(bottom = 4.dp)
+                )
+            }
+
+            Card(
+                shape = RoundedCornerShape(4.dp, 16.dp, 16.dp, 16.dp),
+                colors = CardDefaults.cardColors(
+                    containerColor = MaterialTheme.colorScheme.surfaceVariant
+                ),
+                elevation = CardDefaults.cardElevation(defaultElevation = 1.dp)
+            ) {
+                // Show actual content
+                Text(
+                    modifier = Modifier.padding(12.dp),
+                    text = content,
                     style = MaterialTheme.typography.bodyMedium,
                     color = MaterialTheme.colorScheme.onSurfaceVariant
                 )
+            }
+
+            // Show metrics or generation status below the bubble
+            Row(
+                modifier = Modifier.height(20.dp).padding(top = 4.dp),
+                verticalAlignment = Alignment.CenterVertically
+            ) {
+                if (!isComplete) {
+                    PulsatingDots(small = true)
+
+                    Spacer(modifier = Modifier.width(4.dp))
 
-                        if (message.isComplete) {
                     Text(
-                                text = message.formattedTime,
+                        text = if (isThinking) "Thinking..." else "",
+                        style = MaterialTheme.typography.labelSmall,
+                        color = MaterialTheme.colorScheme.primary
+                    )
+                } else if (metrics != null) {
+                    // Show metrics when message is complete
+                    Text(
+                        text = metrics,
                         style = MaterialTheme.typography.labelSmall,
                         color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.7f),
-                                modifier = Modifier.align(Alignment.End)
-                            )
-                        } else {
-                            CircularProgressIndicator(
-                                modifier = Modifier
-                                    .size(12.dp)
-                                    .align(Alignment.End),
-                                strokeWidth = 2.dp
                     )
                 }
             }
         }
-
+    }
-                Spacer(modifier = Modifier.weight(1f))
+}
-            }
+
-        }
+@Composable
+fun PulsatingDots(small: Boolean = false) {
+    val transition = rememberInfiniteTransition(label = "dots")
+
+    val animations = List(3) { index ->
+        transition.animateFloat(
+            initialValue = 0f,
+            targetValue = 1f,
+            animationSpec = infiniteRepeatable(
+                animation = tween(
+                    durationMillis = 1000,
+                    delayMillis = index * 300,
+                    easing = LinearEasing
+                ),
+                repeatMode = RepeatMode.Reverse
+            ),
+            label = "dot-$index"
+        )
+    }
+
+    Row(verticalAlignment = Alignment.CenterVertically) {
+        animations.forEach { animation ->
+            Spacer(modifier = Modifier.width(2.dp))
+
+            Box(
+                modifier = Modifier
+                    .size(if (small) 5.dp else 8.dp)
+                    .clip(CircleShape)
+                    .background(
+                        color = MaterialTheme.colorScheme.primary.copy(
+                            alpha = 0.3f + (animation.value * 0.7f)
+                        )
+                    )
+            )
+
+            Spacer(modifier = Modifier.width(2.dp))
+        }
+    }
+}
+
+@Composable
+fun ConversationInputField(
+    value: String,
+    onValueChange: (String) -> Unit,
+    onSendClick: () -> Unit,
+    isEnabled: Boolean
+) {
+    Surface(
+        modifier = Modifier
+            .fillMaxWidth(),
+        shadowElevation = 4.dp,
+        color = MaterialTheme.colorScheme.surface
+    ) {
+        Row(
+            modifier = Modifier
+                .fillMaxWidth()
+                .padding(horizontal = 16.dp, vertical = 8.dp),
+            verticalAlignment = Alignment.Bottom
+        ) {
+            OutlinedTextField(
+                value = value,
+                onValueChange = onValueChange,
+                modifier = Modifier
+                    .weight(1f)
+                    .padding(end = 8.dp),
+                placeholder = { Text("Message Kleidi LLaMA...") },
+                maxLines = 5,
+                enabled = isEnabled,
+                colors = TextFieldDefaults.colors(
+                    unfocusedContainerColor = Color.Transparent,
+                    focusedContainerColor = Color.Transparent
+                ),
+                shape = RoundedCornerShape(24.dp)
+            )
+
+            IconButton(
+                onClick = onSendClick,
+                enabled = value.isNotBlank() && isEnabled,
+                modifier = Modifier
+                    .padding(bottom = 4.dp)
+                    .size(48.dp)
+            ) {
+                if (isEnabled) {
+                    Icon(
+                        imageVector = Icons.Default.Send,
+                        contentDescription = "Send message",
+                        tint = if (value.isNotBlank())
+                            MaterialTheme.colorScheme.primary
+                        else
+                            MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.5f)
+                    )
+                } else {
+                    CircularProgressIndicator(
+                        modifier = Modifier.size(24.dp),
+                        strokeWidth = 2.dp,
+                        strokeCap = StrokeCap.Round
+                    )
+                }
+            }
+        }
     }
 }

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

@@ -252,6 +252,7 @@ sealed class Message {
         override val timestamp: Long
     ) : Message()
 
+    // TODO-han.yin: break down into ongoing & completed message subtypes
     data class Assistant(
         override val content: String,
         override val timestamp: Long,