From 0f4a603b72d4f69f903d5825f54f20ac747af3ba Mon Sep 17 00:00:00 2001 From: Hansong Zhang Date: Mon, 6 Apr 2026 11:09:55 -0700 Subject: [PATCH] Add Markdown rendering and collapsible thinking block to LlamaDemo - Render model responses as Markdown using compose-richtext library - Parse /<\/think> tags into a separate collapsible thinking block with animated expand/collapse and vertical bar decoration - Always enable thinking mode for Qwen-3 and remove manual toggle button - User messages remain plain text --- llm/android/LlamaDemo/app/build.gradle.kts | 2 + .../example/executorchllamademo/Message.kt | 7 + .../executorchllamademo/PromptFormat.kt | 4 +- .../ui/components/ChatInput.kt | 15 -- .../ui/components/MessageItem.kt | 134 +++++++++++++++++- .../ui/screens/ChatScreen.kt | 2 - .../ui/viewmodel/ChatViewModel.kt | 50 +++++-- 7 files changed, 175 insertions(+), 39 deletions(-) diff --git a/llm/android/LlamaDemo/app/build.gradle.kts b/llm/android/LlamaDemo/app/build.gradle.kts index d4fba5e693..412f50b743 100644 --- 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 index 0e01ef1e45..b79b0f7da0 100644 --- 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 @@ -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 } } 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 index 1fae20c895..263d5e6d54 100644 --- 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 "\n\n\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|>" -> "" - "" -> "Thinking...\n" - "" -> "\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 index f8fd701bd9..1ab14a155c 100644 --- 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 index 8c911b6d62..6ed3ae7411 100644 --- 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 index b4790b5c30..8cedfdc557 100644 --- 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 @@ -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, 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 index e6b1e62ee1..c15822ebe2 100644 --- 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,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 ... 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 / tags + if (processedResult == "") { + isInThinkingBlock = true + return + } + if (processedResult == "") { + 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) {