Add change from #6925

Author: daymxn

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt

@@ -21,6 +21,7 @@ package com.google.firebase.vertexai.type
 import com.google.firebase.vertexai.common.util.FirstOrdinalSerializer
 import java.util.Calendar
 import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.InternalSerializationApi
 import kotlinx.serialization.KSerializer
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
@@ -55,7 +56,7 @@ internal constructor(
     val groundingMetadata: GroundingMetadata? = null,
   ) {
     internal fun toPublic(): Candidate {
-      val safetyRatings = safetyRatings?.map { it.toPublic() }.orEmpty()
+      val safetyRatings = safetyRatings?.mapNotNull { it.toPublic() }.orEmpty()
       val citations = citationMetadata?.toPublic()
       val finishReason = finishReason?.toPublic()
 
@@ -128,23 +129,32 @@ internal constructor(
   internal data class Internal
   @JvmOverloads
   constructor(
-    val category: HarmCategory.Internal,
-    val probability: HarmProbability.Internal,
+    val category: HarmCategory.Internal? = null,
+    val probability: HarmProbability.Internal? = null,
     val blocked: Boolean? = null, // TODO(): any reason not to default to false?
     val probabilityScore: Float? = null,
     val severity: HarmSeverity.Internal? = null,
     val severityScore: Float? = null,
   ) {
 
     internal fun toPublic() =
-      SafetyRating(
-        category = category.toPublic(),
-        probability = probability.toPublic(),
-        probabilityScore = probabilityScore ?: 0f,
-        blocked = blocked,
-        severity = severity?.toPublic(),
-        severityScore = severityScore
-      )
+      /**
+       * Due to a bug in the backend, it's possible that we receive an invalid `SafetyRating` value,
+       * without either category or probability. We return null in those cases to enable filtering
+       * by the higher level types.
+       */
+      if (category == null || probability == null) {
+        null
+      } else {
+        SafetyRating(
+          category = category.toPublic(),
+          probability = probability.toPublic(),
+          probabilityScore = probabilityScore ?: 0f,
+          blocked = blocked,
+          severity = severity?.toPublic(),
+          severityScore = severityScore
+        )
+      }
   }
 }
 

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/PromptFeedback.kt

@@ -46,7 +46,7 @@ public class PromptFeedback(
   ) {
 
     internal fun toPublic(): PromptFeedback {
-      val safetyRatings = safetyRatings?.map { it.toPublic() }.orEmpty()
+      val safetyRatings = safetyRatings?.mapNotNull { it.toPublic() }.orEmpty()
       return PromptFeedback(blockReason?.toPublic(), safetyRatings, blockReasonMessage)
     }
   }

firebase-vertexai/src/test/java/com/google/firebase/vertexai/VertexAIStreamingSnapshotTests.kt

@@ -36,7 +36,10 @@ import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.withTimeout
 import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
 
+@RunWith(RobolectricTestRunner::class)
 internal class VertexAIStreamingSnapshotTests {
   private val testTimeout = 5.seconds
 
@@ -69,6 +72,18 @@ internal class VertexAIStreamingSnapshotTests {
       }
     }
 
+  @Test
+  fun `invalid safety ratings during image generation`() =
+    goldenVertexStreamingFile("streaming-success-image-invalid-safety-ratings.txt") {
+      val responses = model.generateContentStream("prompt")
+
+      withTimeout(testTimeout) {
+        val responseList = responses.toList()
+
+        responseList.isEmpty() shouldBe false
+      }
+    }
+
   @Test
   fun `unknown enum in safety ratings`() =
     goldenVertexStreamingFile("streaming-success-unknown-safety-enum.txt") {

firebase-vertexai/src/test/java/com/google/firebase/vertexai/VertexAIUnarySnapshotTests.kt

@@ -55,8 +55,11 @@ import kotlinx.serialization.json.jsonObject
 import kotlinx.serialization.json.jsonPrimitive
 import org.json.JSONArray
 import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(PublicPreviewAPI::class)
+@RunWith(RobolectricTestRunner::class)
 internal class VertexAIUnarySnapshotTests {
   private val testTimeout = 5.seconds
 
@@ -125,6 +128,16 @@ internal class VertexAIUnarySnapshotTests {
       }
     }
 
+  @Test
+  fun `invalid safety ratings during image generation`() =
+    goldenVertexUnaryFile("unary-success-image-invalid-safety-ratings.json") {
+      withTimeout(testTimeout) {
+        val response = model.generateContent("prompt")
+
+        response.candidates.isEmpty() shouldBe false
+      }
+    }
+
   @Test
   fun `unknown enum in finish reason`() =
     goldenVertexUnaryFile("unary-failure-unknown-enum-finish-reason.json") {