Bug 2010698 - add concept-llm and a gemini nano lib implementation of it

Author: MatthewTighe

gradle/libs.versions.toml

@@ -218,6 +218,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

@@ -1733,6 +1733,8 @@ projects:
     publish: true
     upstream_dependencies:
     - components:tooling-lint
+    - components:support-base
+    - components:concept-llm
   components:lib-llm-mlpa:
     description: A off-device implementation of the LLM concept that uses the MLPA
       proxy server

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,56 @@
+/* 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
+
+/**
+ * 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 deliver to the LLM for inference.
+     */
+    suspend fun prompt(prompt: Prompt): Response
+
+    /**
+     * Check the availability of a LLM. This may relate to local device or remote service capabilities,
+     * for example.
+     */
+    suspend fun checkAvailability(): AvailabilityStatus
+
+    /**
+     * The response from prompting a LLM.
+     */
+    sealed class Response {
+        /**
+         * A successful response from a LLM.
+         */
+        data class Success(val reply: String) : Response()
+
+        /**
+         * The LLM is already processing another prompt.
+         */
+        data class InProgress(val reason: String) : Response()
+
+        /**
+         * A failure response from a LLM.
+         */
+        data class Failure(val reason: String) : Response()
+    }
+
+    /**
+     * Whether a LLM is available for processing prompts.
+     */
+    enum class AvailabilityStatus {
+        Available,
+        Preparing,
+        Unavailable,
+    }
+}

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,69 @@
+/* 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.first
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.yield
+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 {
+
+    override suspend fun prompt(prompt: Prompt): Llm.Response {
+        val model = buildModel()
+        return when (model.checkStatus()) {
+            FeatureStatus.AVAILABLE -> {
+                model.getPromptResponse(prompt)
+            }
+            FeatureStatus.DOWNLOADING -> {
+                Llm.Response.InProgress("already downloading")
+            }
+            FeatureStatus.DOWNLOADABLE -> {
+                val result = model.download().onEach {
+                    logger("Download update: $it")
+                    yield()
+                }.first { status ->
+                    status == DownloadStatus.DownloadCompleted || status is DownloadStatus.DownloadFailed
+                }
+
+                if (result is DownloadStatus.DownloadFailed) {
+                    Llm.Response.Failure("download failed")
+                } else {
+                    model.getPromptResponse(prompt)
+                }
+            }
+            else -> Llm.Response.Failure("unavailable")
+        }
+    }
+
+    override suspend fun checkAvailability(): Llm.AvailabilityStatus = when (buildModel().checkStatus()) {
+        FeatureStatus.AVAILABLE -> Llm.AvailabilityStatus.Available
+        FeatureStatus.DOWNLOADING,
+        FeatureStatus.DOWNLOADABLE,
+        -> Llm.AvailabilityStatus.Preparing
+        else -> Llm.AvailabilityStatus.Unavailable
+    }
+
+    private suspend fun GenerativeModel.getPromptResponse(prompt: Prompt): Llm.Response = try {
+        val response = generateContent(prompt.value)
+        Llm.Response.Success(response.candidates[0].text)
+    } catch (e: GenAiException) {
+        Llm.Response.Failure("Gemini Nano inference failed: ${e.message}")
+    }
+}