litert-lm

android {
    defaultConfig {
        minSdk = 24
    }
}

dependencies {
    implementation("com.google.ai.edge.litertlm:litertlm-android:0.18.0")
}

dependencies {
    ksp("org.jetbrains.kotlin:kotlin-metadata-jvm:2.4.21")
}

<uses-permission android:name="android.permission.INTERNET"/>

<application>
    <uses-native-library
        android:name="libvndksupport.so"
        android:required="false"/>
    <uses-native-library
        android:name="libOpenCL.so"
        android:required="false"/>
</application>

import android.app.Application
import androidx.lifecycle.AndroidViewModel
import com.google.ai.edge.litertlm.Backend
import com.google.ai.edge.litertlm.Contents
import com.google.ai.edge.litertlm.Conversation
import com.google.ai.edge.litertlm.ConversationConfig
import com.google.ai.edge.litertlm.Engine
import com.google.ai.edge.litertlm.EngineConfig
import java.util.concurrent.Executors
import kotlin.coroutines.cancellation.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.asCoroutineDispatcher
import kotlinx.coroutines.cancel
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

data class ChatState(
    val ready: Boolean = false,
    val busy: Boolean = false,
    val reply: String = "",
    val error: String? = null
)

class ChatViewModel(app: Application) : AndroidViewModel(app) {
    private val executor = Executors.newSingleThreadExecutor()
    private val scope = CoroutineScope(
        SupervisorJob() + executor.asCoroutineDispatcher()
    )
    private val mutex = Mutex()
    private var engine: Engine? = null
    private var conversation: Conversation? = null
    private val _state = MutableStateFlow(ChatState())
    val state: StateFlow<ChatState> = _state.asStateFlow()

    fun load(modelPath: String, backend: Backend = Backend.CPU()) {
        scope.launch {
            mutex.withLock {
                if (engine?.engineConfig?.modelPath == modelPath) return@launch
                release()
                try {
                    val config = EngineConfig(
                        modelPath = modelPath,
                        backend = backend,
                        cacheDir = getApplication<Application>().cacheDir.path
                    )
                    engine = Engine(config).also { it.initialize() }
                    conversation = checkNotNull(engine).createConversation(
                        ConversationConfig(
                            systemInstruction =
                                Contents.of("You are a helpful assistant.")
                        )
                    )
                    _state.value = ChatState(ready = true)
                } catch (e: Exception) {
                    release()
                    _state.value = ChatState(error = e.message)
                }
            }
        }
    }

    fun send(text: String) {
        scope.launch {
            if (!_state.value.ready || _state.value.busy) return@launch
            _state.update { it.copy(busy = true, reply = "", error = null) }
            try {
                val replies = checkNotNull(conversation).sendMessageAsync(text)
                replies.collect { message ->
                    _state.update { it.copy(reply = it.reply + message) }
                }
            } catch (e: CancellationException) {
                throw e
            } catch (e: Exception) {
                _state.update { it.copy(error = e.message) }
            } finally {
                _state.update { it.copy(busy = false) }
            }
        }
    }

    private suspend fun release() {
        _state.update { it.copy(ready = false) }
        while (_state.value.busy) {
            conversation?.cancelProcess()
            delay(100)
        }
        conversation?.close()
        conversation = null
        engine?.close()
        engine = null
        _state.value = ChatState()
    }

    override fun onCleared() {
        scope.launch {
            mutex.withLock { release() }
            executor.shutdown()
        }.invokeOnCompletion { scope.cancel() }
    }
}

import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.rememberScrollState
import androidx.compose.foundation.verticalScroll
import androidx.compose.material3.Button
import androidx.compose.material3.Text
import androidx.compose.material3.TextField
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
import androidx.compose.ui.Modifier
import androidx.lifecycle.compose.collectAsStateWithLifecycle
import androidx.lifecycle.viewmodel.compose.viewModel
import com.google.ai.edge.litertlm.Backend

@Composable
fun ChatScreen(modelPath: String, viewModel: ChatViewModel = viewModel()) {
    val state by viewModel.state.collectAsStateWithLifecycle()
    var input by remember { mutableStateOf("") }
    LaunchedEffect(modelPath) { viewModel.load(modelPath, Backend.GPU()) }
    Column {
        Text(
            state.reply,
            Modifier.weight(1f).verticalScroll(rememberScrollState())
        )
        state.error?.let { Text(it) }
        TextField(value = input, onValueChange = { input = it })
        Button(
            onClick = { viewModel.send(input); input = "" },
            enabled = state.ready && !state.busy
        ) {
            Text("Send")
        }
    }
}