Bug 2010698 - add a basic concept interface for an LLM r=segun

Pull request: #45

Author: MatthewTighe

gradle/libs.versions.toml

@@ -220,6 +220,7 @@ play-services-base = { group = "com.google.android.gms", name = "play-services-b
 play-services-fido = { group = "com.google.android.gms", name = "play-services-fido", version.ref = "play-services-fido" }
 protobuf-compiler = { group = "com.google.protobuf", name = "protoc", version.ref = "protobuf" }
 protobuf-javalite = { group = "com.google.protobuf", name = "protobuf-javalite", version.ref = "protobuf" }
+mlkit-prompt = { group = "com.google.mlkit", name = "genai-prompt", version = "1.0.0-alpha1" }
 
 # Gradle plugins
 semanticdb-java = { group = "com.sourcegraph", name = "semanticdb-javac", version.ref = "semanticdb-javac" }

mobile/android/android-components/.buildconfig.yml

@@ -1750,6 +1750,11 @@ projects:
     path: components/lib/llm-gemininano
     publish: true
     upstream_dependencies:
+    - components:concept-base
+    - components:concept-fetch
+    - components:concept-llm
+    - components:support-base
+    - components:support-test
     - components:tooling-lint
   components:lib-llm-mlpa:
     description: A off-device implementation of the LLM concept that uses the MLPA

mobile/android/android-components/components/concept/llm/build.gradle

@@ -22,6 +22,8 @@ android {
 }
 
 dependencies {
+    implementation libs.kotlinx.coroutines
+
     testImplementation libs.androidx.test.junit
 }
 

mobile/android/android-components/components/concept/llm/src/main/java/mozilla/components/concept/llm/Example.kt

@@ -0,0 +1,58 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package mozilla.components.concept.llm
+
+import kotlinx.coroutines.flow.Flow
+
+/**
+ * A value type representing a prompt that can be delivered to a LLM.
+ */
+@JvmInline
+value class Prompt(val value: String)
+
+/**
+ * An abstract definition of a LLM that can receive prompts.
+ */
+interface Llm {
+    /**
+     * A prompt request delivered to the LLM for inference, which will stream a series
+     * of [Response]s as they are made available.
+     */
+    suspend fun prompt(prompt: Prompt): Flow<Response>
+
+    /**
+     * A response from prompting a LLM.
+     */
+    sealed class Response {
+
+        /**
+         * A successful response from the LLM has occurred. This may include partial data,
+         * or be an indication that the reply has completed.
+         */
+        sealed class Success : Response() {
+            /**
+             * A (potentially) partial reply from the LLM. This may be a complete reply if
+             * it is short or the underlying implementation does not stream responses.
+             */
+            data class ReplyPart(val value: String) : Success()
+
+            /**
+             * An indication that the reply from the LLM is finishes.
+             */
+            data object ReplyFinished : Success()
+        }
+
+        /**
+         * The LLM is engaged in getting ready to receive prompts. This may include actions like
+         * authenticating with a remote or downloading a local model.
+         */
+        data class Preparing(val status: String) : Response()
+
+        /**
+         * A failure response from a LLM.
+         */
+        data class Failure(val reason: String) : Response()
+    }
+}

mobile/android/android-components/components/concept/llm/src/main/java/mozilla/components/concept/llm/Llm.kt

@@ -22,7 +22,13 @@ android {
 }
 
 dependencies {
+    api project(":components:concept-llm")
+    implementation project(':components:support-base')
+    implementation libs.mlkit.prompt
+
+    testImplementation project(':components:support-test')
     testImplementation libs.androidx.test.junit
+    testImplementation libs.kotlinx.coroutines.test
 }
 
 apply from: '../../../common-config.gradle'

mobile/android/android-components/components/concept/llm/src/test/java/mozilla/components/concept/llm/ExampleTest.kt

@@ -0,0 +1,93 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package mozilla.components.lib.llm.gemini.nano
+
+import com.google.mlkit.genai.common.DownloadStatus
+import com.google.mlkit.genai.common.FeatureStatus
+import com.google.mlkit.genai.common.GenAiException
+import com.google.mlkit.genai.prompt.Generation
+import com.google.mlkit.genai.prompt.GenerativeModel
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.FlowCollector
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+import mozilla.components.concept.llm.Llm
+import mozilla.components.concept.llm.Prompt
+import mozilla.components.support.base.log.logger.Logger
+
+/**
+ * An instance of a LLM that uses local, on-device capabilities provided by Gemini Nano to handle
+ * inference.
+ */
+class GeminiNanoLlm(
+    private val buildModel: () -> GenerativeModel = { Generation.getClient() },
+    private val logger: (String) -> Unit = { message -> Logger("mozac/GeminiNanoLlm").info(message) },
+) : Llm {
+
+    private val model by lazy {
+        buildModel()
+    }
+
+    private val downloadMutex = Mutex()
+
+    override suspend fun prompt(prompt: Prompt): Flow<Llm.Response> = flow {
+        val model = model
+        when (model.checkStatus()) {
+            FeatureStatus.AVAILABLE -> {
+                streamPromptResponses(prompt)
+            }
+            FeatureStatus.DOWNLOADING -> {
+                emit(Llm.Response.Preparing("Downloading model"))
+                // await the completion of the ongoing download
+                downloadMutex.withLock {
+                    if (model.checkStatus() == FeatureStatus.AVAILABLE) {
+                        streamPromptResponses(prompt)
+                    } else {
+                        emit(Llm.Response.Failure("Model should be downloaded and is not"))
+                    }
+                }
+            }
+            FeatureStatus.DOWNLOADABLE -> {
+                emit(Llm.Response.Preparing("Downloading model"))
+                val result = downloadMutex.withLock {
+                     model.download().onEach { status ->
+                        logger("Download update: $status")
+                    }.first { status ->
+                        status == DownloadStatus.DownloadCompleted || status is DownloadStatus.DownloadFailed
+                    }
+                }
+                if (result is DownloadStatus.DownloadFailed) {
+                    val message = "Download failed ${result.e.message}"
+                    logger(message)
+                    emit(Llm.Response.Failure(message))
+                } else {
+                    streamPromptResponses(prompt)
+                }
+            }
+            else -> emit(Llm.Response.Failure("Unavailable"))
+        }
+    }
+
+    private suspend fun FlowCollector<Llm.Response>.streamPromptResponses(prompt: Prompt) = try {
+        // consume replies from the model until it provides a finish reason
+        logger("Beginning model response stream")
+        model.generateContentStream(prompt.value).onEach { response ->
+            emit(Llm.Response.Success.ReplyPart(response.candidates[0].text))
+        }.first {
+            val finishReason = it.candidates[0].finishReason
+            (finishReason != null).also {
+                logger("Model stream completed with: $finishReason")
+            }
+        }
+        emit(Llm.Response.Success.ReplyFinished)
+    } catch (e: GenAiException) {
+        val message = "Gemini Nano inference failed: ${e.message}"
+        logger(message)
+        emit(Llm.Response.Failure(message))
+    }
+}