From a0ffcd54518e9af887197720b0b828702827b8e6 Mon Sep 17 00:00:00 2001 From: Mila <107142260+milaGGL@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:41:50 -0400 Subject: [PATCH 1/5] initial implementation --- .../ai/ksp/SchemaSymbolProcessorVisitor.kt | 126 ++++++++++++++++++ .../interop/GenerateObjectResponseInterop.kt | 25 ++++ .../ai/ondevice/interop/GenerativeModel.kt | 14 ++ .../firebase-ai-ondevice.gradle.kts | 7 +- .../ai/ondevice/GenerativeModelImpl.kt | 54 ++++++++ ai-logic/firebase-ai/firebase-ai.gradle.kts | 2 +- .../OnDeviceGenerativeModelProvider.kt | 24 +++- .../com/google/firebase/ai/type/Candidate.kt | 3 +- .../ai/type/GenerateObjectResponse.kt | 3 + .../OnDeviceGenerativeModelProviderTests.kt | 18 ++- gradle/libs.versions.toml | 2 +- settings.gradle.kts | 1 + 12 files changed, 267 insertions(+), 12 deletions(-) create mode 100644 ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponseInterop.kt diff --git a/ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/SchemaSymbolProcessorVisitor.kt b/ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/SchemaSymbolProcessorVisitor.kt index 9109352dbd6..4e33e25c003 100644 --- a/ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/SchemaSymbolProcessorVisitor.kt +++ b/ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/SchemaSymbolProcessorVisitor.kt @@ -25,12 +25,18 @@ import com.google.devtools.ksp.symbol.KSClassDeclaration import com.google.devtools.ksp.symbol.KSType import com.google.devtools.ksp.symbol.KSVisitorVoid import com.google.devtools.ksp.symbol.Modifier +import com.squareup.kotlinpoet.AnnotationSpec import com.squareup.kotlinpoet.ClassName import com.squareup.kotlinpoet.CodeBlock import com.squareup.kotlinpoet.FileSpec import com.squareup.kotlinpoet.FunSpec +import com.squareup.kotlinpoet.KModifier +import com.squareup.kotlinpoet.ParameterSpec import com.squareup.kotlinpoet.ParameterizedTypeName import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy +import com.squareup.kotlinpoet.PropertySpec +import com.squareup.kotlinpoet.TypeName +import com.squareup.kotlinpoet.TypeSpec import com.squareup.kotlinpoet.ksp.toClassName import com.squareup.kotlinpoet.ksp.toTypeName import com.squareup.kotlinpoet.ksp.writeTo @@ -57,6 +63,11 @@ internal class SchemaSymbolProcessorVisitor( codeGenerator, Dependencies(true, containingFile), ) + val companionFile = generateMlKitCompanionFileSpec(classDeclaration) + companionFile.writeTo( + codeGenerator, + Dependencies(true, containingFile), + ) } fun generateFileSpec(classDeclaration: KSClassDeclaration): FileSpec { @@ -248,4 +259,119 @@ internal class SchemaSymbolProcessorVisitor( builder.addStatement("nullable = %L)", className.isNullable).unindent() return builder.build() } + + private fun isGenerableClass(type: KSType): Boolean { + return type.declaration.annotations.any { it.shortName.getShortName() == "Generable" } || + (type.declaration as? KSClassDeclaration)?.modifiers?.contains(Modifier.DATA) == true && + !type.declaration.qualifiedName?.asString()!!.startsWith("kotlin.") + } + + private fun isListOfGenerableClass(type: KSType): Boolean { + val qualifiedName = type.declaration.qualifiedName?.asString() + if (qualifiedName == "kotlin.collections.List" || qualifiedName == "java.util.List") { + val argType = type.arguments.firstOrNull()?.type?.resolve() + if (argType != null) { + return isGenerableClass(argType) + } + } + return false + } + + private fun mapToMlKitCompanionType(type: KSType, packageName: String): TypeName { + if (isListOfGenerableClass(type)) { + val argType = type.arguments.first()!!.type!!.resolve() + val argClassName = + ClassName( + argType.declaration.packageName.asString(), + "${argType.declaration.simpleName.asString()}_MlKitCompanion" + ) + return ClassName("kotlin.collections", "List").parameterizedBy(argClassName) + } else if (isGenerableClass(type)) { + return ClassName( + type.declaration.packageName.asString(), + "${type.declaration.simpleName.asString()}_MlKitCompanion" + ) + } + return type.toTypeName() + } + + fun generateMlKitCompanionFileSpec(classDeclaration: KSClassDeclaration): FileSpec { + val packageName = classDeclaration.packageName.asString() + val companionClassName = "${classDeclaration.simpleName.asString()}_MlKitCompanion" + val fileBuilder = + FileSpec.builder(packageName, companionClassName).addAnnotation(Generated::class) + + val classBuilder = TypeSpec.classBuilder(companionClassName).addModifiers(KModifier.DATA) + + val generableAnn = + classDeclaration.annotations.firstOrNull { it.shortName.getShortName() == "Generable" } + val classDesc = getStringFromAnnotation(generableAnn, "description") + val mlkitGenerableBuilder = + AnnotationSpec.builder( + ClassName("com.google.mlkit.genai.structuredoutput.annotations", "Generable") + ) + if (!classDesc.isNullOrEmpty()) { + mlkitGenerableBuilder.addMember("description = %S", classDesc) + } + classBuilder.addAnnotation(mlkitGenerableBuilder.build()) + + val primaryConstructor = FunSpec.constructorBuilder() + val toSdkBuilder = + FunSpec.builder("toSdk") + .returns(ClassName(packageName, classDeclaration.simpleName.asString())) + val toSdkArgs = mutableListOf() + + classDeclaration.getAllProperties().forEach { property -> + val propName = property.simpleName.asString() + val propType = property.type.resolve() + val typeName = mapToMlKitCompanionType(propType, packageName) + + val paramBuilder = ParameterSpec.builder(propName, typeName) + val propBuilder = PropertySpec.builder(propName, typeName).initializer(propName) + + val guideAnn = property.annotations.firstOrNull { it.shortName.getShortName() == "Guide" } + if (guideAnn != null) { + val guideValues = + getGuideValuesFromAnnotation(guideAnn, getStringFromAnnotation(guideAnn, "description")) + val mlkitGuideBuilder = + AnnotationSpec.builder( + ClassName("com.google.mlkit.genai.structuredoutput.annotations", "Guide") + ) + if (!guideValues.description.isNullOrEmpty()) + mlkitGuideBuilder.addMember("description = %S", guideValues.description) + if (guideValues.minimum != null) + mlkitGuideBuilder.addMember("minimum = %L", guideValues.minimum) + if (guideValues.maximum != null) + mlkitGuideBuilder.addMember("maximum = %L", guideValues.maximum) + if (guideValues.minItems != null) + mlkitGuideBuilder.addMember("minItems = %L", guideValues.minItems) + if (guideValues.maxItems != null) + mlkitGuideBuilder.addMember("maxItems = %L", guideValues.maxItems) + if (!guideValues.format.isNullOrEmpty()) + mlkitGuideBuilder.addMember("format = %S", guideValues.format) + paramBuilder.addAnnotation(mlkitGuideBuilder.build()) + } + + primaryConstructor.addParameter(paramBuilder.build()) + classBuilder.addProperty(propBuilder.build()) + + if (isListOfGenerableClass(propType)) { + toSdkArgs.add("$propName = this.$propName.map { it.toSdk() }") + } else if (isGenerableClass(propType)) { + toSdkArgs.add("$propName = this.$propName.toSdk()") + } else { + toSdkArgs.add("$propName = this.$propName") + } + } + + classBuilder.primaryConstructor(primaryConstructor.build()) + toSdkBuilder.addStatement( + "return %T(\n ${toSdkArgs.joinToString(",\n ")}\n)", + ClassName(packageName, classDeclaration.simpleName.asString()) + ) + classBuilder.addFunction(toSdkBuilder.build()) + + fileBuilder.addType(classBuilder.build()) + return fileBuilder.build() + } } diff --git a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponseInterop.kt b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponseInterop.kt new file mode 100644 index 00000000000..76157fd9b7b --- /dev/null +++ b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponseInterop.kt @@ -0,0 +1,25 @@ +/* + * 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 + +/** + * Represents a structured object generation response from the on-device model interop layer. + * + * @property instance The generated object instance returned directly by the model engine. + * @property rawJson The raw JSON string representation of the object, if available. + */ +public class GenerateObjectResponseInterop(public val instance: T, public val rawJson: String) {} diff --git a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt index e7e03b7e37f..5276b1e28c0 100644 --- a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt +++ b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt @@ -45,6 +45,20 @@ public interface GenerativeModel { */ public suspend fun generateContent(request: GenerateContentRequest): GenerateContentResponse + /** + * Generates a structured object from the input [GenerateContentRequest] given to the model as a + * prompt. + * + * @param request The input given to the model as a prompt. + * @param outputClass The target class type to which the model will generate. + * @throws [FirebaseAIOnDeviceNotAvailableException] if model is not available. + * @return The structured object response generated by the model. + */ + public suspend fun generateObject( + request: GenerateContentRequest, + outputClass: kotlin.reflect.KClass + ): GenerateObjectResponseInterop + /** * Counts the number of tokens in a prompt using the model's tokenizer. * diff --git a/ai-logic/firebase-ai-ondevice/firebase-ai-ondevice.gradle.kts b/ai-logic/firebase-ai-ondevice/firebase-ai-ondevice.gradle.kts index 8fb562972fe..a437e17a267 100644 --- a/ai-logic/firebase-ai-ondevice/firebase-ai-ondevice.gradle.kts +++ b/ai-logic/firebase-ai-ondevice/firebase-ai-ondevice.gradle.kts @@ -62,13 +62,16 @@ android { } kotlin { - compilerOptions { jvmTarget = JvmTarget.JVM_1_8 } + compilerOptions { + jvmTarget = JvmTarget.JVM_1_8 + freeCompilerArgs.add("-Xskip-metadata-version-check") + } explicitApi() } dependencies { implementation(libs.genai.prompt) - implementation("com.google.firebase:firebase-ai-ondevice-interop:16.0.0-beta03") + implementation(project(":ai-logic:firebase-ai-ondevice-interop")) implementation(libs.firebase.common) implementation(libs.firebase.components) diff --git a/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt b/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt index e5ecd485014..762b0bf11ee 100644 --- a/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt +++ b/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt @@ -53,6 +53,60 @@ internal class GenerativeModelImpl( throw getMappingException(e) } + override suspend fun generateObject( + request: GenerateContentRequest, + outputClass: kotlin.reflect.KClass + ): com.google.firebase.ai.ondevice.interop.GenerateObjectResponseInterop = + try { + android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] ==========================================") + android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] Calling ML Kit generateTypedContentRequest") + android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] Input SDK Class: ${outputClass.qualifiedName}") + val companionClassName = "${outputClass.java.name}_MlKitCompanion" + val targetClass = + try { + Class.forName(companionClassName).kotlin + } catch (e: Exception) { + outputClass + } + android.util.Log.i( + "MLKIT_IO", + "[SDK_BRIDGE] Resolved ML Kit Class: ${targetClass.qualifiedName}" + ) + android.util.Log.i( + "MLKIT_IO", + "[SDK_BRIDGE] includeSchemaInPrompt: true (ML Kit manages prompt formatting)" + ) + android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] ==========================================") + + @Suppress("UNCHECKED_CAST") + val mlkitRequest = + com.google.mlkit.genai.prompt.generateTypedContentRequest( + generateContentRequest = request.toMlKit(), + outputClass = targetClass as kotlin.reflect.KClass, + includeSchemaInPrompt = true + ) + val response = mlkitModel.generateContent(mlkitRequest) + val candidate = response.candidates.firstOrNull() + val mlkitInstance = candidate?.response + android.util.Log.i( + "MLKIT_IO", + "[SDK_BRIDGE] ML Kit returned structured instance: $mlkitInstance (${mlkitInstance?.javaClass?.name})" + ) + + @Suppress("UNCHECKED_CAST") + val sdkInstance = + if (mlkitInstance != null && mlkitInstance.javaClass.name.endsWith("_MlKitCompanion")) { + val toSdkMethod = mlkitInstance.javaClass.getMethod("toSdk") + toSdkMethod.invoke(mlkitInstance) as T + } else { + mlkitInstance as T + } + + com.google.firebase.ai.ondevice.interop.GenerateObjectResponseInterop(sdkInstance, "") + } catch (e: GenAiException) { + throw getMappingException(e) + } + override suspend fun countTokens(request: GenerateContentRequest): CountTokensResponse = try { val response = mlkitModel.countTokens(request.toMlKit()) diff --git a/ai-logic/firebase-ai/firebase-ai.gradle.kts b/ai-logic/firebase-ai/firebase-ai.gradle.kts index af6243f6760..5925a75e3d9 100644 --- a/ai-logic/firebase-ai/firebase-ai.gradle.kts +++ b/ai-logic/firebase-ai/firebase-ai.gradle.kts @@ -98,7 +98,7 @@ dependencies { implementation("androidx.concurrent:concurrent-futures:1.2.0") implementation("androidx.concurrent:concurrent-futures-ktx:1.2.0") implementation("com.google.firebase:firebase-auth-interop:18.0.0") - implementation("com.google.firebase:firebase-ai-ondevice-interop:16.0.0-beta03") + implementation(project(":ai-logic:firebase-ai-ondevice-interop")) // Use different logging libraries depending on the variant releaseImplementation(libs.slf4j.nop) diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt index 5471db5dfdb..4902ceeb9f5 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt @@ -27,6 +27,7 @@ import com.google.firebase.ai.ondevice.interop.TextPart as OnDeviceTextPart import com.google.firebase.ai.type.Candidate import com.google.firebase.ai.type.Content import com.google.firebase.ai.type.CountTokensResponse +import com.google.firebase.ai.type.FinishReason import com.google.firebase.ai.type.FirebaseAIException import com.google.firebase.ai.type.GenerateContentResponse import com.google.firebase.ai.type.GenerateObjectResponse @@ -34,6 +35,7 @@ import com.google.firebase.ai.type.ImagePart import com.google.firebase.ai.type.JsonSchema import com.google.firebase.ai.type.PublicPreviewAPI import com.google.firebase.ai.type.TextPart +import com.google.firebase.ai.type.content import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.catch import kotlinx.coroutines.flow.emitAll @@ -143,9 +145,25 @@ internal class OnDeviceGenerativeModelProvider( override suspend fun generateObject( jsonSchema: JsonSchema, prompt: List - ): GenerateObjectResponse { - throw FirebaseAIException.from( - IllegalArgumentException("On-device mode is not supported for `generateObject`") + ): GenerateObjectResponse = withFirebaseAIExceptionHandling { + ensureOnDeviceModelAvailable() + + val request = buildOnDeviceGenerateContentRequest(prompt) + val interopResponse = onDeviceModel.generateObject(request, jsonSchema.clazz) + val candidate = + Candidate( + content = content { text(interopResponse.rawJson) }, + safetyRatings = emptyList(), + citationMetadata = null, + finishReason = FinishReason.STOP, + finishMessage = null, + groundingMetadata = null, + urlContextMetadata = null, + objectResponse = interopResponse.instance + ) + GenerateObjectResponse( + GenerateContentResponse(listOf(candidate), InferenceSource.ON_DEVICE, null, null, "ondevice"), + jsonSchema ) } diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt index 70f1fa3285b..e89ce5e7251 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt @@ -52,7 +52,8 @@ internal constructor( public val finishReason: FinishReason?, public val finishMessage: String?, public val groundingMetadata: GroundingMetadata?, - public val urlContextMetadata: UrlContextMetadata? + public val urlContextMetadata: UrlContextMetadata?, + internal val objectResponse: Any? = null ) { @OptIn(PublicPreviewAPI::class) diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt index 59fe8e5db02..704d2e0a4b1 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt @@ -41,6 +41,9 @@ internal constructor( @OptIn(InternalSerializationApi::class) public fun getObject(candidateIndex: Int = 0): T? { val candidate = response.candidates[candidateIndex] + if (candidate.objectResponse != null) { + @Suppress("UNCHECKED_CAST") return candidate.objectResponse as T? + } val deserializer = schema.getSerializer() val text = diff --git a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/generativeModel/OnDeviceGenerativeModelProviderTests.kt b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/generativeModel/OnDeviceGenerativeModelProviderTests.kt index 9ef22433578..32e9eb33d4e 100644 --- a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/generativeModel/OnDeviceGenerativeModelProviderTests.kt +++ b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/generativeModel/OnDeviceGenerativeModelProviderTests.kt @@ -108,11 +108,21 @@ internal class OnDeviceGenerativeModelProviderTests { } @Test - fun `generateObject always throws FirebaseAIException`(): Unit = runBlocking { - val schema = mockk>() + fun `generateObject delegates to onDeviceModel`(): Unit = runBlocking { + coEvery { onDeviceModel.isAvailable() } returns true + val schema = JsonSchema.string() + val interopResponse = + com.google.firebase.ai.ondevice.interop.GenerateObjectResponseInterop( + instance = "test object", + rawJson = "\"test object\"" + ) + coEvery { onDeviceModel.generateObject(any(), any>()) } returns + interopResponse - val exception = shouldThrow { provider.generateObject(schema, prompt) } - exception.cause!!::class shouldBe IllegalArgumentException::class + val response = provider.generateObject(schema, prompt) + + response.response.inferenceSource shouldBe InferenceSource.ON_DEVICE + response.getObject() shouldBe "test object" } @Test diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index f2df1b48b1d..7bd4c2a4b23 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -28,7 +28,7 @@ firebaseAnnotations = "17.0.0" firebaseCommon = "22.0.1" firebaseComponents = "19.0.0" firebaseCrashlyticsGradle = "3.0.4" -genaiPrompt = "1.0.0-beta2" +genaiPrompt = "1.0.0-beta3" glide = "5.0.5" googleApiClient = "2.8.1" googleServices = "4.3.15" diff --git a/settings.gradle.kts b/settings.gradle.kts index 59f90d402ed..93afb7ab7df 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -36,6 +36,7 @@ dependencyResolutionManagement { google() mavenLocal() mavenCentral() + maven { url = uri("/Users/milamamat/Downloads/mlkit_genai_structured_output") } } } From 64ef2ef98522ab743250c339bb8b45921906a07e Mon Sep 17 00:00:00 2001 From: Mila <107142260+milaGGL@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:09:06 -0400 Subject: [PATCH 2/5] remove objectResponse from Candidate --- .../OnDeviceGenerativeModelProvider.kt | 7 +-- .../com/google/firebase/ai/type/Candidate.kt | 3 +- .../ai/type/GenerateObjectResponse.kt | 50 ++++++++++++------- 3 files changed, 38 insertions(+), 22 deletions(-) diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt index 4902ceeb9f5..95d22234bee 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt @@ -33,6 +33,7 @@ import com.google.firebase.ai.type.GenerateContentResponse import com.google.firebase.ai.type.GenerateObjectResponse import com.google.firebase.ai.type.ImagePart import com.google.firebase.ai.type.JsonSchema +import com.google.firebase.ai.type.ObjectSource import com.google.firebase.ai.type.PublicPreviewAPI import com.google.firebase.ai.type.TextPart import com.google.firebase.ai.type.content @@ -158,12 +159,12 @@ internal class OnDeviceGenerativeModelProvider( finishReason = FinishReason.STOP, finishMessage = null, groundingMetadata = null, - urlContextMetadata = null, - objectResponse = interopResponse.instance + urlContextMetadata = null ) + @Suppress("UNCHECKED_CAST") GenerateObjectResponse( GenerateContentResponse(listOf(candidate), InferenceSource.ON_DEVICE, null, null, "ondevice"), - jsonSchema + ObjectSource.FromInstance(listOf(interopResponse.instance as T)) ) } diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt index e89ce5e7251..70f1fa3285b 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Candidate.kt @@ -52,8 +52,7 @@ internal constructor( public val finishReason: FinishReason?, public val finishMessage: String?, public val groundingMetadata: GroundingMetadata?, - public val urlContextMetadata: UrlContextMetadata?, - internal val objectResponse: Any? = null + public val urlContextMetadata: UrlContextMetadata? ) { @OptIn(PublicPreviewAPI::class) diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt index 704d2e0a4b1..f255dfab2f6 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt @@ -19,6 +19,12 @@ package com.google.firebase.ai.type import kotlinx.serialization.InternalSerializationApi import kotlinx.serialization.json.Json +internal sealed interface ObjectSource { + data class FromText(val schema: JsonSchema) : ObjectSource + + data class FromInstance(val instances: List) : ObjectSource +} + /** * A [GenerateContentResponse] augmented with class information. * @@ -27,9 +33,14 @@ import kotlinx.serialization.json.Json public class GenerateObjectResponse internal constructor( public val response: GenerateContentResponse, - internal val schema: JsonSchema + internal val source: ObjectSource ) { + internal constructor( + response: GenerateContentResponse, + schema: JsonSchema + ) : this(response, ObjectSource.FromText(schema)) + /** * Deserialize a candidate (default first) and convert it into the type associated with this * response. @@ -39,21 +50,26 @@ internal constructor( * @throws SerializationException if an error occurs during deserialization */ @OptIn(InternalSerializationApi::class) - public fun getObject(candidateIndex: Int = 0): T? { - val candidate = response.candidates[candidateIndex] - if (candidate.objectResponse != null) { - @Suppress("UNCHECKED_CAST") return candidate.objectResponse as T? - } - - val deserializer = schema.getSerializer() - val text = - candidate.content.parts - .filter { !it.isThought } - .filterIsInstance() - .joinToString(" ") { it.text } - if (text.isEmpty()) { - return null + public fun getObject(candidateIndex: Int = 0): T? = + when (source) { + is ObjectSource.FromInstance -> source.instances.getOrNull(candidateIndex) + is ObjectSource.FromText -> { + val candidate = response.candidates.getOrNull(candidateIndex) + if (candidate == null) { + null + } else { + val deserializer = source.schema.getSerializer() + val text = + candidate.content.parts + .filter { !it.isThought } + .filterIsInstance() + .joinToString(" ") { it.text } + if (text.isEmpty()) { + null + } else { + Json.decodeFromString(deserializer, text) as T? + } + } + } } - return Json.decodeFromString(deserializer, text) as T? - } } From 35e128aaa14809e604d74c596f50b6e5e8ea3fcb Mon Sep 17 00:00:00 2001 From: Mila <107142260+milaGGL@users.noreply.github.com> Date: Wed, 8 Jul 2026 15:38:49 -0400 Subject: [PATCH 3/5] update naming, the on-device GenerateObjectResponse, and ObjectSource --- .../ai/ksp/SchemaSymbolProcessorVisitor.kt | 3 ++ ...seInterop.kt => GenerateObjectResponse.kt} | 8 ++-- .../ai/ondevice/interop/GenerativeModel.kt | 6 +-- .../ai/ondevice/GenerativeModelImpl.kt | 43 ++++++++++++------- .../OnDeviceGenerativeModelProvider.kt | 26 +++++------ .../ai/type/GenerateObjectResponse.kt | 6 +-- .../OnDeviceGenerativeModelProviderTests.kt | 19 +++++--- 7 files changed, 67 insertions(+), 44 deletions(-) rename ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/{GenerateObjectResponseInterop.kt => GenerateObjectResponse.kt} (73%) diff --git a/ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/SchemaSymbolProcessorVisitor.kt b/ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/SchemaSymbolProcessorVisitor.kt index 4e33e25c003..dfdb0f31c62 100644 --- a/ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/SchemaSymbolProcessorVisitor.kt +++ b/ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/SchemaSymbolProcessorVisitor.kt @@ -302,6 +302,8 @@ internal class SchemaSymbolProcessorVisitor( FileSpec.builder(packageName, companionClassName).addAnnotation(Generated::class) val classBuilder = TypeSpec.classBuilder(companionClassName).addModifiers(KModifier.DATA) + val keepAnnotation = AnnotationSpec.builder(ClassName("androidx.annotation", "Keep")).build() + classBuilder.addAnnotation(keepAnnotation) val generableAnn = classDeclaration.annotations.firstOrNull { it.shortName.getShortName() == "Generable" } @@ -318,6 +320,7 @@ internal class SchemaSymbolProcessorVisitor( val primaryConstructor = FunSpec.constructorBuilder() val toSdkBuilder = FunSpec.builder("toSdk") + .addAnnotation(keepAnnotation) .returns(ClassName(packageName, classDeclaration.simpleName.asString())) val toSdkArgs = mutableListOf() diff --git a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponseInterop.kt b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponse.kt similarity index 73% rename from ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponseInterop.kt rename to ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponse.kt index 76157fd9b7b..e461d53f4d0 100644 --- a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponseInterop.kt +++ b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponse.kt @@ -19,7 +19,9 @@ package com.google.firebase.ai.ondevice.interop /** * Represents a structured object generation response from the on-device model interop layer. * - * @property instance The generated object instance returned directly by the model engine. - * @property rawJson The raw JSON string representation of the object, if available. + * @property instances The list of generated object instances across all candidates returned + * directly by the model engine. */ -public class GenerateObjectResponseInterop(public val instance: T, public val rawJson: String) {} +public class GenerateObjectResponse(public val instances: List) { + public constructor(instance: T) : this(listOf(instance)) +} diff --git a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt index 5276b1e28c0..fa253ee83e5 100644 --- a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt +++ b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt @@ -50,14 +50,14 @@ public interface GenerativeModel { * prompt. * * @param request The input given to the model as a prompt. - * @param outputClass The target class type to which the model will generate. + * @param shadowClass The target or shadow class type to which the model will generate. * @throws [FirebaseAIOnDeviceNotAvailableException] if model is not available. * @return The structured object response generated by the model. */ public suspend fun generateObject( request: GenerateContentRequest, - outputClass: kotlin.reflect.KClass - ): GenerateObjectResponseInterop + shadowClass: kotlin.reflect.KClass + ): GenerateObjectResponse /** * Counts the number of tokens in a prompt using the model's tokenizer. diff --git a/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt b/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt index 762b0bf11ee..54754b037ef 100644 --- a/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt +++ b/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt @@ -55,18 +55,25 @@ internal class GenerativeModelImpl( override suspend fun generateObject( request: GenerateContentRequest, - outputClass: kotlin.reflect.KClass - ): com.google.firebase.ai.ondevice.interop.GenerateObjectResponseInterop = + shadowClass: kotlin.reflect.KClass + ): com.google.firebase.ai.ondevice.interop.GenerateObjectResponse = try { android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] ==========================================") android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] Calling ML Kit generateTypedContentRequest") - android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] Input SDK Class: ${outputClass.qualifiedName}") - val companionClassName = "${outputClass.java.name}_MlKitCompanion" + android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] Input SDK Class: ${shadowClass.qualifiedName}") + val companionClassName = "${shadowClass.java.name}_MlKitCompanion" val targetClass = try { Class.forName(companionClassName).kotlin } catch (e: Exception) { - outputClass + if (shadowClass.java.name.endsWith("_MlKitCompanion")) { + shadowClass + } else { + throw IllegalArgumentException( + "Shadow class $companionClassName not found for structured output. Ensure KSP processor is running and generated the MlKit companion class.", + e + ) + } } android.util.Log.i( "MLKIT_IO", @@ -85,24 +92,28 @@ internal class GenerativeModelImpl( outputClass = targetClass as kotlin.reflect.KClass, includeSchemaInPrompt = true ) - val response = mlkitModel.generateContent(mlkitRequest) - val candidate = response.candidates.firstOrNull() - val mlkitInstance = candidate?.response + val typedResponse = mlkitModel.generateContent(mlkitRequest) + val mlkitInstances = typedResponse.candidates.mapNotNull { it.response } + if (mlkitInstances.isEmpty()) { + throw GenAiException("No response candidate returned by ML Kit", null, 0) + } android.util.Log.i( "MLKIT_IO", - "[SDK_BRIDGE] ML Kit returned structured instance: $mlkitInstance (${mlkitInstance?.javaClass?.name})" + "[SDK_BRIDGE] ML Kit returned structured instances count: ${mlkitInstances.size}" ) @Suppress("UNCHECKED_CAST") - val sdkInstance = - if (mlkitInstance != null && mlkitInstance.javaClass.name.endsWith("_MlKitCompanion")) { - val toSdkMethod = mlkitInstance.javaClass.getMethod("toSdk") - toSdkMethod.invoke(mlkitInstance) as T - } else { - mlkitInstance as T + val sdkInstances = + mlkitInstances.map { mlkitInstance -> + if (mlkitInstance.javaClass.name.endsWith("_MlKitCompanion")) { + val toSdkMethod = mlkitInstance.javaClass.getMethod("toSdk") + toSdkMethod.invoke(mlkitInstance) as T + } else { + mlkitInstance as T + } } - com.google.firebase.ai.ondevice.interop.GenerateObjectResponseInterop(sdkInstance, "") + com.google.firebase.ai.ondevice.interop.GenerateObjectResponse(sdkInstances) } catch (e: GenAiException) { throw getMappingException(e) } diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt index 95d22234bee..15640c0caf0 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt @@ -151,20 +151,22 @@ internal class OnDeviceGenerativeModelProvider( val request = buildOnDeviceGenerateContentRequest(prompt) val interopResponse = onDeviceModel.generateObject(request, jsonSchema.clazz) - val candidate = - Candidate( - content = content { text(interopResponse.rawJson) }, - safetyRatings = emptyList(), - citationMetadata = null, - finishReason = FinishReason.STOP, - finishMessage = null, - groundingMetadata = null, - urlContextMetadata = null - ) + val candidates = + interopResponse.instances.map { + Candidate( + content = content { text("") }, + safetyRatings = emptyList(), + citationMetadata = null, + finishReason = FinishReason.STOP, + finishMessage = null, + groundingMetadata = null, + urlContextMetadata = null + ) + } @Suppress("UNCHECKED_CAST") GenerateObjectResponse( - GenerateContentResponse(listOf(candidate), InferenceSource.ON_DEVICE, null, null, "ondevice"), - ObjectSource.FromInstance(listOf(interopResponse.instance as T)) + GenerateContentResponse(candidates, InferenceSource.ON_DEVICE, null, null, "ondevice"), + ObjectSource.FromInstance(interopResponse.instances as List) ) } diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt index f255dfab2f6..5228805938a 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt @@ -20,7 +20,7 @@ import kotlinx.serialization.InternalSerializationApi import kotlinx.serialization.json.Json internal sealed interface ObjectSource { - data class FromText(val schema: JsonSchema) : ObjectSource + data class FromSchema(val schema: JsonSchema) : ObjectSource data class FromInstance(val instances: List) : ObjectSource } @@ -39,7 +39,7 @@ internal constructor( internal constructor( response: GenerateContentResponse, schema: JsonSchema - ) : this(response, ObjectSource.FromText(schema)) + ) : this(response, ObjectSource.FromSchema(schema)) /** * Deserialize a candidate (default first) and convert it into the type associated with this @@ -53,7 +53,7 @@ internal constructor( public fun getObject(candidateIndex: Int = 0): T? = when (source) { is ObjectSource.FromInstance -> source.instances.getOrNull(candidateIndex) - is ObjectSource.FromText -> { + is ObjectSource.FromSchema -> { val candidate = response.candidates.getOrNull(candidateIndex) if (candidate == null) { null diff --git a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/generativeModel/OnDeviceGenerativeModelProviderTests.kt b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/generativeModel/OnDeviceGenerativeModelProviderTests.kt index 32e9eb33d4e..e548cf492ed 100644 --- a/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/generativeModel/OnDeviceGenerativeModelProviderTests.kt +++ b/ai-logic/firebase-ai/src/test/java/com/google/firebase/ai/generativeModel/OnDeviceGenerativeModelProviderTests.kt @@ -107,22 +107,27 @@ internal class OnDeviceGenerativeModelProviderTests { response.inferenceSource shouldBe InferenceSource.ON_DEVICE } + private data class TestMovieReview(val title: String, val rating: Int) + @Test fun `generateObject delegates to onDeviceModel`(): Unit = runBlocking { coEvery { onDeviceModel.isAvailable() } returns true - val schema = JsonSchema.string() + val dummyInstance = TestMovieReview("Inception", 5) + val schema: JsonSchema = mockk { + every { clazz } returns TestMovieReview::class + } val interopResponse = - com.google.firebase.ai.ondevice.interop.GenerateObjectResponseInterop( - instance = "test object", - rawJson = "\"test object\"" + com.google.firebase.ai.ondevice.interop.GenerateObjectResponse( + instance = dummyInstance ) - coEvery { onDeviceModel.generateObject(any(), any>()) } returns - interopResponse + coEvery { + onDeviceModel.generateObject(any(), any>()) + } returns interopResponse val response = provider.generateObject(schema, prompt) response.response.inferenceSource shouldBe InferenceSource.ON_DEVICE - response.getObject() shouldBe "test object" + response.getObject() shouldBe dummyInstance } @Test From 60b3c1222faae2e9b20aca8ff1da099d44ea756a Mon Sep 17 00:00:00 2001 From: Mila <107142260+milaGGL@users.noreply.github.com> Date: Thu, 9 Jul 2026 13:26:36 -0400 Subject: [PATCH 4/5] update GenerateObjectResponse --- .../OnDeviceGenerativeModelProvider.kt | 3 +- .../ai/type/GenerateObjectResponse.kt | 67 +++++++++++-------- 2 files changed, 39 insertions(+), 31 deletions(-) diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt index 15640c0caf0..439a5b0521d 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/generativemodel/OnDeviceGenerativeModelProvider.kt @@ -33,7 +33,6 @@ import com.google.firebase.ai.type.GenerateContentResponse import com.google.firebase.ai.type.GenerateObjectResponse import com.google.firebase.ai.type.ImagePart import com.google.firebase.ai.type.JsonSchema -import com.google.firebase.ai.type.ObjectSource import com.google.firebase.ai.type.PublicPreviewAPI import com.google.firebase.ai.type.TextPart import com.google.firebase.ai.type.content @@ -166,7 +165,7 @@ internal class OnDeviceGenerativeModelProvider( @Suppress("UNCHECKED_CAST") GenerateObjectResponse( GenerateContentResponse(candidates, InferenceSource.ON_DEVICE, null, null, "ondevice"), - ObjectSource.FromInstance(interopResponse.instances as List) + instances = interopResponse.instances as List ) } diff --git a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt index 5228805938a..af68e25044a 100644 --- a/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt +++ b/ai-logic/firebase-ai/src/main/kotlin/com/google/firebase/ai/type/GenerateObjectResponse.kt @@ -19,12 +19,6 @@ package com.google.firebase.ai.type import kotlinx.serialization.InternalSerializationApi import kotlinx.serialization.json.Json -internal sealed interface ObjectSource { - data class FromSchema(val schema: JsonSchema) : ObjectSource - - data class FromInstance(val instances: List) : ObjectSource -} - /** * A [GenerateContentResponse] augmented with class information. * @@ -33,13 +27,23 @@ internal sealed interface ObjectSource { public class GenerateObjectResponse internal constructor( public val response: GenerateContentResponse, - internal val source: ObjectSource + internal val schema: JsonSchema?, + internal var instances: MutableList? ) { internal constructor( response: GenerateContentResponse, schema: JsonSchema - ) : this(response, ObjectSource.FromSchema(schema)) + ) : this(response, schema = schema, instances = null) + + internal constructor( + response: GenerateContentResponse, + instances: List + ) : this( + response, + schema = null, + instances = (instances as? MutableList) ?: instances.toMutableList() + ) /** * Deserialize a candidate (default first) and convert it into the type associated with this @@ -50,26 +54,31 @@ internal constructor( * @throws SerializationException if an error occurs during deserialization */ @OptIn(InternalSerializationApi::class) - public fun getObject(candidateIndex: Int = 0): T? = - when (source) { - is ObjectSource.FromInstance -> source.instances.getOrNull(candidateIndex) - is ObjectSource.FromSchema -> { - val candidate = response.candidates.getOrNull(candidateIndex) - if (candidate == null) { - null - } else { - val deserializer = source.schema.getSerializer() - val text = - candidate.content.parts - .filter { !it.isThought } - .filterIsInstance() - .joinToString(" ") { it.text } - if (text.isEmpty()) { - null - } else { - Json.decodeFromString(deserializer, text) as T? - } - } - } + public fun getObject(candidateIndex: Int = 0): T? { + // 1. Fast Path / Cache Hit: Return immediately if already resolved or in memory + instances?.getOrNull(candidateIndex)?.let { + return it + } + + // 2. Cache Miss (Cloud response on first access): Deserialize using schema + if (schema == null) return null + val candidate = response.candidates.getOrNull(candidateIndex) ?: return null + val text = + candidate.content.parts + .filter { !it.isThought } + .filterIsInstance() + .joinToString(" ") { it.text } + if (text.isEmpty()) return null + + val deserialized = Json.decodeFromString(schema.getSerializer(), text) as T? + + // 3. Save to instances list for future accesses (lazy loading) and return + if (instances == null) { + instances = MutableList(response.candidates.size) { null } + } + if (candidateIndex < (instances?.size ?: 0)) { + instances?.set(candidateIndex, deserialized) } + return deserialized + } } From 8d670e89986af1468e29ec521a1d7c6b1f1dbfff Mon Sep 17 00:00:00 2001 From: Mila <107142260+milaGGL@users.noreply.github.com> Date: Fri, 10 Jul 2026 17:01:40 -0400 Subject: [PATCH 5/5] update the name of shadowClass property --- .../firebase/ai/ondevice/interop/GenerativeModel.kt | 5 +++-- .../google/firebase/ai/ondevice/GenerativeModelImpl.kt | 10 +++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt index fa253ee83e5..71634edff6b 100644 --- a/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt +++ b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerativeModel.kt @@ -50,13 +50,14 @@ public interface GenerativeModel { * prompt. * * @param request The input given to the model as a prompt. - * @param shadowClass The target or shadow class type to which the model will generate. + * @param schemaClass The target SDK data class (`T`, e.g. `MovieReview::class`), from which the + * underlying engine dynamically resolves and maps the KSP shadow companion class at runtime. * @throws [FirebaseAIOnDeviceNotAvailableException] if model is not available. * @return The structured object response generated by the model. */ public suspend fun generateObject( request: GenerateContentRequest, - shadowClass: kotlin.reflect.KClass + schemaClass: kotlin.reflect.KClass ): GenerateObjectResponse /** diff --git a/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt b/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt index 54754b037ef..e130a7a91b4 100644 --- a/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt +++ b/ai-logic/firebase-ai-ondevice/src/main/kotlin/com/google/firebase/ai/ondevice/GenerativeModelImpl.kt @@ -55,19 +55,19 @@ internal class GenerativeModelImpl( override suspend fun generateObject( request: GenerateContentRequest, - shadowClass: kotlin.reflect.KClass + schemaClass: kotlin.reflect.KClass ): com.google.firebase.ai.ondevice.interop.GenerateObjectResponse = try { android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] ==========================================") android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] Calling ML Kit generateTypedContentRequest") - android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] Input SDK Class: ${shadowClass.qualifiedName}") - val companionClassName = "${shadowClass.java.name}_MlKitCompanion" + android.util.Log.i("MLKIT_IO", "[SDK_BRIDGE] Input SDK Class: ${schemaClass.qualifiedName}") + val companionClassName = "${schemaClass.java.name}_MlKitCompanion" val targetClass = try { Class.forName(companionClassName).kotlin } catch (e: Exception) { - if (shadowClass.java.name.endsWith("_MlKitCompanion")) { - shadowClass + if (schemaClass.java.name.endsWith("_MlKitCompanion")) { + schemaClass } else { throw IllegalArgumentException( "Shadow class $companionClassName not found for structured output. Ensure KSP processor is running and generated the MlKit companion class.",