Skip to content

Commit 840b3f2

Browse files
[AI] Introduce OnDeviceExtension for GenerativeModel (#8086)
Adds `OnDeviceExtension` to `GenerativeModel` in `firebase-ai`, centralizing on-device model management. This extension provides methods for `checkStatus`, `download`, and `warmUp` for on-device models. This removes the need for `firebase-ai-ondevice` SDK to have a public API. Key changes include: * Defining: new `DownloadStatus` and `OnDeviceModelStatus` classes in `firebase-ai` for public consumption, marked as `PublicPreviewAPI`. * Creating: corresponding `DownloadStatusInterop` and `OnDeviceModelStatusInterop` classes in `firebase-ai-ondevice-interop` for internal interop with ML Kit. * Implementing: `checkStatus` and `download` on the `GenerativeModel` interop interface in `firebase-ai-ondevice`, utilizing ML Kit's APIs and new converter functions. * Deprecating: the old `FirebaseAIOnDevice.checkStatus`, `FirebaseAIOnDevice.download`, `DownloadStatus`, `OnDeviceModelStatus`, and `OnDeviceModelOption` in `firebase-ai-ondevice`. * Deprecating: the top-level `GenerativeModel.warmUp()` method in `firebase-ai` in favor of `onDeviceExtension?.warmUp()`. * Adding: new unit tests for the `OnDeviceExtension` functionality. --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent fda7a28 commit 840b3f2

13 files changed

Lines changed: 518 additions & 86 deletions

File tree

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,32 @@ package com.google.firebase.ai.ondevice.interop {
1515
property public final int totalTokens;
1616
}
1717

18+
public class DownloadStatusInterop {
19+
ctor public DownloadStatusInterop();
20+
}
21+
22+
public static final class DownloadStatusInterop.DownloadCompleted extends com.google.firebase.ai.ondevice.interop.DownloadStatusInterop {
23+
ctor public DownloadStatusInterop.DownloadCompleted();
24+
}
25+
26+
public static final class DownloadStatusInterop.DownloadFailed extends com.google.firebase.ai.ondevice.interop.DownloadStatusInterop {
27+
ctor public DownloadStatusInterop.DownloadFailed(com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException exception);
28+
method public com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException getException();
29+
property public final com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException exception;
30+
}
31+
32+
public static final class DownloadStatusInterop.DownloadInProgress extends com.google.firebase.ai.ondevice.interop.DownloadStatusInterop {
33+
ctor public DownloadStatusInterop.DownloadInProgress(long totalBytesDownloaded);
34+
method public long getTotalBytesDownloaded();
35+
property public final long totalBytesDownloaded;
36+
}
37+
38+
public static final class DownloadStatusInterop.DownloadStarted extends com.google.firebase.ai.ondevice.interop.DownloadStatusInterop {
39+
ctor public DownloadStatusInterop.DownloadStarted(long bytesToDownload);
40+
method public long getBytesToDownload();
41+
property public final long bytesToDownload;
42+
}
43+
1844
public final class FinishReason {
1945
method public String getName();
2046
property public final String name;
@@ -83,7 +109,9 @@ package com.google.firebase.ai.ondevice.interop {
83109
}
84110

85111
public interface GenerativeModel {
112+
method public suspend Object? checkStatus(kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.interop.OnDeviceModelStatusInterop>);
86113
method public suspend Object? countTokens(com.google.firebase.ai.ondevice.interop.GenerateContentRequest request, kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.interop.CountTokensResponse>);
114+
method public kotlinx.coroutines.flow.Flow<com.google.firebase.ai.ondevice.interop.DownloadStatusInterop> download();
87115
method public suspend Object? generateContent(com.google.firebase.ai.ondevice.interop.GenerateContentRequest request, kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.interop.GenerateContentResponse>);
88116
method public kotlinx.coroutines.flow.Flow<com.google.firebase.ai.ondevice.interop.GenerateContentResponse> generateContentStream(com.google.firebase.ai.ondevice.interop.GenerateContentRequest request);
89117
method public suspend Object? getBaseModelName(kotlin.coroutines.Continuation<? super java.lang.String>);
@@ -116,6 +144,17 @@ package com.google.firebase.ai.ondevice.interop {
116144
enum_constant public static final com.google.firebase.ai.ondevice.interop.ModelReleaseStage STABLE;
117145
}
118146

147+
public final class OnDeviceModelStatusInterop {
148+
field public static final com.google.firebase.ai.ondevice.interop.OnDeviceModelStatusInterop AVAILABLE;
149+
field public static final com.google.firebase.ai.ondevice.interop.OnDeviceModelStatusInterop.Companion Companion;
150+
field public static final com.google.firebase.ai.ondevice.interop.OnDeviceModelStatusInterop DOWNLOADABLE;
151+
field public static final com.google.firebase.ai.ondevice.interop.OnDeviceModelStatusInterop DOWNLOADING;
152+
field public static final com.google.firebase.ai.ondevice.interop.OnDeviceModelStatusInterop UNAVAILABLE;
153+
}
154+
155+
public static final class OnDeviceModelStatusInterop.Companion {
156+
}
157+
119158
public interface Part {
120159
}
121160

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.firebase.ai.ondevice.interop
18+
19+
/** Represents the status of a model download operation. */
20+
public open class DownloadStatusInterop {
21+
/**
22+
* Indicates that the download has started.
23+
*
24+
* @property bytesToDownload The total number of bytes that need to be downloaded.
25+
*/
26+
public class DownloadStarted(public val bytesToDownload: Long) : DownloadStatusInterop()
27+
28+
/**
29+
* Indicates that the download is in progress.
30+
*
31+
* @property totalBytesDownloaded The total number of bytes downloaded so far.
32+
*/
33+
public class DownloadInProgress(public val totalBytesDownloaded: Long) : DownloadStatusInterop()
34+
35+
/**
36+
* Indicates that the download has failed.
37+
*
38+
* @property exception The exception that caused the download to fail.
39+
*/
40+
public class DownloadFailed(public val exception: FirebaseAIOnDeviceException) :
41+
DownloadStatusInterop()
42+
43+
/** Indicates that the download has completed successfully. */
44+
public class DownloadCompleted : DownloadStatusInterop()
45+
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,4 +86,18 @@ public interface GenerativeModel {
8686
* @throws [FirebaseAIOnDeviceNotAvailableException] if model is not available.
8787
*/
8888
public suspend fun warmup()
89+
90+
/**
91+
* Checks the current status / availability of the on-device AI model.
92+
*
93+
* @return An [OnDeviceModelStatusInterop] indicating the current state of the model.
94+
*/
95+
public suspend fun checkStatus(): OnDeviceModelStatusInterop
96+
97+
/**
98+
* Initiates the download of the on-device AI model.
99+
*
100+
* @return A [Flow] of [DownloadStatusInterop] objects representing the download lifecycle.
101+
*/
102+
public fun download(): Flow<DownloadStatusInterop>
89103
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.firebase.ai.ondevice.interop
18+
19+
/** Represents the status of an on-device AI model. */
20+
public class OnDeviceModelStatusInterop private constructor(private val value: String) {
21+
override fun toString(): String = value
22+
23+
override fun equals(other: Any?): Boolean =
24+
other is OnDeviceModelStatusInterop && value == other.value
25+
26+
override fun hashCode(): Int = value.hashCode()
27+
28+
public companion object {
29+
/** The on-device model is unavailable on the device. */
30+
@JvmField
31+
public val UNAVAILABLE: OnDeviceModelStatusInterop = OnDeviceModelStatusInterop("UNAVAILABLE")
32+
33+
/** The on-device model is available for download. */
34+
@JvmField
35+
public val DOWNLOADABLE: OnDeviceModelStatusInterop = OnDeviceModelStatusInterop("DOWNLOADABLE")
36+
37+
/** The on-device model is currently being downloaded. */
38+
@JvmField
39+
public val DOWNLOADING: OnDeviceModelStatusInterop = OnDeviceModelStatusInterop("DOWNLOADING")
40+
41+
/** The on-device model is available and ready for use. */
42+
@JvmField
43+
public val AVAILABLE: OnDeviceModelStatusInterop = OnDeviceModelStatusInterop("AVAILABLE")
44+
}
45+
}

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

Lines changed: 22 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,54 +1,44 @@
11
// Signature format: 3.0
22
package com.google.firebase.ai.ondevice {
33

4-
public abstract class DownloadStatus {
5-
ctor public DownloadStatus();
4+
@Deprecated public abstract class DownloadStatus {
5+
ctor @Deprecated public DownloadStatus();
66
}
77

8-
public static final class DownloadStatus.DownloadCompleted extends com.google.firebase.ai.ondevice.DownloadStatus {
9-
ctor public DownloadStatus.DownloadCompleted();
8+
@Deprecated public static final class DownloadStatus.DownloadCompleted extends com.google.firebase.ai.ondevice.DownloadStatus {
9+
ctor @Deprecated public DownloadStatus.DownloadCompleted();
1010
}
1111

12-
public static final class DownloadStatus.DownloadFailed extends com.google.firebase.ai.ondevice.DownloadStatus {
13-
method public com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException getException();
14-
property public final com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException exception;
12+
@Deprecated public static final class DownloadStatus.DownloadFailed extends com.google.firebase.ai.ondevice.DownloadStatus {
13+
method @Deprecated public com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException getException();
14+
property @Deprecated public final com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException exception;
1515
}
1616

17-
public static final class DownloadStatus.DownloadInProgress extends com.google.firebase.ai.ondevice.DownloadStatus {
18-
method public long getTotalBytesDownloaded();
19-
property public final long totalBytesDownloaded;
17+
@Deprecated public static final class DownloadStatus.DownloadInProgress extends com.google.firebase.ai.ondevice.DownloadStatus {
18+
method @Deprecated public long getTotalBytesDownloaded();
19+
property @Deprecated public final long totalBytesDownloaded;
2020
}
2121

22-
public static final class DownloadStatus.DownloadStarted extends com.google.firebase.ai.ondevice.DownloadStatus {
23-
method public long getBytesToDownload();
24-
property public final long bytesToDownload;
22+
@Deprecated public static final class DownloadStatus.DownloadStarted extends com.google.firebase.ai.ondevice.DownloadStatus {
23+
method @Deprecated public long getBytesToDownload();
24+
property @Deprecated public final long bytesToDownload;
2525
}
2626

2727
public final class FirebaseAIOnDevice {
28-
method public suspend Object? checkStatus(com.google.firebase.ai.ondevice.OnDeviceModelOption option, kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.OnDeviceModelStatus>);
29-
method public kotlinx.coroutines.flow.Flow<com.google.firebase.ai.ondevice.DownloadStatus> download(com.google.firebase.ai.ondevice.OnDeviceModelOption option);
28+
method @Deprecated public suspend Object? checkStatus(kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.OnDeviceModelStatus>);
29+
method @Deprecated public kotlinx.coroutines.flow.Flow<com.google.firebase.ai.ondevice.DownloadStatus> download();
3030
field public static final com.google.firebase.ai.ondevice.FirebaseAIOnDevice INSTANCE;
3131
}
3232

33-
public final class OnDeviceModelOption {
34-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelOption.Companion Companion;
35-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelOption PREVIEW;
36-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelOption PREVIEW_FAST;
37-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelOption STABLE;
33+
@Deprecated public final class OnDeviceModelStatus {
34+
field @Deprecated public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus AVAILABLE;
35+
field @Deprecated public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus.Companion Companion;
36+
field @Deprecated public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus DOWNLOADABLE;
37+
field @Deprecated public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus DOWNLOADING;
38+
field @Deprecated public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus UNAVAILABLE;
3839
}
3940

40-
public static final class OnDeviceModelOption.Companion {
41-
}
42-
43-
public final class OnDeviceModelStatus {
44-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus AVAILABLE;
45-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus.Companion Companion;
46-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus DOWNLOADABLE;
47-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus DOWNLOADING;
48-
field public static final com.google.firebase.ai.ondevice.OnDeviceModelStatus UNAVAILABLE;
49-
}
50-
51-
public static final class OnDeviceModelStatus.Companion {
41+
@Deprecated public static final class OnDeviceModelStatus.Companion {
5242
}
5343

5444
}

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,15 @@ package com.google.firebase.ai.ondevice
1818

1919
import com.google.firebase.ai.ondevice.interop.Candidate
2020
import com.google.firebase.ai.ondevice.interop.CountTokensResponse
21+
import com.google.firebase.ai.ondevice.interop.DownloadStatusInterop
2122
import com.google.firebase.ai.ondevice.interop.FinishReason
23+
import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException
2224
import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceInvalidRequestException
25+
import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceUnknownException
2326
import com.google.firebase.ai.ondevice.interop.GenerateContentResponse
2427
import com.google.firebase.ai.ondevice.interop.GenerationConfig
2528
import com.google.firebase.ai.ondevice.interop.ModelConfig
29+
import com.google.firebase.ai.ondevice.interop.OnDeviceModelStatusInterop
2630
import com.google.mlkit.genai.prompt.GenerateContentRequest
2731
import com.google.mlkit.genai.prompt.ImagePart
2832
import com.google.mlkit.genai.prompt.ModelPreference
@@ -121,3 +125,34 @@ private fun generateContentRequest(
121125
builder.init()
122126
return builder.build()
123127
}
128+
129+
// ========================================================================
130+
// `DownloadStatus` and `OnDeviceModelStatus` converter extension functions
131+
// ========================================================================
132+
internal fun com.google.mlkit.genai.common.DownloadStatus.toInterop(): DownloadStatusInterop =
133+
when (this) {
134+
is com.google.mlkit.genai.common.DownloadStatus.DownloadStarted ->
135+
DownloadStatusInterop.DownloadStarted(bytesToDownload)
136+
is com.google.mlkit.genai.common.DownloadStatus.DownloadProgress ->
137+
DownloadStatusInterop.DownloadInProgress(totalBytesDownloaded)
138+
is com.google.mlkit.genai.common.DownloadStatus.DownloadCompleted ->
139+
DownloadStatusInterop.DownloadCompleted()
140+
is com.google.mlkit.genai.common.DownloadStatus.DownloadFailed ->
141+
DownloadStatusInterop.DownloadFailed(FirebaseAIOnDeviceException.from(e))
142+
else ->
143+
DownloadStatusInterop.DownloadFailed(
144+
FirebaseAIOnDeviceUnknownException("Unknown download status")
145+
)
146+
}
147+
148+
internal fun Int.toInteropStatus(): OnDeviceModelStatusInterop =
149+
when (this) {
150+
com.google.mlkit.genai.common.FeatureStatus.UNAVAILABLE ->
151+
OnDeviceModelStatusInterop.UNAVAILABLE
152+
com.google.mlkit.genai.common.FeatureStatus.DOWNLOADABLE ->
153+
OnDeviceModelStatusInterop.DOWNLOADABLE
154+
com.google.mlkit.genai.common.FeatureStatus.DOWNLOADING ->
155+
OnDeviceModelStatusInterop.DOWNLOADING
156+
com.google.mlkit.genai.common.FeatureStatus.AVAILABLE -> OnDeviceModelStatusInterop.AVAILABLE
157+
else -> OnDeviceModelStatusInterop.UNAVAILABLE
158+
}

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

Lines changed: 8 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,8 @@ package com.google.firebase.ai.ondevice
1818

1919
import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceException
2020
import com.google.firebase.ai.ondevice.interop.FirebaseAIOnDeviceUnknownException
21-
import com.google.firebase.ai.ondevice.interop.GenerationConfig
2221
import com.google.mlkit.genai.common.FeatureStatus
2322
import com.google.mlkit.genai.prompt.Generation
24-
import com.google.mlkit.genai.prompt.ModelPreference as MlKitModelPreference
25-
import com.google.mlkit.genai.prompt.ModelReleaseStage as MlKitModelReleaseStage
26-
import com.google.mlkit.genai.prompt.generationConfig
27-
import com.google.mlkit.genai.prompt.modelConfig
2823
import kotlinx.coroutines.flow.Flow
2924
import kotlinx.coroutines.flow.map
3025

@@ -38,13 +33,11 @@ public object FirebaseAIOnDevice {
3833
/**
3934
* Checks the current status / availability of the on-device AI model.
4035
*
41-
* @param option The configuration option for the model.
4236
* @return An [OnDeviceModelStatus] object indicating the current state of the model.
4337
*/
44-
public suspend fun checkStatus(option: OnDeviceModelOption): OnDeviceModelStatus {
45-
return OnDeviceModelStatus.fromFeatureStatus(
46-
Generation.getClient(option.toMlKit()).checkStatus()
47-
)
38+
@Deprecated("Use `GenerativeModel.OnDeviceExtension` from `firebase-ai` instead")
39+
public suspend fun checkStatus(): OnDeviceModelStatus {
40+
return OnDeviceModelStatus.fromFeatureStatus(Generation.getClient().checkStatus())
4841
}
4942

5043
/**
@@ -54,52 +47,16 @@ public object FirebaseAIOnDevice {
5447
* Consumers should collect the flow to start the download process, and optionally process any
5548
* updates on the download state, progress, and completion or failure.
5649
*
57-
* @param option The configuration option for the model.
5850
* @return A [Flow] of [DownloadStatus] objects representing the download lifecycle.
5951
*/
60-
public fun download(option: OnDeviceModelOption): Flow<DownloadStatus> {
61-
return Generation.getClient(option.toMlKit()).download().map { DownloadStatus.fromMlKit(it) }
52+
@Deprecated("Use `GenerativeModel.OnDeviceExtension` from `firebase-ai` instead")
53+
public fun download(): Flow<DownloadStatus> {
54+
return Generation.getClient().download().map { DownloadStatus.fromMlKit(it) }
6255
}
6356
}
6457

65-
/** Options for configuring the on-device AI model. */
66-
public class OnDeviceModelOption private constructor(private val value: String) {
67-
override fun toString(): String = value
68-
69-
public companion object {
70-
/** Selects the latest stable model. */
71-
@JvmField public val STABLE: OnDeviceModelOption = OnDeviceModelOption("stable")
72-
73-
/** Selects the latest preview model with full performance. */
74-
@JvmField public val PREVIEW: OnDeviceModelOption = OnDeviceModelOption("preview")
75-
76-
/** Selects the latest preview model optimized for speed. */
77-
@JvmField public val PREVIEW_FAST: OnDeviceModelOption = OnDeviceModelOption("preview_fast")
78-
}
79-
}
80-
81-
internal fun OnDeviceModelOption.toMlKit(): com.google.mlkit.genai.prompt.GenerationConfig =
82-
generationConfig {
83-
modelConfig = modelConfig {
84-
when (this@toMlKit) {
85-
OnDeviceModelOption.STABLE -> {
86-
releaseStage = MlKitModelReleaseStage.STABLE
87-
preference = MlKitModelPreference.FULL
88-
}
89-
OnDeviceModelOption.PREVIEW -> {
90-
releaseStage = MlKitModelReleaseStage.PREVIEW
91-
preference = MlKitModelPreference.FULL
92-
}
93-
OnDeviceModelOption.PREVIEW_FAST -> {
94-
releaseStage = MlKitModelReleaseStage.PREVIEW
95-
preference = MlKitModelPreference.FAST
96-
}
97-
else -> throw IllegalArgumentException("Unknown option: ${this@toMlKit}")
98-
}
99-
}
100-
}
101-
10258
/** Represents the current status of the on-device AI model. */
59+
@Deprecated("Use `GenerativeModel.OnDeviceExtension` from `firebase-ai` instead")
10360
public class OnDeviceModelStatus private constructor(private val status: Int) {
10461
public companion object {
10562
/** The on-device model is unavailable on the device. */
@@ -135,6 +92,7 @@ public class OnDeviceModelStatus private constructor(private val status: Int) {
13592
* This class has several concrete subclasses, each representing a specific stage or outcome of the
13693
* download process.
13794
*/
95+
@Deprecated("Use `GenerativeModel.OnDeviceExtension` from `firebase-ai` instead")
13896
public abstract class DownloadStatus {
13997
/**
14098
* Represents when a download has just started.

0 commit comments

Comments
 (0)