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

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

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

@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/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) provides the scaffold containing the top app bar, scrolling message list, and back-handler logic. Inside the scrolling column, ChatPanel (from 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/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/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.

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

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/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/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:

@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:

@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:

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.

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