# How the LLM Chat Interface Is Built with Jetpack Compose in the AI Edge Gallery

> Discover how the LLM chat interface in the AI Edge Gallery is built with Jetpack Compose using reusable components and a stateful ViewModel for streaming on-device responses.

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

---

**The LLM chat interface in the AI Edge Gallery is built as a fully composable UI using Jetpack Compose, combining reusable components like `LlmChatScreen`, `ChatViewWrapper`, and `MessageInputText` with a stateful `LlmChatViewModel` and the `LlmChatModelHelper` runtime bridge to stream responses from the on-device LiteRT LM engine.**

The **google-ai-edge/gallery** repository demonstrates a production-ready implementation of an on-device large language model (LLM) chat experience using modern Android architecture. This article examines how the **LLM chat interface built with Jetpack Compose** structures its presentation layer, view-model hierarchy, and runtime integration to deliver a reactive, streaming chat UI without traditional Fragments or Activities.

## Architecture Overview

The implementation follows a three-layer composable architecture. The **Presentation** layer hosts UI components like `LlmChatScreen` and `ChatPanel`. The **View-Model** layer manages state and orchestration through `LlmChatViewModelBase`. The **Runtime** layer handles model inference via `LlmChatModelHelper`, which interfaces directly with the LiteRT LM engine.

## Screen Entry Point (`LlmChatScreen`)

The entry composable resides in [`LlmChatScreen.kt`](https://github.com/google-ai-edge/gallery/blob/main/LlmChatScreen.kt) and serves as the navigation destination for the chat feature. It accepts a `ModelManagerViewModel` to observe model download status and uses Hilt to inject the screen-level `LlmChatViewModel`.

```kotlin
@Composable
fun LlmChatScreen(
    modelManagerViewModel: ModelManagerViewModel,
    navigateUp: () -> Unit,
    taskId: String = BuiltInTaskId.LLM_CHAT,
    viewModel: LlmChatViewModel = hiltViewModel(),
) {
    ChatViewWrapper(
        viewModel = viewModel,
        modelManagerViewModel = modelManagerViewModel,
        taskId = taskId,
        navigateUp = navigateUp,
    )
}

```

*Source: [[`LlmChatScreen.kt`](https://github.com/google-ai-edge/gallery/blob/main/LlmChatScreen.kt)](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatScreen.kt)*

By delegating all UI work to `ChatViewWrapper`, the screen remains a thin, testable entry point that purely handles dependency injection and callback wiring.

## UI Orchestration with `ChatViewWrapper`

`ChatViewWrapper` bridges the generic chat UI with the LLM-specific view-model. It retrieves the `Task` configuration (including the `allowThinking` capability) and constructs the `onSendMessage` lambda that triggers inference.

```kotlin
@Composable
fun ChatViewWrapper(
    viewModel: LlmChatViewModelBase,
    modelManagerViewModel: ModelManagerViewModel,
    taskId: String,
    navigateUp: () -> Unit,
) {
    val task = modelManagerViewModel.getTaskById(id = taskId)!!
    val allowThinking = task.allowThinking()

    ChatView(
        task = task,
        viewModel = viewModel,
        modelManagerViewModel = modelManagerViewModel,
        onSendMessage = { model, messages ->
            viewModel.generateResponse(
                model = model,
                input = messages.filterIsInstance<ChatMessageText>()
                    .joinToString(" ") { it.content },
                images = messages.filterIsInstance<ChatMessageImage>()
                    .flatMap { it.bitmaps },
                audioMessages = messages.filterIsInstance<ChatMessageAudioClip>(),
                onFirstToken = { /* analytics */ },
                onDone = { /* analytics */ },
                onError = { err -> viewModel.handleError(/*…*/) },
                allowThinking = allowThinking,
            )
        },
    )
}

```

*Source: [`ChatViewWrapper` in [`LlmChatScreen.kt`](https://github.com/google-ai-edge/gallery/blob/main/LlmChatScreen.kt)](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatScreen.kt#L69-L95)*

This composable filters the `List<ChatMessage>` into text, image, and audio components before passing them to `generateResponse()`, ensuring the view-model receives structured inputs.

## Core Chat UI Components

### `ChatView` and `ChatPanel`

The `ChatView` composable (located in [`ui/common/chat/ChatView.kt`](https://github.com/google-ai-edge/gallery/blob/main/ui/common/chat/ChatView.kt)) provides the scaffold containing the top app bar, scrolling message list, and back-handler logic. Inside the scrolling column, `ChatPanel` (from [`ChatPanel.kt`](https://github.com/google-ai-edge/gallery/blob/main/ChatPanel.kt)) renders the actual message history using a `LazyColumn`, handling empty states, benchmarking dialogs, and the dynamic **stop** button that appears only during active generation.

Each `ChatMessage` type—text, image, audio clip, thinking tokens, and errors—uses a dedicated composable renderer, allowing the UI to recompose individual rows efficiently as streaming text arrives.

*Sources: [[`ChatView.kt`](https://github.com/google-ai-edge/gallery/blob/main/ChatView.kt)](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/common/chat/ChatView.kt) & [[`ChatPanel.kt`](https://github.com/google-ai-edge/gallery/blob/main/ChatPanel.kt)](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/common/chat/ChatPanel.kt)*

### Message Input Handling (`MessageInputText`)

`MessageInputText` manages the bottom input row, including the text field, image picker, camera capture, and audio recording. It uses Activity Result APIs (`PickMultipleVisualMedia`, `RequestPermission`) and CameraX for media capture.

```kotlin
MessageInputText(
    task = task,
    modelManagerViewModel = modelManagerViewModel,
    curMessage = curMessage,
    inProgress = uiState.inProgress,
    onSendMessage = { msgs -> onSendMessage(selectedModel, msgs) },
)

```

When the user taps **Send**, `createMessagesToSend()` assembles the payload—attaching images, audio, and text—before invoking the `onSendMessage` callback passed down from `ChatViewWrapper`.

*Source: [[`MessageInputText.kt`](https://github.com/google-ai-edge/gallery/blob/main/MessageInputText.kt)](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/common/chat/MessageInputText.kt)*

## ViewModel State Management (`LlmChatViewModelBase`)

The `LlmChatViewModelBase` class (extended by `LlmChatViewModel`) exposes `StateFlow` properties for `uiState`, `messagesByModel`, and `inProgress` status. The core `generateResponse()` method launches within `viewModelScope` using `Dispatchers.Default` to offload work from the main thread.

```kotlin
fun generateResponse(
    model: Model,
    input: String,
    images: List<Bitmap> = listOf(),
    audioMessages: List<ChatMessageAudioClip> = listOf(),
    onFirstToken: (Model) -> Unit = {},
    onDone: () -> Unit = {},
    onError: (String) -> Unit,
    allowThinking: Boolean = false,
) = viewModelScope.launch(Dispatchers.Default) {
    setInProgress(true)
    setPreparing(true)
    addMessage(model, ChatMessageLoading(accelerator = accelerator))

    while (model.instance == null) delay(100)
    delay(500)

    val audioClips = audioMessages.map { it.genByteArrayForWav() }

    model.runtimeHelper.runInference(
        model = model,
        input = input,
        images = images,
        audioClips = audioClips,
        resultListener = resultListener,
        cleanUpListener = cleanUpListener,
        onError = onError,
        coroutineScope = viewModelScope,
        extraContext = if (allowThinking) mapOf("enable_thinking" to "true") else null,
    )
}

```

The method waits for `model.instance` (initialized asynchronously by `LlmChatModelHelper`) to be non-null, converts audio clips to byte arrays, and delegates inference to the runtime helper. Partial results stream back via `resultListener`, which updates the message list incrementally.

*Source: [[`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)*

## Runtime Integration (`LlmChatModelHelper`)

`LlmChatModelHelper` implements the `LlmModelHelper` contract required by `ModelManagerViewModel`. It manages the LiteRT LM `Engine` and `Conversation` objects, providing three critical operations:

- **`initialize`**: Builds the `Engine` with a hardware-specific `Backend` (CPU, GPU, or NPU), creates a `Conversation`, and stores the pair in `model.instance`.
- **`runInference`**: Assembles a `Contents` object containing text, image bytes, and audio bytes, then calls `conversation.sendMessageAsync()` to stream responses.
- **Clean-up**: Releases engine and conversation resources when models are removed.

Function calling support is enabled by passing a list of `ToolProvider` objects through `initialize` into `ConversationConfig.tools`.

*Source: [[`LlmChatModelHelper.kt`](https://github.com/google-ai-edge/gallery/blob/main/LlmChatModelHelper.kt)](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatModelHelper.kt)*

## End-to-End Data Flow

The reactive data flow ensures the UI updates automatically as inference progresses:

1. **User Action**: Tap **Send** in `MessageInputText` → build `List<ChatMessage>`.
2. **Delegation**: `LlmChatScreen` → `ChatViewWrapper` → `viewModel.generateResponse()`.
3. **Inference**: `LlmChatViewModel` → `LlmChatModelHelper.runInference()` → LiteRT LM `Engine`.
4. **Streaming**: `Conversation` streams partial results → `ResultListener` → `viewModel.updateLastTextMessageContentIncrementally()`.
5. **UI Update**: `ChatPanel` observes `StateFlow` changes via `collectAsState()` and recomposes the `LazyColumn` rows.

Because Jetpack Compose observes `StateFlow` changes, only the affected message rows recompose, maintaining smooth scrolling even with large media attachments.

## Implementation Examples

### Embedding the Chat Screen in Navigation

Since `LlmChatScreen` is a pure composable, it requires no Activity or Fragment hosting:

```kotlin
@Composable
fun GalleryNavHost(navController: NavHostController) {
    NavHost(navController, startDestination = "home") {
        composable("home") { HomeScreen(navController) }
        composable("llmChat") {
            val modelMgr = hiltViewModel<ModelManagerViewModel>()
            LlmChatScreen(
                modelManagerViewModel = modelMgr,
                navigateUp = { navController.popBackStack() }
            )
        }
    }
}

```

### Triggering Programmatic Chat

You can invoke the view-model directly from any composable to start a conversation without user input:

```kotlin
@Composable
fun QuickAskButton(modelMgr: ModelManagerViewModel) {
    val llmVm = hiltViewModel<LlmChatViewModel>()
    Button(onClick = {
        val model = modelMgr.uiState.value.selectedModel
        llmVm.generateResponse(
            model = model,
            input = "What is the weather today?",
            onError = { Log.e("QuickAsk", it) },
            onFirstToken = { Log.d("QuickAsk", "first token") },
            onDone = { Log.d("QuickAsk", "response finished") }
        )
    }) {
        Text("Ask LLM")
    }
}

```

### Adding Function Calling (Tools)

Extend the model initialization with custom tools:

```kotlin
val myTool = object : ToolProvider {
    override val name = "weather_lookup"
    override val description = "Fetches the current weather for a city"
    override suspend fun call(arguments: String): String {
        return "Sunny, 23°C in ${arguments.trim()}"
    }
}

modelMgr.initializeModel(
    context = LocalContext.current,
    task = task,
    model = model,
    systemInstruction = null,
    tools = listOf(myTool)
)

```

## Summary

- **Composable-First Architecture**: The LLM chat interface is built entirely with Jetpack Compose, using `LlmChatScreen` as the entry point and `ChatViewWrapper` to bridge UI and state.
- **Reactive State Management**: `LlmChatViewModelBase` uses `StateFlow` to stream UI states and partial inference results, enabling automatic recomposition of the message list.
- **Multimodal Support**: `MessageInputText` handles text, images (via CameraX and photo picker), and audio recording, passing structured data to the view-model.
- **Runtime Bridge**: `LlmChatModelHelper` abstracts the LiteRT LM `Engine` and `Conversation` APIs, handling initialization, inference, and resource cleanup.
- **Function Calling**: The architecture supports tool use through `ToolProvider` implementations passed during model initialization.

## Frequently Asked Questions

### How does the chat interface handle multimodal inputs like images and audio?

The `MessageInputText` composable manages media capture through Activity Result APIs and CameraX. When the user sends a message, it constructs a `List<ChatMessage>` containing `ChatMessageText`, `ChatMessageImage`, and `ChatMessageAudioClip` instances. The `ChatViewWrapper` filters these into separate `images` and `audioMessages` lists before passing them to `LlmChatViewModelBase.generateResponse()`, which converts audio clips to byte arrays using `genByteArrayForWav()` for the LiteRT LM runtime.

### What is the purpose of `LlmChatModelHelper` in the Jetpack Compose architecture?

`LlmChatModelHelper` implements the `LlmModelHelper` interface to bridge the ViewModel layer with the native LiteRT LM engine. It initializes the `Engine` with hardware-specific backends (CPU/GPU/NPU), creates a `Conversation` session, and exposes `runInference()` to stream responses asynchronously. This separation keeps composables free of native engine dependencies while allowing the ViewModel to trigger inference and handle callbacks.

### How is the LLM inference triggered from the UI components?

When the user taps the send button, `MessageInputText` invokes the `onSendMessage` callback provided by `ChatViewWrapper`. This callback calls `viewModel.generateResponse()` with the selected model, filtered text input, and any attached media. The ViewModel then launches a coroutine on `Dispatchers.Default` and calls `model.runtimeHelper.runInference()`, passing `resultListener` and `onError` callbacks that update the UI state as tokens arrive.

### How does the chat UI manage state during streaming responses?

The `LlmChatViewModelBase` maintains an `inProgress` boolean and a list of messages exposed as `StateFlow`. During inference, the `resultListener` callback invokes `updateLastTextMessageContentIncrementally()` to append partial tokens to the latest assistant message. `ChatPanel` observes these state changes using `collectAsState()`, causing Jetpack Compose to recompose only the affected message row in the `LazyColumn` rather than redrawing the entire list.