add tests and minor adjustments to support them

David Motsonashvili · David Motsonashvili · commit a9ea9e6b6aeb · 2025-01-22T14:22:24.000-08:00
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/FirebaseVertexAI.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/FirebaseVertexAI.kt
@@ -26,7 +26,6 @@ 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.ImagenModelConfig
 import com.google.firebase.vertexai.type.InvalidLocationException
 import com.google.firebase.vertexai.type.RequestOptions
 import com.google.firebase.vertexai.type.SafetySetting
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/ImageModel.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/ImageModel.kt
@@ -5,24 +5,24 @@ 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.HeaderProvider
+import com.google.firebase.vertexai.common.PromptBlockedException
 import com.google.firebase.vertexai.internal.GenerateImageRequest
+import com.google.firebase.vertexai.internal.GenerateImageResponse
 import com.google.firebase.vertexai.internal.ImagenParameters
 import com.google.firebase.vertexai.internal.ImagenPromptInstance
 import com.google.firebase.vertexai.internal.util.toInternal
 import com.google.firebase.vertexai.internal.util.toPublicGCS
 import com.google.firebase.vertexai.internal.util.toPublicInline
 import com.google.firebase.vertexai.type.FirebaseVertexAIException
-import com.google.firebase.vertexai.type.ImagenSafetySettings
 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.ImagenModelConfig
-import com.google.firebase.vertexai.type.PromptBlockedException
+import com.google.firebase.vertexai.type.ImagenSafetySettings
 import com.google.firebase.vertexai.type.RequestOptions
-import kotlinx.coroutines.tasks.await
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
+import kotlinx.coroutines.tasks.await
 
 public class ImageModel
 internal constructor(
@@ -87,11 +87,12 @@ internal constructor(
     ),
   )
 
-  public suspend fun generateImage(
-    prompt: String,
-  ): ImagenGenerationResponse<ImagenInlineImage> =
+  public suspend fun generateImage(prompt: String): ImagenGenerationResponse<ImagenInlineImage> =
     try {
-      controller.generateImage(constructRequest(prompt, null, generationConfig)).toPublicInline().validate()
+      controller
+        .generateImage(constructRequest(prompt, null, generationConfig))
+        .validate()
+        .toPublicInline()
     } catch (e: Throwable) {
       throw FirebaseVertexAIException.from(e)
     }
@@ -101,7 +102,10 @@ internal constructor(
     gcsUri: String,
   ): ImagenGenerationResponse<ImagenGCSImage> =
     try {
-      controller.generateImage(constructRequest(prompt, gcsUri, generationConfig)).toPublicGCS().validate()
+      controller
+        .generateImage(constructRequest(prompt, gcsUri, generationConfig))
+        .validate()
+        .toPublicGCS()
     } catch (e: Throwable) {
       throw FirebaseVertexAIException.from(e)
     }
@@ -129,14 +133,17 @@ internal constructor(
 
   internal companion object {
     private val TAG = ImageModel::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."
+    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 <T> ImagenGenerationResponse<T>.validate():
-  ImagenGenerationResponse<T> {
-  if (images.isEmpty()) {
-    throw PromptBlockedException(message = filteredReason ?: ImageModel.DEFAULT_FILTERED_ERROR)
+private fun GenerateImageResponse.validate(): GenerateImageResponse {
+  if (predictions.none { it.mimeType != null }) {
+    throw PromptBlockedException(
+      message = predictions.first { it.raiFilteredReason != null }.raiFilteredReason
+          ?: ImageModel.DEFAULT_FILTERED_ERROR
+    )
   }
   return this
 }
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/APIController.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/APIController.kt
@@ -271,6 +271,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") {
@@ -309,6 +312,4 @@ private fun GenerateContentResponse.validate() = apply {
     ?.let { throw ResponseStoppedException(this) }
 }
 
-private fun GenerateImageResponse.validate() = apply {
-
-}
+private fun GenerateImageResponse.validate() = apply {}
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/Exceptions.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/Exceptions.kt
@@ -65,14 +65,18 @@ internal class InvalidAPIKeyException(message: String, cause: Throwable? = null)
  *
  * @property response the full server response for the request.
  */
-internal class PromptBlockedException(
-  val response: GenerateContentResponse,
-  cause: Throwable? = null
+internal class PromptBlockedException
+internal constructor(
+  val response: GenerateContentResponse?,
+  cause: Throwable? = null,
+  message: String? = null,
 ) :
   FirebaseCommonAIException(
-    "Prompt was blocked: ${response.promptFeedback?.blockReason?.name}",
+    "Prompt was blocked: ${response?.promptFeedback?.blockReason?.name?: message}",
     cause,
-  )
+  ) {
+  internal constructor(message: String, cause: Throwable? = null) : this(null, cause, message)
+}
 
 /**
  * The user's location (region) is not supported by the API.
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/GenerateImageResponse.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/GenerateImageResponse.kt
@@ -9,5 +9,6 @@ internal data class GenerateImageResponse(val predictions: List<ImagenImageRespo
 internal data class ImagenImageResponse(
   val bytesBase64Encoded: String? = null,
   val gcsUri: String? = null,
-  val mimeType: String,
+  val mimeType: String? = null,
+  val raiFilteredReason: String? = null,
 )
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt
@@ -368,16 +368,22 @@ internal fun com.google.firebase.vertexai.common.CountTokensResponse.toPublic()
   CountTokensResponse(totalTokens, totalBillableCharacters ?: 0)
 
 internal fun com.google.firebase.vertexai.internal.GenerateImageResponse.toPublicInline() =
-  ImagenGenerationResponse(images = predictions.map { it.toPublicInline() }, null)
+  ImagenGenerationResponse(
+    images = predictions.filter { it.mimeType != null }.map { it.toPublicInline() },
+    null,
+  )
 
 internal fun com.google.firebase.vertexai.internal.ImagenImageResponse.toPublicInline() =
-  ImagenInlineImage(bytesBase64Encoded!!.toByteArray(), mimeType)
+  ImagenInlineImage(bytesBase64Encoded!!.toByteArray(), mimeType!!)
 
 internal fun com.google.firebase.vertexai.internal.GenerateImageResponse.toPublicGCS() =
-  ImagenGenerationResponse(images = predictions.map { it.toPublicGCS() }, null)
+  ImagenGenerationResponse(
+    images = predictions.filter { it.mimeType != null }.map { it.toPublicGCS() },
+    null,
+  )
 
 internal fun com.google.firebase.vertexai.internal.ImagenImageResponse.toPublicGCS() =
-  ImagenGCSImage(gcsUri!!, mimeType)
+  ImagenGCSImage(gcsUri!!, mimeType!!)
 
 internal fun JsonObject.toPublic() = JSONObject(toString())
 
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Exceptions.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Exceptions.kt
@@ -45,7 +45,7 @@ internal constructor(message: String, cause: Throwable? = null) : RuntimeExcepti
             is com.google.firebase.vertexai.common.InvalidAPIKeyException ->
               InvalidAPIKeyException(cause.message ?: "")
             is com.google.firebase.vertexai.common.PromptBlockedException ->
-              PromptBlockedException(cause.response.toPublic(), cause.cause)
+              PromptBlockedException(cause.response?.toPublic(), cause.cause)
             is com.google.firebase.vertexai.common.UnsupportedUserLocationException ->
               UnsupportedUserLocationException(cause.cause)
             is com.google.firebase.vertexai.common.InvalidStateException ->
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenAspectRatio.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenAspectRatio.kt
@@ -2,9 +2,9 @@ package com.google.firebase.vertexai.type
 
 @Suppress("EnumEntryName")
 public enum class ImagenAspectRatio(internal val internalVal: String) {
-    SQUARE_1x1("1:1"),
-    PORTRAIT_3x4("3:4"),
-    LANDSCAPE_4x3("4:3"),
-    PORTRAIT_9x16("9:16"),
-    LANDSCAPE_16x9("16:9")
+  SQUARE_1x1("1:1"),
+  PORTRAIT_3x4("3:4"),
+  LANDSCAPE_4x3("4:3"),
+  PORTRAIT_9x16("9:16"),
+  LANDSCAPE_16x9("16:9")
 }
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenGenerationConfig.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenGenerationConfig.kt
@@ -1,9 +1,9 @@
 package com.google.firebase.vertexai.type
 
 public class ImagenGenerationConfig(
-    public val negativePrompt: String? = null,
-    public val numberOfImages: Int? = 1,
-    public val aspectRatio: ImagenAspectRatio? = null,
-    public val imageFormat: ImagenImageFormat? = null,
-    public val addWatermark: Boolean? = null
+  public val negativePrompt: String? = null,
+  public val numberOfImages: Int? = 1,
+  public val aspectRatio: ImagenAspectRatio? = null,
+  public val imageFormat: ImagenImageFormat? = null,
+  public val addWatermark: Boolean? = null
 ) {}
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenModelConfig.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenModelConfig.kt
@@ -1,7 +1,3 @@
 package com.google.firebase.vertexai.type
 
-public class ImagenModelConfig private constructor(
-
-) {
-
-}
+public class ImagenModelConfig private constructor() {}
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenPersonFilter.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenPersonFilter.kt
@@ -1,7 +1,7 @@
 package com.google.firebase.vertexai.type
 
 public enum class ImagenPersonFilter(internal val internalVal: String) {
-    ALLOW_ALL("allow_all"),
-    ALLOW_ADULT("allow_adult"),
-    BLOCK_ALL("dont_allow"),
+  ALLOW_ALL("allow_all"),
+  ALLOW_ADULT("allow_adult"),
+  BLOCK_ALL("dont_allow"),
 }
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenSafetyFilter.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/ImagenSafetyFilter.kt
@@ -1,8 +1,8 @@
-package com.google.firebase.vertexai.type;
+package com.google.firebase.vertexai.type
 
 public enum class ImagenSafetyFilter(internal val internalVal: String) {
-    BLOCK_LOW_AND_ABOVE("block_low_and_above"),
-    BLOCK_MEDIUM_AND_ABOVE("block_medium_and_above"),
-    BLOCK_ONLY_HIGH("block_only_high"),
-    BLOCK_NONE("block_none")
+  BLOCK_LOW_AND_ABOVE("block_low_and_above"),
+  BLOCK_MEDIUM_AND_ABOVE("block_medium_and_above"),
+  BLOCK_ONLY_HIGH("block_only_high"),
+  BLOCK_NONE("block_none")
 }
diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt
@@ -119,7 +119,7 @@ internal class StreamingSnapshotTests {
 
       withTimeout(testTimeout) {
         val exception = shouldThrow<PromptBlockedException> { responses.collect() }
-        exception.response.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
+        exception.response?.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
       }
     }
 
@@ -130,8 +130,8 @@ internal class StreamingSnapshotTests {
 
       withTimeout(testTimeout) {
         val exception = shouldThrow<PromptBlockedException> { responses.collect() }
-        exception.response.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
-        exception.response.promptFeedback?.blockReasonMessage shouldBe "Reasons"
+        exception.response?.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
+        exception.response?.promptFeedback?.blockReasonMessage shouldBe "Reasons"
       }
     }
 
diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt
@@ -110,7 +110,7 @@ internal class UnarySnapshotTests {
       withTimeout(testTimeout) {
         shouldThrow<PromptBlockedException> { model.generateContent("prompt") } should
           {
-            it.response.promptFeedback?.blockReason shouldBe BlockReason.UNKNOWN
+            it.response?.promptFeedback?.blockReason shouldBe BlockReason.UNKNOWN
           }
       }
     }
@@ -165,7 +165,7 @@ internal class UnarySnapshotTests {
       withTimeout(testTimeout) {
         shouldThrow<PromptBlockedException> { model.generateContent("prompt") } should
           {
-            it.response.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
+            it.response?.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
           }
       }
     }
@@ -176,8 +176,8 @@ internal class UnarySnapshotTests {
       withTimeout(testTimeout) {
         shouldThrow<PromptBlockedException> { model.generateContent("prompt") } should
           {
-            it.response.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
-            it.response.promptFeedback?.blockReasonMessage shouldContain "Reasons"
+            it.response?.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
+            it.response?.promptFeedback?.blockReasonMessage shouldContain "Reasons"
           }
       }
     }
@@ -200,7 +200,7 @@ internal class UnarySnapshotTests {
   fun `user location error`() =
     goldenUnaryFile(
       "unary-failure-unsupported-user-location.json",
-      HttpStatusCode.PreconditionFailed
+      HttpStatusCode.PreconditionFailed,
     ) {
       withTimeout(testTimeout) {
         shouldThrow<UnsupportedUserLocationException> { model.generateContent("prompt") }
@@ -473,4 +473,39 @@ internal class UnarySnapshotTests {
     goldenUnaryFile("unary-failure-model-not-found.json", HttpStatusCode.NotFound) {
       withTimeout(testTimeout) { shouldThrow<ServerException> { model.countTokens("prompt") } }
     }
+
+  @Test
+  fun `generateImages should throw when all images filtered`() =
+    goldenUnaryFile("unary-failure-generate-images-all-filtered.json") {
+      withTimeout(testTimeout) {
+        shouldThrow<PromptBlockedException> { imageModel.generateImage("prompt") }
+      }
+    }
+
+  @Test
+  fun `generateImages should return when some images are filtered -- gcs`() =
+    goldenUnaryFile("unary-failure-generate-images-gcs-some-filtered.json") {
+      withTimeout(testTimeout) {
+        imageModel.generateImage("prompt", "gcsBucket").images.isEmpty() shouldBe false
+      }
+    }
+
+  @Test
+  fun `generateImages should throw when prompt blocked`() =
+    goldenUnaryFile(
+      "unary-failure-generate-images-prompt-blocked.json",
+      HttpStatusCode.BadRequest,
+    ) {
+      withTimeout(testTimeout) {
+        shouldThrow<PromptBlockedException> { imageModel.generateImage("prompt") }
+      }
+    }
+
+  @Test
+  fun `generateImages gcs should succeed`() =
+    goldenUnaryFile("unary-success-generate-images-gcs.json") {
+      withTimeout(testTimeout) {
+        imageModel.generateImage("prompt", "gcsBucket").images.isEmpty() shouldBe false
+      }
+    }
 }
diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/common/StreamingSnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/common/StreamingSnapshotTests.kt
@@ -104,7 +104,7 @@ internal class StreamingSnapshotTests {
 
       withTimeout(testTimeout) {
         val exception = shouldThrow<PromptBlockedException> { responses.collect() }
-        exception.response.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
+        exception.response?.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
       }
     }
 
diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/common/UnarySnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/common/UnarySnapshotTests.kt
@@ -111,7 +111,7 @@ internal class UnarySnapshotTests {
       withTimeout(testTimeout) {
         shouldThrow<PromptBlockedException> {
           apiController.generateContent(textGenerateContentRequest("prompt"))
-        } should { it.response.promptFeedback?.blockReason shouldBe BlockReason.SAFETY }
+        } should { it.response?.promptFeedback?.blockReason shouldBe BlockReason.SAFETY }
       }
     }
 
diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/util/tests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/util/tests.kt
@@ -17,6 +17,7 @@
 package com.google.firebase.vertexai.util
 
 import com.google.firebase.vertexai.GenerativeModel
+import com.google.firebase.vertexai.ImageModel
 import com.google.firebase.vertexai.common.APIController
 import com.google.firebase.vertexai.type.RequestOptions
 import io.kotest.matchers.collections.shouldNotBeEmpty
@@ -57,7 +58,11 @@ internal suspend fun ByteChannel.send(bytes: ByteArray) {
  * @see commonTest
  * @see send
  */
-internal data class CommonTestScope(val channel: ByteChannel, val model: GenerativeModel)
+internal data class CommonTestScope(
+  val channel: ByteChannel,
+  val model: GenerativeModel,
+  val imageModel: ImageModel,
+)
 
 /** A test that runs under a [CommonTestScope]. */
 internal typealias CommonTest = suspend CommonTestScope.() -> Unit
@@ -104,7 +109,8 @@ internal fun commonTest(
       null,
     )
   val model = GenerativeModel("cool-model-name", controller = apiController)
-  CommonTestScope(channel, model).block()
+  val imageModel = ImageModel("cooler-model-name", controller = apiController)
+  CommonTestScope(channel, model, imageModel).block()
 }
 
 /**
diff --git a/firebase-vertexai/update_responses.sh b/firebase-vertexai/update_responses.sh
@@ -17,7 +17,7 @@
 # This script replaces mock response files for Vertex AI unit tests with a fresh
 # clone of the shared repository of Vertex AI test data.
 
-RESPONSES_VERSION='v3.*' # The major version of mock responses to use
+RESPONSES_VERSION='v5.3' # The major version of mock responses to use
 REPO_NAME="vertexai-sdk-test-data"
 REPO_LINK="https://github.com/FirebaseExtended/$REPO_NAME.git"
 

Original file line number	Diff line number	Diff line change
`@@ -9,5 +9,6 @@ internal data class GenerateImageResponse(val predictions: List<ImagenImageRespo`
`9`	`9`	`internal data class ImagenImageResponse(`
`10`	`10`	`val bytesBase64Encoded: String? = null,`
`11`	`11`	`val gcsUri: String? = null,`
`12`		`- val mimeType: String,`
	`12`	`+ val mimeType: String? = null,`
	`13`	`+ val raiFilteredReason: String? = null,`
`13`	`14`	`)`