Skip to content

[AI] support on-device structured output - #8395

Open
milaGGL wants to merge 14 commits into
mainfrom
mila-support-structured-ouput-locally
Open

[AI] support on-device structured output#8395
milaGGL wants to merge 14 commits into
mainfrom
mila-support-structured-ouput-locally

Conversation

@milaGGL

@milaGGL milaGGL commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

This PR introduces Structured Output support on device, allowing developers to generate strongly-typed objects using ML Kit's backend. It defines the public API (GenerateObjectResponse and generateObject()), handles the necessary Kotlin Symbol Processing (KSP) logic for schema translation.

How to use On-Device Structured Output

  1. Define a data class annotated with @generable (and @Guide). Note: You must provide an empty companion object so the KSP processor can generate the necessary schema extensions.
  2. Call generateObject
import com.google.firebase.ai.annotations.Generable
@Generable
data class Recipe(
    val recipeName: String,
    @Guide(description = "A list of ingredients", minItems = 2)
    val ingredients: List<String>,
    @Guide(description = "Preparation time in minutes", minimum = 1.0)
    val prepTimeMinutes: Int
) {
    companion object
}

val schema = Recipe.firebaseAISchema()
val prompt = "Give me a recipe for chocolate chip cookies"

val response = model.generateObject(schema, prompt)
val recipe: Recipe? = response.getObject()
recipe?.let {
    println("Recipe: ${it.recipeName}")
    println("Ingredients: ${it.ingredients.joinToString()}")
    println("Prep time: ${it.prepTimeMinutes} mins")
}

@github-actions

github-actions Bot commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

📝 PRs merging into main branch

Our main branch should always be in a releasable state. If you are working on a larger change, or if you don't want this change to see the light of the day just yet, consider using a feature branch first, and only merge into the main branch when the code complete and ready to be released.

@google-oss-bot

Copy link
Copy Markdown
Collaborator
1 Warning
⚠️ Did you forget to add a changelog entry? (Add the 'no-changelog' label to the PR to silence this warning.)

Generated by 🚫 Danger

@gemini-code-assist

Copy link
Copy Markdown
Contributor
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

@google-oss-bot

Copy link
Copy Markdown
Collaborator

The public api surface has changed for the subproject ai-logic_firebase-ai-ondevice-interop:
error: Added class com.google.firebase.ai.ondevice.interop.GenerateObjectResponse [AddedClass]
error: Added method com.google.firebase.ai.ondevice.interop.GenerativeModel.generateObject(com.google.firebase.ai.ondevice.interop.GenerateContentRequest,kotlin.reflect.KClass,kotlin.coroutines.Continuation<? super com.google.firebase.ai.ondevice.interop.GenerateObjectResponse>) [AddedAbstractMethod]

Please update the api.txt files for the subprojects being affected by this change by running ./gradlew ${subproject}:generateApiTxtFile. Also perform a major/minor bump accordingly.

@milaGGL
milaGGL marked this pull request as ready for review July 28, 2026 14:25
@milaGGL milaGGL changed the title initial implementation [AI] support on-device structured output Jul 28, 2026
@milaGGL

milaGGL commented Jul 28, 2026

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for structured object generation (generateObject) in the on-device model interop and implementation layers, including updates to the KSP processor to generate ML Kit companion classes and support enum values in the @Guide annotation. Feedback on these changes highlights a type-safety issue in GenerateObjectResponse.kt regarding MutableList invariance, as well as an opportunity to replace unsafe non-null assertions (!!) with idiomatic safe calls in SchemaSymbolProcessorVisitor.kt.

@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@firebase firebase deleted a comment from google-oss-bot Jul 28, 2026
@milaGGL

milaGGL commented Jul 28, 2026

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for structured object generation on-device by integrating with ML Kit's schema-based generation. It updates the KSP processor to generate companion classes (_MlKitCompanion) for @Generable classes, adds the generateObject API to the on-device interop and implementation layers, and updates GenerateObjectResponse to support lazy-loading and caching of deserialized instances. The review feedback highlights several important improvements: handling nested classes correctly when resolving companion class names, refining the caching logic in GenerateObjectResponse to prevent redundant deserialization of null values, safely escaping enum values using KotlinPoet's %S specifier, and propagating KDoc-based class and property descriptions to the generated companion annotations.

schemaClass: kotlin.reflect.KClass<T>
): com.google.firebase.ai.ondevice.interop.GenerateObjectResponse<T> =
try {
val companionClassName = "${schemaClass.java.name}_MlKitCompanion"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

If schemaClass is a nested class (e.g., Outer.Inner), schemaClass.java.name will be Outer$Inner. However, KSP generates the companion class as a top-level class named Inner_MlKitCompanion in the same package. Using ${schemaClass.java.name}_MlKitCompanion will resolve to Outer$Inner_MlKitCompanion, which does not exist and will cause a runtime IllegalArgumentException. Construct the companion class name using the package name and the simple name of the class to support nested classes correctly.

      val packageName = schemaClass.java.name.substringBeforeLast('.', "")
      val simpleName = schemaClass.java.simpleName
      val companionClassName = if (packageName.isEmpty()) "${simpleName}_MlKitCompanion" else "$packageName.${simpleName}_MlKitCompanion"

Comment on lines 27 to +34
public class GenerateObjectResponse<T : Any>
internal constructor(
public val response: GenerateContentResponse,
internal val schema: JsonSchema<T>
internal val schema: JsonSchema<T>?,
internal var instances: MutableList<T?>?
) {

internal constructor(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Introduce a resolvedIndices set to track which candidate indices have already been resolved. This is necessary to distinguish between an unresolved candidate and a candidate that has been resolved to null (e.g., due to empty text or deserialization failure), preventing redundant deserialization attempts.

Suggested change
public class GenerateObjectResponse<T : Any>
internal constructor(
public val response: GenerateContentResponse,
internal val schema: JsonSchema<T>
internal val schema: JsonSchema<T>?,
internal var instances: MutableList<T?>?
) {
internal constructor(
public class GenerateObjectResponse<T : Any>
internal constructor(
public val response: GenerateContentResponse,
internal val schema: JsonSchema<T>?,
internal var instances: MutableList<T?>?
) {
private val resolvedIndices = HashSet<Int>()
internal constructor(

Comment on lines +57 to +60
// 1. Fast Path / Cache Hit: Return immediately if already resolved or in memory
insts?.getOrNull(candidateIndex)?.let {
return it
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Check if the candidate index has already been resolved or if we are in on-device mode (where schema is null and instances are pre-populated), and return the cached value directly. This avoids skipping the cache when the resolved value is null.

Suggested change
// 1. Fast Path / Cache Hit: Return immediately if already resolved or in memory
insts?.getOrNull(candidateIndex)?.let {
return it
}
// 1. Fast Path / Cache Hit: Return immediately if already resolved or in memory
if (insts != null && (schema == null || candidateIndex in resolvedIndices)) {
return insts.getOrNull(candidateIndex)
}

Comment on lines +86 to +88
if (candidateIndex < currentInstances.size) {
currentInstances[candidateIndex] = deserialized
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Mark the candidate index as resolved when saving the deserialized value to the cache.

Suggested change
if (candidateIndex < currentInstances.size) {
currentInstances[candidateIndex] = deserialized
}
if (candidateIndex < currentInstances.size) {
currentInstances[candidateIndex] = deserialized
resolvedIndices.add(candidateIndex)
}

.indent()
.addStatement("values = listOf(")
.indent()
.addStatement(guideValues.enumValues.joinToString { "\"$it\"" })

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using joinToString { "\"$it\"" } to embed raw string values directly into the generated code block can result in invalid Kotlin code if any enum value contains double quotes or other special characters. Use KotlinPoet's %S format specifier to safely escape and format the string values.

Suggested change
.addStatement(guideValues.enumValues.joinToString { "\"$it\"" })
.addStatement(
guideValues.enumValues.joinToString { "%S" },
*guideValues.enumValues.toTypedArray()
)

Comment on lines +314 to +315
val packageName = classDeclaration.packageName.asString()
val companionClassName = "${classDeclaration.simpleName.asString()}_MlKitCompanion"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Extract the property KDocs once at the beginning of the function so they can be used as fallbacks for property descriptions in the generated companion class.

Suggested change
val packageName = classDeclaration.packageName.asString()
val companionClassName = "${classDeclaration.simpleName.asString()}_MlKitCompanion"
val packageName = classDeclaration.packageName.asString()
val companionClassName = "${classDeclaration.simpleName.asString()}_MlKitCompanion"
val propertyDocs = extractPropertyKdocs(classDeclaration.docString ?: "")

Comment on lines +323 to +325
val generableAnn =
classDeclaration.annotations.firstOrNull { it.shortName.getShortName() == "Generable" }
val classDesc = getStringFromAnnotation(generableAnn, "description")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Class descriptions documented via KDocs are not propagated to the generated companion class's @Generable annotation. Since ML Kit's on-device structured output engine relies on these annotations on the companion class to build the schema, any KDoc-based class descriptions will be lost on-device. We should fall back to the base KDoc description if the annotation description is missing.

Suggested change
val generableAnn =
classDeclaration.annotations.firstOrNull { it.shortName.getShortName() == "Generable" }
val classDesc = getStringFromAnnotation(generableAnn, "description")
val generableAnn =
classDeclaration.annotations.firstOrNull { it.shortName.getShortName() == "Generable" }
val classDesc = getStringFromAnnotation(generableAnn, "description") ?: extractBaseKdoc(classDeclaration.docString ?: "")

Comment on lines +348 to +351
val guideAnn = property.annotations.firstOrNull { it.shortName.getShortName() == "Guide" }
if (guideAnn != null) {
val guideValues =
getGuideValuesFromAnnotation(guideAnn, getStringFromAnnotation(guideAnn, "description"))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Use the extracted property KDocs as a fallback description for the @Guide annotation on the companion class property, and ensure the @Guide annotation is generated if a description exists even if the original property lacked a @Guide annotation.

Suggested change
val guideAnn = property.annotations.firstOrNull { it.shortName.getShortName() == "Guide" }
if (guideAnn != null) {
val guideValues =
getGuideValuesFromAnnotation(guideAnn, getStringFromAnnotation(guideAnn, "description"))
val guideAnn = property.annotations.firstOrNull { it.shortName.getShortName() == "Guide" }
val propDescription = getStringFromAnnotation(guideAnn, "description") ?: propertyDocs[propName]
if (guideAnn != null || !propDescription.isNullOrEmpty()) {
val guideValues =
getGuideValuesFromAnnotation(guideAnn, propDescription)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants