This is Part 2 of a series of articles where I explain how to implement GenAI on Android. [Click here to view the full series.]

Let’s test out the proofreading feature on MLKit

After building the summarisation feature in SmartWriter, integrating proofreading was incredibly fast — I reused almost all the same logic and had a working feature in under an hour.

This time, instead of focusing on UI, I’ll walk you through how the entire ViewModel works — including model download handling, inference, and graceful error fallback — using the new Proofreading API from ML Kit GenAI.

👉 Full app available here:

https://github.com/josegbel/smartwriter

📌 What we’re building

We’re using ML Kit’s on-device GenAI model to take raw user input and return grammatically improved alternatives. For example:

Input:

i think this idea could works, but not sure if make sense

Output:

I think this idea could work, but I’m not sure if it makes sense.

Let’s look at how the ViewModel makes this possible — no Compose needed to understand this one 😉

🧠 The ViewModel in full

Here’s the complete code for the ProofreadingViewModel. I’ll break it down below:

package com.example.smartwriter.viewmodel

import android.content.Context
import android.util.Log
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.example.smartwriter.ui.model.ProofreadingUiEvent
import com.example.smartwriter.ui.model.ProofreadingUiState
import com.google.mlkit.genai.common.DownloadCallback
import com.google.mlkit.genai.common.FeatureStatus
import com.google.mlkit.genai.common.GenAiException
import com.google.mlkit.genai.proofreading.Proofreader
import com.google.mlkit.genai.proofreading.ProofreaderOptions
import com.google.mlkit.genai.proofreading.Proofreading
import com.google.mlkit.genai.proofreading.ProofreadingRequest
import kotlinx.coroutines.flow.*
import kotlinx.coroutines.guava.await
import kotlinx.coroutines.launch
import javax.inject.Inject

class ProofreadingViewModel
    @Inject
    constructor() : ViewModel() {

    companion object {
        private val TAG = ProofreadingViewModel::class.java.simpleName
    }

    private val _uiState = MutableStateFlow(ProofreadingUiState())
    val uiState: StateFlow<ProofreadingUiState> = _uiState.asStateFlow()

    private val _uiEvent = MutableSharedFlow<ProofreadingUiEvent>()
    val uiEvent: SharedFlow<ProofreadingUiEvent> = _uiEvent.asSharedFlow()

    private var proofreader: Proofreader? = null

    override fun onCleared() {
        proofreader?.close()
        super.onCleared()
    }

    fun onInputTextChanged(newText: String) {
        _uiState.update { it.copy(inputText = newText) }
    }

    fun onProofreadClicked(context: Context) {
        viewModelScope.launch {
            try {
                val options =
                    ProofreaderOptions
                        .builder(context)
                        .setInputType(ProofreaderOptions.InputType.KEYBOARD)
                        .setLanguage(ProofreaderOptions.Language.ENGLISH)
                        .build()

                proofreader = Proofreading.getClient(options)

                prepareAndStartProofreading()
            } catch (e: Exception) {
                Log.e(TAG, "Error in onProofreadClicked: ${e.message}", e)
                _uiEvent.emit(ProofreadingUiEvent.Error(message = "Error: ${e.message}"))
            }
        }
    }

🔍 Breakdown — Initialisation

• We use ProofreaderOptions to configure the input type (keyboard text vs voice) and language (English).
• Proofreading.getClient(options) creates the proofreader instance.
• Then we delegate to prepareAndStartProofreading().

 

suspend fun prepareAndStartProofreading() {
    val featureStatus = proofreader?.checkFeatureStatus()?.await()
    Log.d(TAG, "Feature status: $featureStatus")

    when (featureStatus) {
        FeatureStatus.DOWNLOADABLE -> {
            Log.d(TAG, "Feature DOWNLOADABLE – starting download")
            downloadFeature()
        }

        FeatureStatus.DOWNLOADING -> {
            Log.d(TAG, "Feature DOWNLOADING – will start once ready")
            proofreader?.let { startProofreadingRequest(uiState.value.inputText, it) }
        }

        FeatureStatus.AVAILABLE -> {
            Log.d(TAG, "Feature AVAILABLE – running inference")
            _uiState.update { it.copy(isLoading = true) }
            proofreader?.let {
                Log.d(TAG, "starting proofreading request")
                startProofreadingRequest(uiState.value.inputText, it)
            }
        }

        FeatureStatus.UNAVAILABLE, null -> {
            Log.e(TAG, "Feature UNAVAILABLE")
            _uiEvent.emit(ProofreadingUiEvent.Error(message = "Your device does not support this feature."))
        }
    }
}

 

🔍 Breakdown — Feature availability

• checkFeatureStatus() tells us whether the Gemini Nano model is available, downloading, or needs to be downloaded.
• If it’s downloadable, we call downloadFeature().
• If it’s already available, we immediately run the inference.
• We also handle DOWNLOADING and UNAVAILABLE.

This flexibility ensures the user gets feedback even when the model isn’t ready yet.

private fun downloadFeature() {
    proofreader?.downloadFeature(
        object : DownloadCallback {
            override fun onDownloadStarted(bytesToDownload: Long) {
                _uiState.update { it.copy(isLoading = true) }
                Log.d(TAG, "Download started – bytesToDownload=$bytesToDownload")
            }

            override fun onDownloadProgress(totalBytesDownloaded: Long) {
                _uiState.update { it.copy(isLoading = true) }
                Log.d(TAG, "Download progress – totalBytesDownloaded=$totalBytesDownloaded")
            }

            override fun onDownloadCompleted() {
                _uiState.update { it.copy(isLoading = false) }
                Log.d(TAG, "Download completed – starting inference")
                proofreader?.let { startProofreadingRequest(uiState.value.inputText, it) }
            }

            override fun onDownloadFailed(e: GenAiException) {
                _uiState.update { it.copy(isLoading = false) }
                Log.e(TAG, "Download failed: ${e.message}", e)
                _uiEvent.tryEmit(
                    ProofreadingUiEvent.Error(
                        message = "Download failed: ${e.message}",
                    ),
                )
            }
        },
    )
}

🔍 Breakdown — Download logic

• This handles model download events and updates the UI accordingly.
• On successful download, we trigger the inference again — using the same input.

 

 fun startProofreadingRequest(
        text: String,
        proofreader: Proofreader,
    ) {
        val proofreadingRequest = ProofreadingRequest.builder(text).build()
        _uiState.update { it.copy(isLoading = true) }

        viewModelScope.launch {
            try {
                val results = proofreader.runInference(proofreadingRequest).await().results
                _uiState.update { state ->
                    state.copy(correctionSuggestions = results.map { it.text })
                }
            } catch (e: Exception) {
                _uiEvent.emit(
                    ProofreadingUiEvent.Error(
                        message = "Error during proofreading: ${e.message}",
                    ),
                )
            } finally {
                _uiState.update { it.copy(isLoading = false) }
            }
        }
    }
}

🔍 Breakdown — Running the model

• The input string is wrapped in a ProofreadingRequest.
• runInference(…).await() gives us the corrected suggestions.
• We extract and display result.text for each suggestion in the UI.

🧪 Key takeaways

• ✅ If you’ve implemented summarisation already, adding proofreading is fast.
• ✅ The Gemini Nano model gives usable, human-like corrections.
• 🧠 You get multiple suggestions, not diffs — you can choose how to present them.
• 📵 Emulators and unsupported phones won’t work. I tested on a Galaxy S25 Ultra.
• 🌍 Currently only supports English.

🚀 Try it yourself

This is Part 2 of my SmartWriter blog series. You can find the full app (including Compose UI and previews) here:

👉 https://github.com/josegbel/smartwriter

Next up: Text Rewriting — where we’ll teach the app to rephrase user input in different tones and lengths.

Follow along and let me know what you’d improve.

This article was previously published on proandroiddev.com.