LLM Inference Android Sample Issue: collectIndexed in sendMessage Causes Duplicate Collections on Multiple Invocations

[icools]

in android sample

ChatViewModel.kt

The collectIndexed function for inferenceModel.partialResults should not be placed inside sendMessage. Doing so will cause collect to be called again on partialResults the second time sendMessage is pressed. As a result, when result returns, duplicate data will be added to the message list. If sendMessage is pressed a third time, it will trigger a third collect on inferenceModel.partialResults, leading to receiving three identical results.

https://github.com/google-ai-edge/mediapipe-samples/blob/main/examples/llm_inference/android/app/src/main/java/com/google/mediapipe/examples/llminference/ChatViewModel.kt

fun sendMessage(userMessage: String) {
        viewModelScope.launch(Dispatchers.IO) {
            _uiState.value.addMessage(userMessage, USER_PREFIX)
            var currentMessageId: String? = _uiState.value.createLoadingMessage()
            setInputEnabled(false)
            try {
                val fullPrompt = _uiState.value.fullPrompt
                inferenceModel.generateResponseAsync(fullPrompt)
                inferenceModel.partialResults
                    .collectIndexed { index, (partialResult, done) ->  // The collect function should not be executed in sendMessage 
 function
                        currentMessageId?.let {
                            if (index == 0) {
                                _uiState.value.appendFirstMessage(it, partialResult)
                            } else {
                                _uiState.value.appendMessage(it, partialResult, done)
                            }
                            if (done) {
                                currentMessageId = null
                                // Re-enable text input
                                setInputEnabled(true)
                            }
                        }
                    }
            } catch (e: Exception) {
                _uiState.value.addMessage(e.localizedMessage ?: "Unknown Error", MODEL_PREFIX)
                setInputEnabled(true)
            }
        }
    }a

Instead, the inferenceModel.partialResults.collectIndexed should be moved and organized within the init block.

init() {
        viewModelScope.launch(Dispatchers.IO) {
            try {
                inferenceModel.partialResults
                    .collectIndexed { index, (partialResult, done) ->  
                        currentMessageId?.let {
                            if (index == 0) {
                                _uiState.value.appendFirstMessage(it, partialResult)
                            } else {
                                _uiState.value.appendMessage(it, partialResult, done)
                            }
                            if (done) {
                                currentMessageId = null
                                // Re-enable text input
                                setInputEnabled(true)
                            }
                        }
                    }
            } catch (e: Exception) {
                _uiState.value.addMessage(e.localizedMessage ?: "Unknown Error", MODEL_PREFIX)
                setInputEnabled(true)
            }
        }
    }
}