Add Markdown rendering and collapsible thinking block to LlamaDemo

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 {

llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/Message.kt

@@ -44,10 +44,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
@@ -59,6 +65,7 @@ class Message(
         return Message(sourceText, isSent, messageType, promptID, timestamp).also {
             it.tokensPerSecond = tokensPerSecond
             it.totalGenerationTime = totalGenerationTime
+            it.thinkingContent = thinkingContent
         }
     }
 

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

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

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)
+                }
+            }
+        }
+    }
+}

llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/screens/ChatScreen.kt

@@ -195,8 +195,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,

llm/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ui/viewmodel/ChatViewModel.kt

@@ -60,6 +60,9 @@ 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
@@ -643,6 +646,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
@@ -756,6 +760,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 &&
@@ -773,21 +787,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)
             }
-            // Increment scroll trigger to auto-scroll during generation
-            scrollTrigger++
         }
+
+        // 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) {