[Ai] Add workaround for invalid SafetyRating from the backend. (#6925)

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.

Author: rlazo

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

@@ -51,7 +51,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()
 
@@ -120,23 +120,31 @@ 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-ai/src/main/kotlin/com/google/firebase/ai/type/PromptFeedback.kt

@@ -42,7 +42,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-ai/src/test/java/com/google/firebase/ai/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
 
@@ -85,6 +88,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 finish reason`() =
     goldenVertexStreamingFile("streaming-failure-unknown-finish-enum.txt") {

firebase-ai/src/test/java/com/google/firebase/ai/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") {