Merge origin/main into android/fix-chat-scroll-ui

kirklandsign · kirklandsign · commit c0f24d3c6b8f · 2026-04-09T16:34:52.000-07:00
diff --git a/llm/android/LlamaDemo/app/build.gradle.kts b/llm/android/LlamaDemo/app/build.gradle.kts
@@ -270,6 +270,8 @@ dependencies {
     implementation("androidx.constraintlayout:constraintlayout:2.2.0")
     implementation("com.facebook.fbjni:fbjni:0.5.1")
     implementation("com.google.code.gson:gson:2.8.6")
+    implementation("com.halilibo.compose-richtext:richtext-commonmark:1.0.0-alpha02")
+    implementation("com.halilibo.compose-richtext:richtext-ui-material3:1.0.0-alpha02")
     if (useLocalAar == true) {
         implementation(files("libs/executorch.aar"))
     } else {
diff --git a/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/Message.kt b/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/Message.kt
@@ -46,10 +46,16 @@ class Message(
 
     var totalGenerationTime: Long = 0L
 
+    var thinkingContent: String = ""
+
     fun appendText(text: String) {
         this.text += text
     }
 
+    fun appendThinkingText(text: String) {
+        thinkingContent += text
+    }
+
     /**
      * Creates a new Message instance with the same state.
      * Required for Compose strong skipping mode (default since Kotlin 2.0): composable functions
@@ -61,6 +67,7 @@ class Message(
         return Message(sourceText, isSent, messageType, promptID, timestamp, id).also {
             it.tokensPerSecond = tokensPerSecond
             it.totalGenerationTime = totalGenerationTime
+            it.thinkingContent = thinkingContent
         }
     }
 
diff --git a/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.kt b/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.kt
@@ -64,7 +64,7 @@ object PromptFormat {
     @JvmStatic
     fun getThinkingModeToken(modelType: ModelType, thinkingMode: Boolean): String {
         return when (modelType) {
-            ModelType.QWEN_3 -> if (thinkingMode) "" else "<think>\n\n</think>\n\n\n"
+            ModelType.QWEN_3 -> "" // Always enable thinking for Qwen-3
             else -> ""
         }
     }
@@ -74,8 +74,6 @@ object PromptFormat {
         return when (modelType) {
             ModelType.QWEN_3 -> when (token) {
                 "<|im_end|>" -> ""
-                "<think>" -> "Thinking...\n"
-                "</think>" -> "\nDone thinking"
                 else -> token
             }
             else -> token
diff --git a/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/components/ChatInput.kt b/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/components/ChatInput.kt
@@ -35,7 +35,6 @@ import androidx.compose.foundation.text.BasicTextField
 import androidx.compose.material.icons.Icons
 import androidx.compose.material.icons.filled.Add
 import androidx.compose.material.icons.filled.Close
-import androidx.compose.material.icons.filled.Lightbulb
 import androidx.compose.material.icons.filled.Send
 import androidx.compose.material.icons.filled.Stop
 import androidx.compose.material.icons.outlined.AudioFile
@@ -70,8 +69,6 @@ fun ChatInput(
     onInputTextChange: (String) -> Unit,
     isModelReady: Boolean,
     isGenerating: Boolean,
-    thinkMode: Boolean,
-    onThinkModeToggle: () -> Unit,
     onSendClick: () -> Unit,
     onStopClick: () -> Unit,
     showMediaButtons: Boolean,
@@ -151,18 +148,6 @@ fun ChatInput(
                 }
             }
 
-            // Think mode button
-            IconButton(
-                onClick = onThinkModeToggle,
-                modifier = Modifier.size(40.dp)
-            ) {
-                Icon(
-                    imageVector = Icons.Filled.Lightbulb,
-                    contentDescription = "Think mode",
-                    tint = if (thinkMode) Color(0xFFFFD54F) else appColors.textOnNavBar
-                )
-            }
-
             // Text input
             Box(
                 modifier = Modifier
diff --git a/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/components/MessageItem.kt b/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/components/MessageItem.kt
@@ -9,25 +9,44 @@
 package com.example.executorchllamademo.ui.components
 
 import android.net.Uri
+import androidx.compose.animation.AnimatedVisibility
+import androidx.compose.animation.core.animateFloatAsState
+import androidx.compose.animation.expandVertically
+import androidx.compose.animation.shrinkVertically
 import androidx.compose.foundation.background
+import androidx.compose.foundation.clickable
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.IntrinsicSize
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.fillMaxHeight
 import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.size
 import androidx.compose.foundation.layout.width
 import androidx.compose.foundation.layout.widthIn
 import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.filled.KeyboardArrowDown
+import androidx.compose.material3.Icon
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+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.draw.rotate
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.layout.ContentScale
 import androidx.compose.ui.platform.LocalConfiguration
+import androidx.compose.ui.text.TextStyle
+import androidx.compose.ui.text.font.FontFamily
 import androidx.compose.ui.text.font.FontStyle
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.unit.sp
@@ -36,6 +55,10 @@ import com.example.executorchllamademo.Message
 import com.example.executorchllamademo.MessageType
 import com.example.executorchllamademo.ui.theme.LocalAppColors
 import com.example.executorchllamademo.ui.theme.MessageBubbleSent
+import com.halilibo.richtext.commonmark.Markdown
+import com.halilibo.richtext.ui.CodeBlockStyle
+import com.halilibo.richtext.ui.RichTextStyle
+import com.halilibo.richtext.ui.material3.RichText
 
 @Composable
 fun MessageItem(
@@ -151,12 +174,39 @@ private fun TextMessage(
                 )
                 .padding(horizontal = 12.dp, vertical = 8.dp)
         ) {
-            Text(
-                text = message.text,
-                fontSize = 16.sp,
-                letterSpacing = 0.sp,
-                color = textColor
-            )
+            if (isSent) {
+                // User messages: plain text
+                Text(
+                    text = message.text,
+                    fontSize = 16.sp,
+                    letterSpacing = 0.sp,
+                    color = textColor
+                )
+            } else {
+                // Thinking block (collapsible, shown before response)
+                if (message.thinkingContent.isNotEmpty()) {
+                    ThinkingBlock(
+                        content = message.thinkingContent,
+                        textColor = textColor
+                    )
+                }
+                // Model responses: Markdown rendering
+                if (message.text.isNotEmpty()) {
+                    RichText(
+                        style = RichTextStyle(
+                            codeBlockStyle = CodeBlockStyle(
+                                textStyle = TextStyle(
+                                    fontSize = 14.sp,
+                                    fontFamily = FontFamily.Monospace,
+                                    color = textColor
+                                )
+                            )
+                        )
+                    ) {
+                        Markdown(content = message.text)
+                    }
+                }
+            }
 
             // Show metrics and timestamp on the same row
             val hasMetrics = message.tokensPerSecond > 0 || message.totalGenerationTime > 0
@@ -207,3 +257,75 @@ private fun TextMessage(
         }
     }
 }
+
+@Composable
+private fun ThinkingBlock(
+    content: String,
+    textColor: Color,
+    modifier: Modifier = Modifier
+) {
+    var expanded by remember { mutableStateOf(false) }
+    val rotationAngle by animateFloatAsState(
+        targetValue = if (expanded) 180f else 0f,
+        label = "thinking_expand"
+    )
+
+    Column(modifier = modifier.padding(bottom = 8.dp)) {
+        // Header: clickable to expand/collapse
+        Row(
+            modifier = Modifier
+                .clickable { expanded = !expanded }
+                .padding(vertical = 4.dp),
+            verticalAlignment = Alignment.CenterVertically
+        ) {
+            Icon(
+                imageVector = Icons.Default.KeyboardArrowDown,
+                contentDescription = if (expanded) "Collapse" else "Expand",
+                tint = textColor.copy(alpha = 0.5f),
+                modifier = Modifier
+                    .size(18.dp)
+                    .rotate(rotationAngle)
+            )
+            Spacer(modifier = Modifier.width(4.dp))
+            Text(
+                text = "Thinking",
+                fontSize = 13.sp,
+                fontStyle = FontStyle.Italic,
+                color = textColor.copy(alpha = 0.5f)
+            )
+        }
+
+        // Collapsible thinking content
+        AnimatedVisibility(
+            visible = expanded,
+            enter = expandVertically(),
+            exit = shrinkVertically()
+        ) {
+            Row(modifier = Modifier.height(IntrinsicSize.Min)) {
+                // Left vertical line decoration
+                Box(
+                    modifier = Modifier
+                        .width(2.dp)
+                        .fillMaxHeight()
+                        .background(textColor.copy(alpha = 0.15f))
+                )
+                Spacer(modifier = Modifier.width(8.dp))
+                // Thinking content with Markdown rendering
+                RichText(
+                    modifier = Modifier.weight(1f),
+                    style = RichTextStyle(
+                        codeBlockStyle = CodeBlockStyle(
+                            textStyle = TextStyle(
+                                fontSize = 13.sp,
+                                fontFamily = FontFamily.Monospace,
+                                color = textColor.copy(alpha = 0.7f)
+                            )
+                        )
+                    )
+                ) {
+                    Markdown(content = content)
+                }
+            }
+        }
+    }
+}
diff --git a/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/screens/ChatScreen.kt b/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/screens/ChatScreen.kt
@@ -255,8 +255,6 @@ fun ChatScreen(
                 onInputTextChange = { viewModel.inputText = it },
                 isModelReady = viewModel.isModelReady,
                 isGenerating = viewModel.isGenerating,
-                thinkMode = viewModel.thinkMode,
-                onThinkModeToggle = { viewModel.toggleThinkMode() },
                 onSendClick = { viewModel.sendMessage() },
                 onStopClick = { viewModel.stopGeneration() },
                 showMediaButtons = viewModel.showMediaButtons,
diff --git a/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/viewmodel/ChatViewModel.kt b/llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/viewmodel/ChatViewModel.kt
@@ -60,6 +60,14 @@ class ChatViewModel(application: Application) : AndroidViewModel(application), L
     var supportsImageInput by mutableStateOf(false)
     var supportsAudioInput by mutableStateOf(false)
 
+    // Thinking mode state: tracks whether we're inside a <think>...</think> block
+    private var isInThinkingBlock = false
+
+    // Counter that increments on each token to trigger auto-scroll during generation
+    var scrollTrigger by mutableStateOf(0)
+        private set
+
+
     private val _selectedImages = mutableStateListOf<Uri>()
     val selectedImages: List<Uri> = _selectedImages
 
@@ -639,6 +647,7 @@ class ChatViewModel(application: Application) : AndroidViewModel(application), L
 
         // Create result message placeholder
         resultMessage = Message("", false, MessageType.TEXT, promptID)
+        isInThinkingBlock = false
         _messages.add(resultMessage!!)
 
         // Clear selected images after adding to chat
@@ -752,6 +761,16 @@ class ChatViewModel(application: Application) : AndroidViewModel(application), L
             return
         }
 
+        // Thinking mode state machine: intercept <think> / </think> tags
+        if (processedResult == "<think>") {
+            isInThinkingBlock = true
+            return
+        }
+        if (processedResult == "</think>") {
+            isInThinkingBlock = false
+            return
+        }
+
         processedResult = PromptFormat.replaceSpecialToken(currentSettingsFields.modelType, processedResult)
 
         if (currentSettingsFields.modelType == ModelType.LLAMA_3 &&
@@ -769,19 +788,31 @@ class ChatViewModel(application: Application) : AndroidViewModel(application), L
             return
         }
 
-        val keepResult = !(processedResult == "\n" || processedResult == "\n\n") ||
-                resultMessage?.text?.isNotEmpty() == true
-        if (keepResult) {
-            resultMessage?.appendText(processedResult)
-            // Create a new Message reference to trigger recomposition under Compose strong
-            // skipping mode, which compares unstable parameters by reference equality (===).
-            val index = _messages.indexOfLast { it === resultMessage }
-            if (index >= 0) {
-                val updated = resultMessage!!.copy()
-                _messages[index] = updated
-                resultMessage = updated
+        if (isInThinkingBlock) {
+            // Skip leading newlines in thinking content
+            val keepThinking = !(processedResult == "\n" || processedResult == "\n\n") ||
+                    resultMessage?.thinkingContent?.isNotEmpty() == true
+            if (keepThinking) {
+                resultMessage?.appendThinkingText(processedResult)
             }
+        } else {
+            val keepResult = !(processedResult == "\n" || processedResult == "\n\n") ||
+                    resultMessage?.text?.isNotEmpty() == true
+            if (keepResult) {
+                resultMessage?.appendText(processedResult)
+            }
+        }
+
+        // Create a new Message reference to trigger recomposition under Compose strong
+        // skipping mode, which compares unstable parameters by reference equality (===).
+        val index = _messages.indexOfLast { it === resultMessage }
+        if (index >= 0) {
+            val updated = resultMessage!!.copy()
+            _messages[index] = updated
+            resultMessage = updated
         }
+        // Increment scroll trigger to auto-scroll during generation
+        scrollTrigger++
     }
 
     override fun onStats(stats: String) {

Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ object PromptFormat {`
`64`	`64`	`@JvmStatic`
`65`	`65`	`fun getThinkingModeToken(modelType: ModelType, thinkingMode: Boolean): String {`
`66`	`66`	`return when (modelType) {`
`67`		`- ModelType.QWEN_3 -> if (thinkingMode) "" else "<think>\n\n</think>\n\n\n"`
	`67`	`+ ModelType.QWEN_3 -> "" // Always enable thinking for Qwen-3`
`68`	`68`	`else -> ""`
`69`	`69`	`}`
`70`	`70`	`}`
`@@ -74,8 +74,6 @@ object PromptFormat {`
`74`	`74`	`return when (modelType) {`
`75`	`75`	`ModelType.QWEN_3 -> when (token) {`
`76`	`76`	`"<\|im_end\|>" -> ""`
`77`		`- "<think>" -> "Thinking...\n"`
`78`		`- "</think>" -> "\nDone thinking"`
`79`	`77`	`else -> token`
`80`	`78`	`}`
`81`	`79`	`else -> token`