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..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 @@ -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,122 @@ 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 keepAnnotation = AnnotationSpec.builder(ClassName("androidx.annotation", "Keep")).build() + classBuilder.addAnnotation(keepAnnotation) + + 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") + .addAnnotation(keepAnnotation) + .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/GenerateObjectResponse.kt b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponse.kt new file mode 100644 index 00000000000..e461d53f4d0 --- /dev/null +++ b/ai-logic/firebase-ai-ondevice-interop/src/main/kotlin/com/google/firebase/ai/ondevice/interop/GenerateObjectResponse.kt @@ -0,0 +1,27 @@ +/* + * 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 instances The list of generated object instances across all candidates returned + * directly by the model engine. + */ +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 e7e03b7e37f..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 @@ -45,6 +45,21 @@ 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 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, + schemaClass: 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/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..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 @@ -53,6 +53,71 @@ internal class GenerativeModelImpl( throw getMappingException(e) } + override suspend fun generateObject( + request: GenerateContentRequest, + 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: ${schemaClass.qualifiedName}") + val companionClassName = "${schemaClass.java.name}_MlKitCompanion" + val targetClass = + try { + Class.forName(companionClassName).kotlin + } catch (e: Exception) { + 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.", + e + ) + } + } + 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 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 instances count: ${mlkitInstances.size}" + ) + + @Suppress("UNCHECKED_CAST") + 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.GenerateObjectResponse(sdkInstances) + } 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..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 @@ -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,27 @@ 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 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(candidates, InferenceSource.ON_DEVICE, null, null, "ondevice"), + 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 59fe8e5db02..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 @@ -27,9 +27,24 @@ import kotlinx.serialization.json.Json public class GenerateObjectResponse internal constructor( public val response: GenerateContentResponse, - internal val schema: JsonSchema + internal val schema: JsonSchema?, + internal var instances: MutableList? ) { + internal constructor( + response: GenerateContentResponse, + schema: JsonSchema + ) : 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 * response. @@ -40,17 +55,30 @@ internal constructor( */ @OptIn(InternalSerializationApi::class) public fun getObject(candidateIndex: Int = 0): T? { - val candidate = response.candidates[candidateIndex] + // 1. Fast Path / Cache Hit: Return immediately if already resolved or in memory + instances?.getOrNull(candidateIndex)?.let { + return it + } - val deserializer = schema.getSerializer() + // 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 + 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 Json.decodeFromString(deserializer, text) as T? + return deserialized } } 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..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,12 +107,27 @@ internal class OnDeviceGenerativeModelProviderTests { response.inferenceSource shouldBe InferenceSource.ON_DEVICE } + private data class TestMovieReview(val title: String, val rating: Int) + @Test - fun `generateObject always throws FirebaseAIException`(): Unit = runBlocking { - val schema = mockk>() + fun `generateObject delegates to onDeviceModel`(): Unit = runBlocking { + coEvery { onDeviceModel.isAvailable() } returns true + val dummyInstance = TestMovieReview("Inception", 5) + val schema: JsonSchema = mockk { + every { clazz } returns TestMovieReview::class + } + val interopResponse = + com.google.firebase.ai.ondevice.interop.GenerateObjectResponse( + instance = dummyInstance + ) + 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 dummyInstance } @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") } } }