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

> Uncover how Google AI Edge Gallery's Thinking Mode reveals LLM reasoning during chat. Get a complete technical guide to this UI-driven feature.

- Repository: [google-ai-edge/gallery](https://github.com/google-ai-edge/gallery)
- Tags: deep-dive
- Published: 2026-04-06

---

**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/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.

```kotlin
// 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/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.

```kotlin
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/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.

```kotlin
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`](https://github.com/google-ai-edge/gallery/blob/main/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/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.

```kotlin
@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`](https://github.com/google-ai-edge/gallery/blob/main/ChatMessage.kt)) stores reasoning content separately from standard chat messages with an `inProgress` state flag
- The **partialThinkingResult** callback parameter in [`LlmChatViewModel.kt`](https://github.com/google-ai-edge/gallery/blob/main/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`](https://github.com/google-ai-edge/gallery/blob/main/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`](https://github.com/google-ai-edge/gallery/blob/main/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.