[AI] support on-device structured output

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 {
@@ -142,7 +153,18 @@ internal class SchemaSymbolProcessorVisitor(
         builder.addStatement("JsonSchema.double(").indent()
       }
       "kotlin.String" -> {
-        builder.addStatement("JsonSchema.string(").indent()
+        if (!guideValues.enumValues.isNullOrEmpty()) {
+          builder
+            .addStatement("JsonSchema.enumeration(")
+            .indent()
+            .addStatement("values = listOf(")
+            .indent()
+            .addStatement(guideValues.enumValues.joinToString { "\"$it\"" })
+            .unindent()
+            .addStatement("),")
+        } else {
+          builder.addStatement("JsonSchema.string(").indent()
+        }
       }
       "kotlin.collections.List" -> {
 
@@ -155,7 +177,12 @@ internal class SchemaSymbolProcessorVisitor(
           throw RuntimeException()
         }
         val listParamCodeBlock =
-          generateCodeBlockForSchema(type = listTypeParam.resolve(), parentType = type)
+          generateCodeBlockForSchema(
+            type = listTypeParam.resolve(),
+            parentType = type,
+            guideAnnotation =
+              if (!guideValues.enumValues.isNullOrEmpty()) guideAnnotation else null,
+          )
         builder
           .addStatement("JsonSchema.array(")
           .indent()
@@ -248,4 +275,124 @@ 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" }
+  }
+
+  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)) {
+      type.arguments.firstOrNull()?.type?.resolve()?.let { argType ->
+        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.schema.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<String>()
+
+    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.schema.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)
+        if (!guideValues.enumValues.isNullOrEmpty()) {
+          val enumElements = guideValues.enumValues.joinToString { "%S" }
+          mlkitGuideBuilder.addMember(
+            "enumValues = arrayOf($enumElements)",
+            *guideValues.enumValues.toTypedArray()
+          )
+        }
+        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()
+  }
 }

ai-logic/firebase-ai-ksp-processor/src/main/kotlin/com/google/firebase/ai/ksp/Util.kt

@@ -32,7 +32,8 @@ internal data class GuideValues(
   val minItems: Int?,
   val maxItems: Int?,
   val format: String?,
-  val description: String?
+  val description: String?,
+  val enumValues: List<String>? = null,
 )
 
 internal fun getGuideValuesFromAnnotation(
@@ -45,7 +46,8 @@ internal fun getGuideValuesFromAnnotation(
     minItems = getIntFromAnnotation(guideAnnotation, "minItems"),
     maxItems = getIntFromAnnotation(guideAnnotation, "maxItems"),
     format = getStringFromAnnotation(guideAnnotation, "format"),
-    description = description
+    description = description,
+    enumValues = getStringListFromAnnotation(guideAnnotation, "enumValues"),
   )
 
 internal fun getDescriptionFromAnnotations(
@@ -103,6 +105,27 @@ internal fun getStringFromAnnotation(
   return guidePropertyStringValue
 }
 
+internal fun getStringListFromAnnotation(
+  guideAnnotation: KSAnnotation?,
+  listName: String,
+): List<String>? {
+  val rawValue =
+    guideAnnotation
+      ?.arguments
+      ?.firstOrNull { it.name?.getShortName()?.equals(listName) == true }
+      ?.value
+  val list =
+    when (rawValue) {
+      is List<*> -> rawValue.mapNotNull { it as? String }
+      is Array<*> -> rawValue.mapNotNull { it as? String }
+      else -> null
+    }
+  if (list.isNullOrEmpty()) {
+    return null
+  }
+  return list
+}
+
 internal fun extractBaseKdoc(kdoc: String): String? {
   return baseKdocRegex.matchEntire(kdoc)?.groups?.get(1)?.value?.trim().let {
     if (it.isNullOrEmpty()) null else it

ai-logic/firebase-ai-ksp-processor/test-app/src/main/kotlin/com/google/firebase/testing/processor/SchemaTest.kt

@@ -32,14 +32,17 @@ data class RootSchemaTestClass(
   val listTest: List<Int>,
   @Guide(description = "most likely true, very rarely false") val booleanTest: Boolean,
   val stringTest: String,
-  val objTest: SecondarySchemaTestClass,
+  val nestedSchemaTest: SecondarySchemaTestClass,
   val enumTest: EnumTest,
+  @Guide(enumValues = ["NORTH", "SOUTH", "EAST", "WEST"]) val stringEnumTest: String,
 ) {
   companion object
 }
 
-/** class kdoc should be used if property kdocs aren't present */
-data class SecondarySchemaTestClass(val testInt: Int)
+@Generable
+data class SecondarySchemaTestClass(val testInt: Int) {
+  companion object
+}
 
 enum class EnumTest {
   A,

ai-logic/firebase-ai-ksp-processor/test-app/src/test/kotlin/com/google/firebase/testing/processor/FirebaseKspProcessorTest.kt

@@ -77,13 +77,18 @@ class FirebaseKspProcessorTest {
     assertThat(enumSchema.title).isEqualTo("enumTest")
     assertThat(enumSchema.nullable).isEqualTo(false)
 
-    assertThat(rootSchema.properties?.get("objTest")).isNotNull()
-    val objSchema = rootSchema.properties?.get("objTest")!!
+    assertThat(rootSchema.properties?.get("nestedSchemaTest")).isNotNull()
+    val objSchema = rootSchema.properties?.get("nestedSchemaTest")!!
     assertThat(objSchema.clazz).isEqualTo(SecondarySchemaTestClass::class)
     assertThat(objSchema.properties).isNotNull()
-    assertThat(objSchema.description)
-      .isEqualTo("class kdoc should be used if property kdocs aren't present")
-    assertThat(objSchema.title).isEqualTo("objTest")
+    assertThat(objSchema.description).isNull()
+    assertThat(objSchema.title).isEqualTo("nestedSchemaTest")
     assertThat(objSchema.nullable).isEqualTo(false)
+
+    assertThat(rootSchema.properties?.get("stringEnumTest")).isNotNull()
+    val stringEnumSchema = rootSchema.properties?.get("stringEnumTest")!!
+    assertThat(stringEnumSchema.enum).isEqualTo(listOf("NORTH", "SOUTH", "EAST", "WEST"))
+    assertThat(stringEnumSchema.title).isEqualTo("stringEnumTest")
+    assertThat(stringEnumSchema.nullable).isEqualTo(false)
   }
 }

ai-logic/firebase-ai-ksp-processor/test-app/test-app.gradle.kts

@@ -28,7 +28,7 @@ android {
   compileSdk = 36
   defaultConfig {
     applicationId = "com.google.firebase.testing.processor"
-    minSdk = 23
+    minSdk = 26
     targetSdk = 36
     versionCode = 1
     versionName = "1.0"
@@ -41,10 +41,18 @@ android {
   }
 }
 
-kotlin { compilerOptions { jvmTarget = JvmTarget.JVM_1_8 } }
+kotlin {
+  compilerOptions {
+    jvmTarget = JvmTarget.JVM_1_8
+    freeCompilerArgs.add("-Xskip-metadata-version-check")
+  }
+}
 
 dependencies {
   implementation(project(":ai-logic:firebase-ai"))
+  implementation(project(":ai-logic:firebase-ai-ondevice"))
+
+  implementation(libs.genai.schema)
   ksp(project(":ai-logic:firebase-ai-ksp-processor"))
 
   implementation("com.google.firebase:firebase-common:22.0.0")

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

@@ -105,6 +105,13 @@ package com.google.firebase.ai.ondevice.interop {
     property public final String? modelVersion;
   }
 
+  public final class GenerateObjectResponse<T> {
+    ctor public GenerateObjectResponse(java.util.List<? extends T> instances);
+    ctor public GenerateObjectResponse(T instance);
+    method public java.util.List<T> getInstances();
+    property public final java.util.List<T> instances;
+  }
+
   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();
@@ -117,6 +124,7 @@ package com.google.firebase.ai.ondevice.interop {
     method public kotlinx.coroutines.flow.Flow<com.google.firebase.ai.ondevice.interop.DownloadStatusInterop> download();
     method public suspend Object? generateContent(com.google.firebase.ai.ondevice.interop.GenerateContentRequest request, kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.interop.GenerateContentResponse>);
     method public kotlinx.coroutines.flow.Flow<com.google.firebase.ai.ondevice.interop.GenerateContentResponse> generateContentStream(com.google.firebase.ai.ondevice.interop.GenerateContentRequest request);
+    method public suspend <T> Object? generateObject(com.google.firebase.ai.ondevice.interop.GenerateContentRequest request, kotlin.reflect.KClass<T> schemaClass, kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.interop.GenerateObjectResponse<T>>);
     method public suspend Object? getBaseModelName(kotlin.coroutines.Continuation<? super java.lang.String>);
     method public suspend Object? getTokenLimit(kotlin.coroutines.Continuation<? super java.lang.Integer>);
     method public suspend Object? isAvailable(kotlin.coroutines.Continuation<? super java.lang.Boolean>);

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<T>(public val instances: List<T>) {
+  public constructor(instance: T) : this(listOf(instance))
+}

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 <T : Any> generateObject(
+    request: GenerateContentRequest,
+    schemaClass: kotlin.reflect.KClass<T>
+  ): GenerateObjectResponse<T>
+
   /**
    * Counts the number of tokens in a prompt using the model's tokenizer.
    *