Add support for streaming tests

Author: rlazo

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

@@ -17,6 +17,7 @@
 package com.google.firebase.ai.type
 
 import android.graphics.Bitmap
+import kotlin.collections.filterNot
 import kotlinx.serialization.EncodeDefault
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Serializable
@@ -90,14 +91,21 @@ constructor(public val role: String? = "user", public val parts: List<Part>) {
   @Serializable
   internal data class Internal(
     @EncodeDefault val role: String? = "user",
-    val parts: List<InternalPart>
+    val parts: List<InternalPart>? = null
   ) {
     internal fun toPublic(): Content {
+      // Return empty if none of the parts is a known part
+      if (parts == null || parts.filterNot { it is UnknownPart.Internal }.isEmpty()) {
+        return Content(role, emptyList())
+      }
+      // From all the known parts, if they are all text and empty, we coalesce them into a single
+      // one-character string part so the backend doesn't fail if we send this back as part of a
+      // multi-turn interaction.
       val returnedParts =
-        parts.filterNot { it is UnknownPart.Internal }.map { it.toPublic() }.filterNot { it is TextPart && it.text.isEmpty() }
-      // If all returned parts were text and empty, we coalesce them into a single one-character
-      // string
-      // part so the backend doesn't fail if we send this back as part of a multi-turn interaction.
+        parts
+          .filterNot { it is UnknownPart.Internal }
+          .map { it.toPublic() }
+          .filterNot { it is TextPart && it.text.isEmpty() }
       return Content(role, returnedParts.ifEmpty { listOf(TextPart(" ")) })
     }
   }

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

@@ -271,10 +271,8 @@ internal constructor(
   }
 }
 
-
-internal data class UnknownPart(public override val isThought: Boolean = false): Part {
-  @Serializable
-  internal data class Internal(val thought: Boolean? = null): InternalPart
+internal data class UnknownPart(public override val isThought: Boolean = false) : Part {
+  @Serializable internal data class Internal(val thought: Boolean? = null) : InternalPart
 }
 
 /** Returns the part as a [String] if it represents text, and null otherwise */

firebase-ai/src/test/java/com/google/firebase/ai/DevAPIStreamingSnapshotTests.kt

@@ -23,14 +23,18 @@ import com.google.firebase.ai.type.ResponseStoppedException
 import com.google.firebase.ai.type.ServerException
 import com.google.firebase.ai.util.goldenDevAPIStreamingFile
 import io.kotest.assertions.throwables.shouldThrow
+import io.kotest.matchers.collections.shouldBeEmpty
 import io.kotest.matchers.shouldBe
 import io.ktor.http.HttpStatusCode
 import kotlin.time.Duration.Companion.seconds
 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 DevAPIStreamingSnapshotTests {
   private val testTimeout = 5.seconds
 
@@ -64,6 +68,23 @@ internal class DevAPIStreamingSnapshotTests {
       }
     }
 
+  @Test
+  fun `reply with a mostly empty part`() =
+    goldenDevAPIStreamingFile("streaming-success-empty-parts.txt") {
+      val responses = model.generateContentStream("prompt")
+
+      withTimeout(testTimeout) {
+        val responseList = responses.toList()
+        responseList.isEmpty() shouldBe false
+        // Second to last response has no parts
+        responseList[5].candidates.first().content.parts.shouldBeEmpty()
+        responseList.last().candidates.first().apply {
+          finishReason shouldBe FinishReason.STOP
+          content.parts.isEmpty() shouldBe false
+        }
+      }
+    }
+
   @Test
   fun `prompt blocked for safety`() =
     goldenDevAPIStreamingFile("streaming-failure-prompt-blocked-safety.txt") {