Initial implementation to API spec

firebase-vertexai/CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased
 - [changed] Internal improvements to correctly handle empty model responses.
+- [feature] Added support for generating images with Imagen models.
 
 
 # 16.0.2

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/FirebaseVertexAI.kt

@@ -24,6 +24,8 @@ import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.inject.Provider
 import com.google.firebase.vertexai.type.Content
 import com.google.firebase.vertexai.type.GenerationConfig
+import com.google.firebase.vertexai.type.ImagenGenerationConfig
+import com.google.firebase.vertexai.type.ImagenSafetySettings
 import com.google.firebase.vertexai.type.InvalidLocationException
 import com.google.firebase.vertexai.type.RequestOptions
 import com.google.firebase.vertexai.type.SafetySetting
@@ -79,6 +81,36 @@ internal constructor(
     )
   }
 
+  /**
+   * Instantiates a new [ImagenModel] given the provided parameters.
+   *
+   * @param modelName The name of the model to use, for example `"imagen-3.0-generate-001"`.
+   * @param generationConfig The configuration parameters to use for image generation.
+   * @param safetySettings The safety bounds the model will abide by during image generation.
+   * @param requestOptions Configuration options for sending requests to the backend.
+   * @return The initialized [ImagenModel] instance.
+   */
+  @JvmOverloads
+  public fun imagenModel(
+    modelName: String,
+    generationConfig: ImagenGenerationConfig? = null,
+    safetySettings: ImagenSafetySettings? = null,
+    requestOptions: RequestOptions = RequestOptions(),
+  ): ImagenModel {
+    if (location.trim().isEmpty() || location.contains("/")) {
+      throw InvalidLocationException(location)
+    }
+    return ImagenModel(
+      "projects/${firebaseApp.options.projectId}/locations/${location}/publishers/google/models/${modelName}",
+      firebaseApp.options.apiKey,
+      generationConfig,
+      safetySettings,
+      requestOptions,
+      appCheckProvider.get(),
+      internalAuthProvider.get(),
+    )
+  }
+
   public companion object {
     /** The [FirebaseVertexAI] instance for the default [FirebaseApp] */
     @JvmStatic

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/GenerativeModel.kt

@@ -17,13 +17,12 @@
 package com.google.firebase.vertexai
 
 import android.graphics.Bitmap
-import android.util.Log
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.vertexai.common.APIController
+import com.google.firebase.vertexai.common.AppCheckHeaderProvider
 import com.google.firebase.vertexai.common.CountTokensRequest
 import com.google.firebase.vertexai.common.GenerateContentRequest
-import com.google.firebase.vertexai.common.HeaderProvider
 import com.google.firebase.vertexai.type.Content
 import com.google.firebase.vertexai.type.CountTokensResponse
 import com.google.firebase.vertexai.type.FinishReason
@@ -38,12 +37,9 @@ import com.google.firebase.vertexai.type.SerializationException
 import com.google.firebase.vertexai.type.Tool
 import com.google.firebase.vertexai.type.ToolConfig
 import com.google.firebase.vertexai.type.content
-import kotlin.time.Duration
-import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.catch
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.tasks.await
 
 /**
  * Represents a multimodal model (like Gemini), capable of generating content based on various input
@@ -57,10 +53,8 @@ internal constructor(
   private val tools: List<Tool>? = null,
   private val toolConfig: ToolConfig? = null,
   private val systemInstruction: Content? = null,
-  private val controller: APIController
+  private val controller: APIController,
 ) {
-
-  @JvmOverloads
   internal constructor(
     modelName: String,
     apiKey: String,
@@ -84,42 +78,8 @@ internal constructor(
       modelName,
       requestOptions,
       "gl-kotlin/${KotlinVersion.CURRENT} fire/${BuildConfig.VERSION_NAME}",
-      object : HeaderProvider {
-        override val timeout: Duration
-          get() = 10.seconds
-
-        override suspend fun generateHeaders(): Map<String, String> {
-          val headers = mutableMapOf<String, String>()
-          if (appCheckTokenProvider == null) {
-            Log.w(TAG, "AppCheck not registered, skipping")
-          } else {
-            val token = appCheckTokenProvider.getToken(false).await()
-
-            if (token.error != null) {
-              Log.w(TAG, "Error obtaining AppCheck token", token.error)
-            }
-            // The Firebase App Check backend can differentiate between apps without App Check, and
-            // wrongly configured apps by verifying the value of the token, so it always needs to be
-            // included.
-            headers["X-Firebase-AppCheck"] = token.token
-          }
-
-          if (internalAuthProvider == null) {
-            Log.w(TAG, "Auth not registered, skipping")
-          } else {
-            try {
-              val token = internalAuthProvider.getAccessToken(false).await()
-
-              headers["Authorization"] = "Firebase ${token.token!!}"
-            } catch (e: Exception) {
-              Log.w(TAG, "Error getting Auth token ", e)
-            }
-          }
-
-          return headers
-        }
-      }
-    )
+      AppCheckHeaderProvider(TAG, appCheckTokenProvider, internalAuthProvider),
+    ),
   )
 
   /**
@@ -247,7 +207,7 @@ internal constructor(
       generationConfig?.toInternal(),
       tools?.map { it.toInternal() },
       toolConfig?.toInternal(),
-      systemInstruction?.copy(role = "system")?.toInternal()
+      systemInstruction?.copy(role = "system")?.toInternal(),
     )
 
   private fun constructCountTokensRequest(vararg prompt: Content) =

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/ImagenModel.kt

@@ -0,0 +1,137 @@
+/*
+ * Copyright 2025 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.vertexai
+
+import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
+import com.google.firebase.auth.internal.InternalAuthProvider
+import com.google.firebase.vertexai.common.APIController
+import com.google.firebase.vertexai.common.AppCheckHeaderProvider
+import com.google.firebase.vertexai.common.ContentBlockedException
+import com.google.firebase.vertexai.common.GenerateImageRequest
+import com.google.firebase.vertexai.type.FirebaseVertexAIException
+import com.google.firebase.vertexai.type.ImagenGCSImage
+import com.google.firebase.vertexai.type.ImagenGenerationConfig
+import com.google.firebase.vertexai.type.ImagenGenerationResponse
+import com.google.firebase.vertexai.type.ImagenInlineImage
+import com.google.firebase.vertexai.type.ImagenSafetySettings
+import com.google.firebase.vertexai.type.RequestOptions
+
+/**
+ * Represents a generative model (like Imagen), capable of generating images based on various input
+ * types.
+ */
+public class ImagenModel
+internal constructor(
+  private val modelName: String,
+  private val generationConfig: ImagenGenerationConfig? = null,
+  private val safetySettings: ImagenSafetySettings? = null,
+  private val controller: APIController,
+) {
+  @JvmOverloads
+  internal constructor(
+    modelName: String,
+    apiKey: String,
+    generationConfig: ImagenGenerationConfig? = null,
+    safetySettings: ImagenSafetySettings? = null,
+    requestOptions: RequestOptions = RequestOptions(),
+    appCheckTokenProvider: InteropAppCheckTokenProvider? = null,
+    internalAuthProvider: InternalAuthProvider? = null,
+  ) : this(
+    modelName,
+    generationConfig,
+    safetySettings,
+    APIController(
+      apiKey,
+      modelName,
+      requestOptions,
+      "gl-kotlin/${KotlinVersion.CURRENT} fire/${BuildConfig.VERSION_NAME}",
+      AppCheckHeaderProvider(TAG, appCheckTokenProvider, internalAuthProvider),
+    ),
+  )
+
+  /**
+   * Generates an image, returning the result directly to the caller.
+   *
+   * @param prompt The input(s) given to the model as a prompt.
+   */
+  public suspend fun generateImages(prompt: String): ImagenGenerationResponse<ImagenInlineImage> =
+    try {
+      controller
+        .generateImage(constructRequest(prompt, null, generationConfig))
+        .validate()
+        .toPublicInline()
+    } catch (e: Throwable) {
+      throw FirebaseVertexAIException.from(e)
+    }
+
+  /**
+   * Generates an image, storing the result in Google Cloud Storage and returning a URL
+   *
+   * @param prompt The input(s) given to the model as a prompt.
+   * @param gcsUri Specifies where in Google Cloud Storage to store the image (for example, a
+   * specific bucket or folder).
+   */
+  public suspend fun generateImages(
+    prompt: String,
+    gcsUri: String,
+  ): ImagenGenerationResponse<ImagenGCSImage> =
+    try {
+      controller
+        .generateImage(constructRequest(prompt, gcsUri, generationConfig))
+        .validate()
+        .toPublicGCS()
+    } catch (e: Throwable) {
+      throw FirebaseVertexAIException.from(e)
+    }
+
+  private fun constructRequest(
+    prompt: String,
+    gcsUri: String?,
+    config: ImagenGenerationConfig?,
+  ): GenerateImageRequest {
+    return GenerateImageRequest(
+      listOf(GenerateImageRequest.ImagenPrompt(prompt)),
+      GenerateImageRequest.ImagenParameters(
+        sampleCount = config?.numberOfImages ?: 1,
+        includeRaiReason = true,
+        addWatermark = generationConfig?.addWatermark,
+        personGeneration = safetySettings?.personFilterLevel?.internalVal,
+        negativePrompt = config?.negativePrompt,
+        safetySetting = safetySettings?.safetyFilterLevel?.internalVal,
+        storageUri = gcsUri,
+        aspectRatio = config?.aspectRatio?.internalVal,
+        imageOutputOptions = generationConfig?.imageFormat?.toInternal(),
+      ),
+    )
+  }
+
+  internal companion object {
+    private val TAG = ImagenModel::class.java.simpleName
+    internal const val DEFAULT_FILTERED_ERROR =
+      "Unable to show generated images. All images were filtered out because they violated Vertex AI's usage guidelines. You will not be charged for blocked images. Try rephrasing the prompt. If you think this was an error, send feedback."
+  }
+}
+
+private fun ImagenGenerationResponse.Internal.validate(): ImagenGenerationResponse.Internal {
+  if (predictions.none { it.mimeType != null }) {
+    throw ContentBlockedException(
+      message = predictions.first { it.raiFilteredReason != null }.raiFilteredReason
+          ?: ImagenModel.DEFAULT_FILTERED_ERROR
+    )
+  }
+  return this
+}

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/APIController.kt

@@ -25,6 +25,7 @@ import com.google.firebase.vertexai.type.CountTokensResponse
 import com.google.firebase.vertexai.type.FinishReason
 import com.google.firebase.vertexai.type.GRpcErrorResponse
 import com.google.firebase.vertexai.type.GenerateContentResponse
+import com.google.firebase.vertexai.type.ImagenGenerationResponse
 import com.google.firebase.vertexai.type.RequestOptions
 import com.google.firebase.vertexai.type.Response
 import io.ktor.client.HttpClient
@@ -58,12 +59,15 @@ import kotlinx.coroutines.flow.channelFlow
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withTimeout
+import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.json.Json
 
+@OptIn(ExperimentalSerializationApi::class)
 internal val JSON = Json {
   ignoreUnknownKeys = true
   prettyPrint = false
   isLenient = true
+  explicitNulls = false
 }
 
 /**
@@ -122,6 +126,19 @@ internal constructor(
       throw FirebaseCommonAIException.from(e)
     }
 
+  suspend fun generateImage(request: GenerateImageRequest): ImagenGenerationResponse.Internal =
+    try {
+      client
+        .post("${requestOptions.endpoint}/${requestOptions.apiVersion}/$model:predict") {
+          applyCommonConfiguration(request)
+          applyHeaderProvider()
+        }
+        .also { validateResponse(it) }
+        .body<ImagenGenerationResponse.Internal>()
+    } catch (e: Throwable) {
+      throw FirebaseCommonAIException.from(e)
+    }
+
   fun generateContentStream(
     request: GenerateContentRequest
   ): Flow<GenerateContentResponse.Internal> =
@@ -151,6 +168,7 @@ internal constructor(
     when (request) {
       is GenerateContentRequest -> setBody<GenerateContentRequest>(request)
       is CountTokensRequest -> setBody<CountTokensRequest>(request)
+      is GenerateImageRequest -> setBody<GenerateImageRequest>(request)
     }
     contentType(ContentType.Application.Json)
     header("x-goog-api-key", key)
@@ -258,6 +276,9 @@ private suspend fun validateResponse(response: HttpResponse) {
   if (message.contains("quota")) {
     throw QuotaExceededException(message)
   }
+  if (message.contains("The prompt could not be submitted")) {
+    throw PromptBlockedException(message)
+  }
   getServiceDisabledErrorDetailsOrNull(error)?.let {
     val errorMessage =
       if (it.metadata?.get("service") == "firebasevertexai.googleapis.com") {