[AI] Add configurable model generation for AI On-Device (#8043)

Introduced `GenerationConfig`, `ModelConfig`, `ModelReleaseStage`, and
`ModelPreference` to `firebase-ai-ondevice-interop` to allow for
configurable model selection.

Updated the internal `genaiPrompt` dependency to `1.0.0-beta2` to
support the new configuration options.

Deprecated the parameter-less
`FirebaseAIOnDeviceGenerativeModelFactory.newGenerativeModel()` method
in favor of a new overload that accepts a `GenerationConfig`.

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: rlazo

ai-logic/firebase-ai-ondevice-interop/CHANGELOG.md

@@ -1,5 +1,7 @@
 # Unreleased
 
+- [feature] Added support for model selection, required for nano-v4. (#8043)
+
 # 16.0.0-beta01
 
 - [feature] Initial release.

ai-logic/firebase-ai-ondevice-interop/api.txt

@@ -36,7 +36,8 @@ package com.google.firebase.ai.ondevice.interop {
   }
 
   public interface FirebaseAIOnDeviceGenerativeModelFactory {
-    method public com.google.firebase.ai.ondevice.interop.GenerativeModel newGenerativeModel();
+    method @Deprecated public com.google.firebase.ai.ondevice.interop.GenerativeModel newGenerativeModel();
+    method public com.google.firebase.ai.ondevice.interop.GenerativeModel newGenerativeModel(com.google.firebase.ai.ondevice.interop.GenerationConfig? generationConfig);
   }
 
   public final class FirebaseAIOnDeviceInvalidRequestException extends com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException {
@@ -75,6 +76,12 @@ package com.google.firebase.ai.ondevice.interop {
     property public final java.util.List<com.google.firebase.ai.ondevice.interop.Candidate> candidates;
   }
 
+  public final class GenerationConfig {
+    ctor public GenerationConfig(com.google.firebase.ai.ondevice.interop.ModelConfig modelConfig);
+    method public com.google.firebase.ai.ondevice.interop.ModelConfig getModelConfig();
+    property public final com.google.firebase.ai.ondevice.interop.ModelConfig modelConfig;
+  }
+
   public interface GenerativeModel {
     method public suspend Object? countTokens(com.google.firebase.ai.ondevice.interop.GenerateContentRequest request, kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.interop.CountTokensResponse>);
     method public suspend Object? generateContent(com.google.firebase.ai.ondevice.interop.GenerateContentRequest request, kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.interop.GenerateContentResponse>);
@@ -91,6 +98,24 @@ package com.google.firebase.ai.ondevice.interop {
     property public final android.graphics.Bitmap bitmap;
   }
 
+  public final class ModelConfig {
+    ctor public ModelConfig(com.google.firebase.ai.ondevice.interop.ModelReleaseStage releaseStage = com.google.firebase.ai.ondevice.interop.ModelReleaseStage.STABLE, com.google.firebase.ai.ondevice.interop.ModelPreference preference = com.google.firebase.ai.ondevice.interop.ModelPreference.FULL);
+    method public com.google.firebase.ai.ondevice.interop.ModelPreference getPreference();
+    method public com.google.firebase.ai.ondevice.interop.ModelReleaseStage getReleaseStage();
+    property public final com.google.firebase.ai.ondevice.interop.ModelPreference preference;
+    property public final com.google.firebase.ai.ondevice.interop.ModelReleaseStage releaseStage;
+  }
+
+  public enum ModelPreference {
+    enum_constant public static final com.google.firebase.ai.ondevice.interop.ModelPreference FAST;
+    enum_constant public static final com.google.firebase.ai.ondevice.interop.ModelPreference FULL;
+  }
+
+  public enum ModelReleaseStage {
+    enum_constant public static final com.google.firebase.ai.ondevice.interop.ModelReleaseStage PREVIEW;
+    enum_constant public static final com.google.firebase.ai.ondevice.interop.ModelReleaseStage STABLE;
+  }
+
   public interface Part {
   }
 

ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/FirebaseAIOnDeviceGenerativeModelFactory.kt

@@ -28,5 +28,19 @@ public interface FirebaseAIOnDeviceGenerativeModelFactory {
    *
    * @return A new [GenerativeModel] instance ready for use.
    */
+  @Deprecated(
+    message = "Use newGenerativeModel(GenerationConfig?) instead",
+    replaceWith = ReplaceWith("newGenerativeModel(null)")
+  )
   public fun newGenerativeModel(): GenerativeModel
+
+  /**
+   * Creates and returns a new instance of [GenerativeModel] optionally configured with
+   * [GenerationConfig].
+   *
+   * @param generationConfig The configuration for the model, `null` if the default configuration
+   * should be used.
+   * @return A new [GenerativeModel] instance ready for use.
+   */
+  public fun newGenerativeModel(generationConfig: GenerationConfig?): GenerativeModel
 }

ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerationConfig.kt

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.firebase.ai.ondevice.interop
+
+public class GenerationConfig(public val modelConfig: ModelConfig) {
+  override fun toString(): String = "GenerationConfig(modelConfig=$modelConfig)"
+
+  override fun equals(other: Any?): Boolean =
+    other is GenerationConfig && modelConfig == other.modelConfig
+
+  override fun hashCode(): Int = modelConfig.hashCode() ?: 0
+}
+
+/**
+ * Configuration parameters for model selection.
+ *
+ * @property releaseStage The release stage of the model to use.
+ * @property preference The performance preference for the model.
+ */
+public class ModelConfig(
+  public val releaseStage: ModelReleaseStage = ModelReleaseStage.STABLE,
+  public val preference: ModelPreference = ModelPreference.FULL,
+) {
+  override fun equals(other: Any?): Boolean =
+    other is ModelConfig && releaseStage == other.releaseStage && preference == other.preference
+
+  override fun hashCode(): Int {
+    var result = releaseStage.hashCode()
+    result = 31 * result + preference.hashCode()
+    return result
+  }
+
+  override fun toString(): String {
+    return "ModelConfig(releaseStage=$releaseStage, preference=$preference)"
+  }
+}
+/** Defines the release stage of the model. */
+public enum class ModelReleaseStage {
+  /**
+   * Selects the latest model version that is fully tested and on consumer devices. This is the
+   * default setting.
+   */
+  STABLE,
+
+  /**
+   * Selects the latest model version in the preview stage. This stage lets you test beta features
+   * or newer model architectures before they are widely deployed.
+   */
+  PREVIEW,
+}
+
+/** Defines the performance preference for the model. */
+public enum class ModelPreference {
+  /** Recommended when model accuracy and full capabilities are prioritized over speed. */
+  FULL,
+
+  /** Recommended for latency-sensitive apps that require minimal response times. */
+  FAST,
+}

ai-logic/firebase-ai-ondevice/CHANGELOG.md

@@ -1,5 +1,7 @@
 # Unreleased
 
+- [feature] Added support for model selection, required for nano-v4. (#8043)
+
 # 16.0.0-beta01
 
 - [feature] Initial release.

ai-logic/firebase-ai-ondevice/api.txt

@@ -25,11 +25,21 @@ package com.google.firebase.ai.ondevice {
   }
 
   public final class FirebaseAIOnDevice {
-    method public suspend Object? checkStatus(kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.OnDeviceModelStatus>);
-    method public kotlinx.coroutines.flow.Flow<com.google.firebase.ai.ondevice.DownloadStatus> download();
+    method public suspend Object? checkStatus(com.google.firebase.ai.ondevice.OnDeviceModelOption option, kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.OnDeviceModelStatus>);
+    method public kotlinx.coroutines.flow.Flow<com.google.firebase.ai.ondevice.DownloadStatus> download(com.google.firebase.ai.ondevice.OnDeviceModelOption option);
     field public static final com.google.firebase.ai.ondevice.FirebaseAIOnDevice INSTANCE;
   }
 
+  public final class OnDeviceModelOption {
+    field public static final com.google.firebase.ai.ondevice.OnDeviceModelOption.Companion Companion;
+    field public static final com.google.firebase.ai.ondevice.OnDeviceModelOption PREVIEW;
+    field public static final com.google.firebase.ai.ondevice.OnDeviceModelOption PREVIEW_FAST;
+    field public static final com.google.firebase.ai.ondevice.OnDeviceModelOption STABLE;
+  }
+
+  public static final class OnDeviceModelOption.Companion {
+  }
+
   public final class OnDeviceModelStatus {
     field public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus AVAILABLE;
     field public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus.Companion Companion;

ai-logic/firebase-ai-ondevice/firebase-ai-ondevice.gradle.kts

@@ -68,7 +68,7 @@ kotlin {
 
 dependencies {
   implementation(libs.genai.prompt)
-  implementation("com.google.firebase:firebase-ai-ondevice-interop:16.0.0-beta01")
+  implementation(project(":ai-logic:firebase-ai-ondevice-interop"))
 
   implementation(libs.firebase.common)
   implementation(libs.firebase.components)

ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/Converters.kt

@@ -21,9 +21,15 @@ import com.google.firebase.ai.ondevice.interop.CountTokensResponse
 import com.google.firebase.ai.ondevice.interop.FinishReason
 import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceInvalidRequestException
 import com.google.firebase.ai.ondevice.interop.GenerateContentResponse
+import com.google.firebase.ai.ondevice.interop.GenerationConfig
+import com.google.firebase.ai.ondevice.interop.ModelConfig
 import com.google.mlkit.genai.prompt.GenerateContentRequest
 import com.google.mlkit.genai.prompt.ImagePart
+import com.google.mlkit.genai.prompt.ModelPreference
+import com.google.mlkit.genai.prompt.ModelReleaseStage
 import com.google.mlkit.genai.prompt.TextPart
+import com.google.mlkit.genai.prompt.generationConfig
+import com.google.mlkit.genai.prompt.modelConfig
 
 // ====================================
 // `Part` converter extension functions
@@ -73,6 +79,34 @@ internal fun com.google.firebase.ai.ondevice.interop.GenerateContentRequest.toMl
 internal fun com.google.mlkit.genai.prompt.GenerateContentResponse.toInterop():
   GenerateContentResponse = GenerateContentResponse(candidates.map { it.toInterop() })
 
+// ================================================
+// `GenerationConfig` converter extension functions
+// ================================================
+internal fun GenerationConfig.toMlKit(): com.google.mlkit.genai.prompt.GenerationConfig =
+  generationConfig {
+    this@toMlKit.modelConfig?.let { modelConfig = it.toMlKit() }
+  }
+
+// ===========================================
+// `ModelConfig` converter extension functions
+// ===========================================
+internal fun ModelConfig.toMlKit(): com.google.mlkit.genai.prompt.ModelConfig = modelConfig {
+  releaseStage = this@toMlKit.releaseStage.toMlKit()
+  preference = this@toMlKit.preference.toMlKit()
+}
+
+private fun com.google.firebase.ai.ondevice.interop.ModelReleaseStage.toMlKit(): Int =
+  when (this) {
+    com.google.firebase.ai.ondevice.interop.ModelReleaseStage.PREVIEW -> ModelReleaseStage.PREVIEW
+    com.google.firebase.ai.ondevice.interop.ModelReleaseStage.STABLE -> ModelReleaseStage.STABLE
+  }
+
+private fun com.google.firebase.ai.ondevice.interop.ModelPreference.toMlKit(): Int =
+  when (this) {
+    com.google.firebase.ai.ondevice.interop.ModelPreference.FULL -> ModelPreference.FULL
+    com.google.firebase.ai.ondevice.interop.ModelPreference.FAST -> ModelPreference.FAST
+  }
+
 private fun generateContentRequest(
   text: com.google.firebase.ai.ondevice.interop.TextPart,
   image: com.google.firebase.ai.ondevice.interop.ImagePart? = null,

ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/FirebaseAIOnDevice.kt

@@ -18,8 +18,13 @@ package com.google.firebase.ai.ondevice
 
 import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException
 import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceUnknownException
+import com.google.firebase.ai.ondevice.interop.GenerationConfig
 import com.google.mlkit.genai.common.FeatureStatus
 import com.google.mlkit.genai.prompt.Generation
+import com.google.mlkit.genai.prompt.ModelPreference as MlKitModelPreference
+import com.google.mlkit.genai.prompt.ModelReleaseStage as MlKitModelReleaseStage
+import com.google.mlkit.genai.prompt.generationConfig
+import com.google.mlkit.genai.prompt.modelConfig
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.map
 
@@ -33,10 +38,13 @@ public object FirebaseAIOnDevice {
   /**
    * Checks the current status / availability of the on-device AI model.
    *
+   * @param option The configuration option for the model.
    * @return An [OnDeviceModelStatus] object indicating the current state of the model.
    */
-  public suspend fun checkStatus(): OnDeviceModelStatus {
-    return OnDeviceModelStatus.fromFeatureStatus(Generation.getClient().checkStatus())
+  public suspend fun checkStatus(option: OnDeviceModelOption): OnDeviceModelStatus {
+    return OnDeviceModelStatus.fromFeatureStatus(
+      Generation.getClient(option.toMlKit()).checkStatus()
+    )
   }
 
   /**
@@ -46,12 +54,51 @@ public object FirebaseAIOnDevice {
    * Consumers should collect the flow to start the download process, and optionally process any
    * updates on the download state, progress, and completion or failure.
    *
+   * @param option The configuration option for the model.
    * @return A [Flow] of [DownloadStatus] objects representing the download lifecycle.
    */
-  public fun download(): Flow<DownloadStatus> =
-    Generation.getClient().download().map { DownloadStatus.fromMlKit(it) }
+  public fun download(option: OnDeviceModelOption): Flow<DownloadStatus> {
+    return Generation.getClient(option.toMlKit()).download().map { DownloadStatus.fromMlKit(it) }
+  }
 }
 
+/** Options for configuring the on-device AI model. */
+public class OnDeviceModelOption private constructor(private val value: String) {
+  override fun toString(): String = value
+
+  public companion object {
+    /** Selects the latest stable model. */
+    @JvmField public val STABLE: OnDeviceModelOption = OnDeviceModelOption("stable")
+
+    /** Selects the latest preview model with full performance. */
+    @JvmField public val PREVIEW: OnDeviceModelOption = OnDeviceModelOption("preview")
+
+    /** Selects the latest preview model optimized for speed. */
+    @JvmField public val PREVIEW_FAST: OnDeviceModelOption = OnDeviceModelOption("preview_fast")
+  }
+}
+
+internal fun OnDeviceModelOption.toMlKit(): com.google.mlkit.genai.prompt.GenerationConfig =
+  generationConfig {
+    modelConfig = modelConfig {
+      when (this@toMlKit) {
+        OnDeviceModelOption.STABLE -> {
+          releaseStage = MlKitModelReleaseStage.STABLE
+          preference = MlKitModelPreference.FULL
+        }
+        OnDeviceModelOption.PREVIEW -> {
+          releaseStage = MlKitModelReleaseStage.PREVIEW
+          preference = MlKitModelPreference.FULL
+        }
+        OnDeviceModelOption.PREVIEW_FAST -> {
+          releaseStage = MlKitModelReleaseStage.PREVIEW
+          preference = MlKitModelPreference.FAST
+        }
+        else -> throw IllegalArgumentException("Unknown option: ${this@toMlKit}")
+      }
+    }
+  }
+
 /** Represents the current status of the on-device AI model. */
 public class OnDeviceModelStatus private constructor(private val status: Int) {
   public companion object {

ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/FirebaseAIOnDeviceComponent.kt

@@ -17,7 +17,9 @@
 package com.google.firebase.ai.ondevice
 
 import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceGenerativeModelFactory
+import com.google.firebase.ai.ondevice.interop.GenerationConfig
 import com.google.firebase.ai.ondevice.interop.GenerativeModel
+import com.google.mlkit.genai.prompt.Generation
 
 /**
  * Factory class for Firebase AI OnDevice.
@@ -26,5 +28,16 @@ import com.google.firebase.ai.ondevice.interop.GenerativeModel
  */
 internal class FirebaseAIOnDeviceComponent : FirebaseAIOnDeviceGenerativeModelFactory {
 
-  override fun newGenerativeModel(): GenerativeModel = GenerativeModelImpl()
+  @Deprecated(
+    "Use newGenerativeModel(GenerationConfig?) instead",
+    replaceWith = ReplaceWith("newGenerativeModel(null)")
+  )
+  override fun newGenerativeModel(): GenerativeModel = newGenerativeModel(null)
+
+  override fun newGenerativeModel(generationConfig: GenerationConfig?): GenerativeModel =
+    if (generationConfig == null) {
+      GenerativeModelImpl(Generation.getClient())
+    } else {
+      GenerativeModelImpl(Generation.getClient(generationConfig.toMlKit()))
+    }
 }