How Thinking Mode Works in Google AI Edge Gallery: A Complete Technical Guide

Thinking Mode is a UI-driven feature that streams the LLM's internal reasoning process in a separate collapsible pane while simultaneously generating the final answer.

The Google AI Edge Gallery implements this capability through a configuration-driven pipeline that intercepts the model's secondary reasoning stream and renders it in real-time. This article examines the complete implementation across the Android Kotlin codebase, from the configuration flag to the Jetpack Compose UI components.

Architecture Overview

The feature operates through a six-stage pipeline that bridges the inference engine with the chat interface. When Thinking Mode is active, the system sends an extraContext parameter to the model, consumes a parallel reasoning stream, and maintains a distinct message type to track the progressive disclosure of thought processes.

The flow follows this sequence:

1. Configuration Check – The app reads ConfigKeys.ENABLE_THINKING from user settings
2. Context Injection – The LlmChatViewModel builds an extraContext map containing "enable_thinking" → "true"
3. Stream Consumption – The inference helper emits reasoning tokens via partialThinkingResult callbacks
4. Message Creation – A ChatMessageThinking instance is inserted with inProgress = true
5. UI Rendering – The MessageBodyThinking composable displays the expanding reasoning panel
6. State Completion – The message transitions to inProgress = false when the model finishes reasoning

Enabling Thinking Mode in Configuration

The feature toggle is defined in [Config.kt](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/data/Config.kt#L60-L62) as a boolean preference that persists across sessions.

// Settings screen retrieval pattern
val enableThinking = model.getBooleanConfigValue(
    key = ConfigKeys.ENABLE_THINKING,
    defaultValue = false
)

When enabled, this boolean propagates through the view-model to conditionally populate the inference request metadata.

Injecting Extra Context Before Inference

Before invoking the model, the [LlmChatViewModel.kt](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatViewModel.kt#L213-L217) constructs the extraContext map. This map signals the underlying runtime to emit the secondary reasoning stream alongside the primary response.

val enableThinking = allowThinking &&
    model.getBooleanConfigValue(
        key = ConfigKeys.ENABLE_THINKING,
        defaultValue = false
    )

val extraContext = if (enableThinking) {
    mapOf("enable_thinking" to "true")
} else {
    null
}

model.runtimeHelper.runInference(
    model = model,
    input = userPrompt,
    extraContext = extraContext,
    // additional parameters...
)

The extraContext parameter acts as a contract with the inference engine: when present and set to "true", the model exposes its chain-of-thought through a separate callback channel.

Handling the Thinking Stream

The view-model consumes the reasoning stream through a specialized overload of the resultListener callback defined in [LlmChatViewModel.kt](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatViewModel.kt#L84-L106). This callback receives three parameters: the partial result text, a completion boolean, and the optional partialThinkingResult string.

val resultListener: (String, Boolean, String?) -> Unit = { 
    partialResult, done, partialThinkingResult ->
    
    // Handle primary response stream...
    
    // ---- Thinking stream processing ----
    val thinkingText = partialThinkingResult
    val isThinking = thinkingText != null && thinkingText.isNotEmpty()
    var currentLastMessage = getLastMessage(model)

    if (isThinking) {
        // Initialize thinking message if absent
        if (currentLastMessage?.type != ChatMessageType.THINKING) {
            addMessage(
                model = model,
                message = ChatMessageThinking(
                    content = "",
                    inProgress = true,
                    side = ChatSide.AGENT,
                    accelerator = accelerator,
                    hideSenderLabel = currentLastMessage?.type == 
                        ChatMessageType.COLLAPSABLE_PROGRESS_PANEL
                )
            )
        }
        
        // Append incremental reasoning tokens
        updateLastThinkingMessageContentIncrementally(
            model = model,
            partialContent = thinkingText!!
        )
    } else {
        // Finalize thinking state
        if (currentLastMessage?.type == ChatMessageType.THINKING) {
            val thinkingMsg = currentLastMessage as ChatMessageThinking
            if (thinkingMsg.inProgress) {
                replaceLastMessage(
                    model = model,
                    message = thinkingMsg.copy(inProgress = false),
                    type = ChatMessageType.THINKING
                )
            }
        }
    }
}

The implementation distinguishes between streaming and completion states using the ChatMessageThinking data class (lines 99-122 in ChatMessage.kt), which tracks inProgress to manage UI animations and auto-expansion behavior.

Rendering the Thinking Pane UI

The composable [MessageBodyThinking.kt](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/common/chat/MessageBodyThinking.kt#L49-L80) handles the presentation layer. This component manages its own expansion state while respecting the stream's progress status.

@Composable
fun MessageBodyThinking(thinkingText: String, inProgress: Boolean) {
    var isExpanded by remember { mutableStateOf(false) }

    // Force expansion while model is reasoning
    if (inProgress) isExpanded = true

    Column(
        Modifier
            .fillMaxWidth()
            .padding(horizontal = 12.dp, vertical = 8.dp)
    ) {
        Row(
            Modifier
                .clickable { isExpanded = !isExpanded }
                .padding(vertical = 4.dp),
            verticalAlignment = Alignment.CenterVertically,
            horizontalArrangement = Arrangement.spacedBy(4.dp)
        ) {
            Text(
                text = stringResource(R.string.show_thinking),
                style = MaterialTheme.typography.bodyMedium
            )
            Icon(
                imageVector = if (isExpanded) {
                    Icons.Filled.ArrowDropUp
                } else {
                    Icons.Filled.ArrowDropDown
                },
                contentDescription = if (isExpanded) "Hide thinking" else "Show thinking"
            )
        }

        AnimatedVisibility(
            visible = isExpanded,
            enter = expandVertically(),
            exit = shrinkVertically()
        ) {
            MarkdownText(
                text = thinkingText,
                smallFontSize = true
            )
        }
    }
}

The component uses AnimatedVisibility to smoothly expand and collapse the reasoning content, while the inProgress parameter ensures the panel remains visible during active generation. Once the model completes reasoning and inProgress becomes false, users can toggle the panel independently.

Summary

• Thinking Mode relies on the ENABLE_THINKING configuration key to toggle the feature globally across the app
• The extraContext map passes "enable_thinking" → "true" to the inference runtime, triggering the secondary reasoning stream
• ChatMessageThinking (defined in ChatMessage.kt) stores reasoning content separately from standard chat messages with an inProgress state flag
• The partialThinkingResult callback parameter in LlmChatViewModel.kt delivers incremental reasoning tokens distinct from the final answer stream
• MessageBodyThinking (implemented in Jetpack Compose) provides the collapsible UI surface with auto-expand behavior during stream ingestion

Frequently Asked Questions

How do I enable Thinking Mode in the Google AI Edge Gallery app?

Navigate to the Settings panel and toggle the switch labeled with the string resource defined in strings.xml (R.string.show_thinking). This updates the ConfigKeys.ENABLE_THINKING boolean preference. The change takes effect immediately for subsequent chat messages without requiring an app restart.

What happens if the model does not support reasoning streams?

If the underlying LLM runtime does not recognize the "enable_thinking" key in the extraContext map, it simply ignores the parameter. The partialThinkingResult callback will receive null values, and the UI will render only the standard ChatMessageText without creating a ChatMessageThinking instance. The app handles this gracefully by checking isThinking before inserting any thinking-related UI components.

Is the reasoning content persisted in chat history?

Yes. The ChatMessageThinking data class is a first-class citizen in the message model (lines 99-122 of ChatMessage.kt). When replaceLastMessage is called with inProgress = false, the final reasoning content is stored in the chat history alongside regular messages. Users can expand or collapse this content when reviewing past conversations, though the persistence layer stores the complete reasoning text regardless of the UI collapse state.

Does Thinking Mode impact inference performance or battery life?

The feature itself adds minimal overhead to the client-side codebase, as it primarily consumes an existing stream from the model. However, generating reasoning tokens requires additional computation from the LLM, which may increase latency and power consumption depending on the model's implementation. The extraContext flag merely signals the runtime to expose its internal reasoning; the cost depends on the specific model weights deployed via the Google AI Edge Gallery runtime.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}