Gemini Nano and ML Kit GenAI on Android: How On-Device AI Works

dependencies {
    // Beta at the time of writing: check the ML Kit release notes for the current version.
    implementation("com.google.mlkit:genai-summarization:1.0.0-beta1")
    // await() for ListenableFuture, used by the GenAI APIs
    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-guava:<latest>")
}

import android.content.Context
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.summarization.Summarization
import com.google.mlkit.genai.summarization.SummarizationRequest
import com.google.mlkit.genai.summarization.SummarizerOptions
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.guava.await

sealed interface ModelState {
    data object Unavailable : ModelState
    data object NeedsDownload : ModelState
    data class Downloading(val downloadedBytes: Long, val totalBytes: Long) : ModelState
    data object Ready : ModelState
}

class OnDeviceSummarizer(context: Context) : AutoCloseable {

    private val summarizer = Summarization.getClient(
        SummarizerOptions.builder(context)
            .setInputType(SummarizerOptions.InputType.ARTICLE)
            .setOutputType(SummarizerOptions.OutputType.THREE_BULLETS)
            .setLanguage(SummarizerOptions.Language.ENGLISH)
            .build()
    )

    suspend fun status(): ModelState = when (summarizer.checkFeatureStatus().await()) {
        FeatureStatus.AVAILABLE -> ModelState.Ready
        FeatureStatus.DOWNLOADABLE -> ModelState.NeedsDownload
        FeatureStatus.DOWNLOADING -> ModelState.Downloading(0, 0)
        else -> ModelState.Unavailable
    }

    /** Emits download progress, then Ready. Fails the flow if the download fails. */
    fun download(): Flow<ModelState> = callbackFlow {
        var total = 0L
        summarizer.downloadFeature(object : DownloadCallback {
            override fun onDownloadStarted(bytesToDownload: Long) {
                total = bytesToDownload
                trySend(ModelState.Downloading(0, total))
            }
            override fun onDownloadProgress(totalBytesDownloaded: Long) {
                trySend(ModelState.Downloading(totalBytesDownloaded, total))
            }
            override fun onDownloadCompleted() {
                trySend(ModelState.Ready)
                close()
            }
            override fun onDownloadFailed(e: GenAiException) {
                close(e)
            }
        })
        awaitClose { }
    }

    /** Streams partial text through [onPartial] and returns the final summary. */
    suspend fun summarize(text: String, onPartial: (String) -> Unit): String {
        val request = SummarizationRequest.builder(text).build()
        val result = summarizer.runInference(request) { newText -> onPartial(newText) }.await()
        return result.summary
    }

    override fun close() = summarizer.close()
}

data class SummaryUiState(
    val model: ModelState = ModelState.Unavailable,
    val summary: String = "",
    val isSummarizing: Boolean = false,
    val error: String? = null,
)

class SummaryViewModel(private val summarizer: OnDeviceSummarizer) : ViewModel() {

    private val _uiState = MutableStateFlow(SummaryUiState())
    val uiState = _uiState.asStateFlow()

    init {
        viewModelScope.launch { _uiState.update { it.copy(model = summarizer.status()) } }
    }

    fun downloadModel() {
        viewModelScope.launch {
            summarizer.download()
                .catch { _uiState.update { s -> s.copy(error = "Download failed. Check storage and connection.") } }
                .collect { state -> _uiState.update { it.copy(model = state) } }
        }
    }

    fun summarize(note: String) {
        // The API documents input limits (roughly 4,000 tokens); very short inputs give poor summaries.
        if (note.length < 400) {
            _uiState.update { it.copy(error = "Write a little more before summarising.") }
            return
        }
        viewModelScope.launch {
            _uiState.update { it.copy(isSummarizing = true, summary = "", error = null) }
            try {
                val final = summarizer.summarize(note) { partial ->
                    _uiState.update { it.copy(summary = it.summary + partial) }
                }
                _uiState.update { it.copy(isSummarizing = false, summary = final) }
            } catch (e: GenAiException) {
                _uiState.update { it.copy(isSummarizing = false, error = "Couldn't summarise right now. Please try again.") }
            }
        }
    }

    override fun onCleared() = summarizer.close()
}