From ef3e9afa28272801b88133e677277242f97c6a4a Mon Sep 17 00:00:00 2001 From: Rodrigo Lazo Paz Date: Thu, 13 Nov 2025 12:15:54 -0500 Subject: [PATCH 1/3] [AI] Chat history must store all the parts (streaming) The chat history in streaming mode reconstructs the parts from their contents, rather than storing the parts themselves. This causes non-visible elements, like `thoughtSignature` to get lost. --- .../kotlin/com/google/firebase/ai/Chat.kt | 26 ++++++++++++------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt b/firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt index 73d304d3885..b491d5ec535 100644 --- a/firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt +++ b/firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt @@ -22,6 +22,7 @@ import com.google.firebase.ai.type.GenerateContentResponse import com.google.firebase.ai.type.ImagePart import com.google.firebase.ai.type.InlineDataPart import com.google.firebase.ai.type.InvalidStateException +import com.google.firebase.ai.type.Part import com.google.firebase.ai.type.TextPart import com.google.firebase.ai.type.content import java.util.LinkedList @@ -133,6 +134,7 @@ public class Chat( val bitmaps = LinkedList() val inlineDataParts = LinkedList() val text = StringBuilder() + val parts = mutableListOf() /** * TODO: revisit when images and inline data are returned. This will cause issues with how @@ -147,6 +149,7 @@ public class Chat( is ImagePart -> bitmaps.add(part.image) is InlineDataPart -> inlineDataParts.add(part) } + parts.add(part) } } .onCompletion { @@ -154,15 +157,20 @@ public class Chat( if (it == null) { val content = content("model") { - for (bitmap in bitmaps) { - image(bitmap) - } - for (inlineDataPart in inlineDataParts) { - inlineData(inlineDataPart.inlineData, inlineDataPart.mimeType) - } - if (text.isNotBlank()) { - text(text.toString()) - } + setParts( + parts + .filterNot { part -> part is TextPart && part.text.isNotEmpty() } + .toMutableList() + ) + // for (bitmap in bitmaps) { + // image(bitmap) + // } + // for (inlineDataPart in inlineDataParts) { + // inlineData(inlineDataPart.inlineData, inlineDataPart.mimeType) + // } + // if (text.isNotBlank()) { + // text(text.toString()) + // } } history.add(prompt) From 2bf06b1ca3c42384c61884acaf2dcc3a0213d75f Mon Sep 17 00:00:00 2001 From: Rodrigo Lazo Paz Date: Fri, 14 Nov 2025 11:06:08 -0500 Subject: [PATCH 2/3] Add testing --- .../kotlin/com/google/firebase/ai/Chat.kt | 22 ++--- .../ai/DevAPIStreamingSnapshotTests.kt | 91 +++++++++++++++++++ .../java/com/google/firebase/ai/util/tests.kt | 26 +++++- 3 files changed, 123 insertions(+), 16 deletions(-) diff --git a/firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt b/firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt index b491d5ec535..4ff2c3c8e82 100644 --- a/firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt +++ b/firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt @@ -158,19 +158,8 @@ public class Chat( val content = content("model") { setParts( - parts - .filterNot { part -> part is TextPart && part.text.isNotEmpty() } - .toMutableList() + parts.filterNot { part -> part is TextPart && !part.hasContent() }.toMutableList() ) - // for (bitmap in bitmaps) { - // image(bitmap) - // } - // for (inlineDataPart in inlineDataParts) { - // inlineData(inlineDataPart.inlineData, inlineDataPart.mimeType) - // } - // if (text.isNotBlank()) { - // text(text.toString()) - // } } history.add(prompt) @@ -232,3 +221,12 @@ public class Chat( } } } + +/** + * Returns true if the [TextPart] contains any content, either in its [TextPart.text] property or + * its [TextPart.thoughtSignature] property. + */ +private fun TextPart.hasContent(): Boolean { + if (text.isNotEmpty()) return true + return !thoughtSignature.isNullOrBlank() +} diff --git a/firebase-ai/src/test/java/com/google/firebase/ai/DevAPIStreamingSnapshotTests.kt b/firebase-ai/src/test/java/com/google/firebase/ai/DevAPIStreamingSnapshotTests.kt index 2aac8f7a0d2..37cab5615e5 100644 --- a/firebase-ai/src/test/java/com/google/firebase/ai/DevAPIStreamingSnapshotTests.kt +++ b/firebase-ai/src/test/java/com/google/firebase/ai/DevAPIStreamingSnapshotTests.kt @@ -18,18 +18,29 @@ package com.google.firebase.ai import com.google.firebase.ai.type.BlockReason import com.google.firebase.ai.type.FinishReason +import com.google.firebase.ai.type.FunctionCallPart import com.google.firebase.ai.type.PromptBlockedException import com.google.firebase.ai.type.ResponseStoppedException import com.google.firebase.ai.type.ServerException +import com.google.firebase.ai.type.content import com.google.firebase.ai.util.goldenDevAPIStreamingFile import io.kotest.assertions.throwables.shouldThrow import io.kotest.matchers.collections.shouldBeEmpty +import io.kotest.matchers.collections.shouldNotBeEmpty +import io.kotest.matchers.nulls.shouldNotBeNull import io.kotest.matchers.shouldBe +import io.kotest.matchers.string.shouldStartWith +import io.ktor.client.engine.mock.toByteArray +import io.ktor.client.request.HttpRequestData 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 kotlinx.serialization.json.Json +import kotlinx.serialization.json.jsonArray +import kotlinx.serialization.json.jsonObject +import kotlinx.serialization.json.jsonPrimitive import org.junit.Test import org.junit.runner.RunWith import org.robolectric.RobolectricTestRunner @@ -85,6 +96,86 @@ internal class DevAPIStreamingSnapshotTests { } } + @Test + fun `success call with thought summary and signature`() = + goldenDevAPIStreamingFile( + "streaming-success-thinking-function-call-thought-summary-signature.txt" + ) { + val responses = model.generateContentStream("prompt") + + withTimeout(testTimeout) { + val responseList = responses.toList() + responseList.isEmpty() shouldBe false + val functionCallResponse = responseList.find { it.functionCalls.isNotEmpty() } + functionCallResponse.shouldNotBeNull() + functionCallResponse.functionCalls.first().let { + it.thoughtSignature.shouldNotBeNull() + it.thoughtSignature.shouldStartWith("CiIBVKhc7vB") + } + } + } + + @Test + fun `chat call with history including thought summary and signature`() { + var capturedRequest: HttpRequestData? = null + goldenDevAPIStreamingFile( + "streaming-success-thinking-function-call-thought-summary-signature.txt", + requestHandler = { capturedRequest = it } + ) { + val chat = model.startChat() + val firstPrompt = content { text("first prompt") } + val secondPrompt = content { text("second prompt") } + val responses = chat.sendMessageStream(firstPrompt) + + withTimeout(testTimeout) { + val responseList = responses.toList() + responseList.shouldNotBeEmpty() + + chat.history.let { history -> + history.contains(firstPrompt) + val functionCallPart = + history.flatMap { it.parts }.first { it is FunctionCallPart } as FunctionCallPart + functionCallPart.let { + it.thoughtSignature.shouldNotBeNull() + it.thoughtSignature.shouldStartWith("CiIBVKhc7vB") + } + } + + // Reset the request so we can be sure we capture the latest version + capturedRequest = null + + // We don't care about the response, only the request + val unused = chat.sendMessageStream(secondPrompt).toList() + + // Make sure the history contains all prompts seen so far + chat.history.contains(firstPrompt) + chat.history.contains(secondPrompt) + + // Put the captured request into a `val` to enable smart casting + val request = capturedRequest + request.shouldNotBeNull() + val bodyAsString = request.body.toByteArray().decodeToString() + bodyAsString.shouldNotBeNull() + + val rootElement = Json.parseToJsonElement(bodyAsString).jsonObject + + // Traverse the tree: contents -> parts -> thoughtSignature + val contents = rootElement["contents"]?.jsonArray + + val signature = + contents?.firstNotNullOfOrNull { content -> + content.jsonObject["parts"]?.jsonArray?.firstNotNullOfOrNull { part -> + // resulting value is a JsonPrimitive, so we use .content to get the string + part.jsonObject["thoughtSignature"]?.jsonPrimitive?.content + } + } + + signature.shouldNotBeNull() + signature.shouldStartWith("CiIBVKhc7vB") + } + } + } + @Test fun `prompt blocked for safety`() = goldenDevAPIStreamingFile("streaming-failure-prompt-blocked-safety.txt") { diff --git a/firebase-ai/src/test/java/com/google/firebase/ai/util/tests.kt b/firebase-ai/src/test/java/com/google/firebase/ai/util/tests.kt index 393a2a16adc..1394bb07e86 100644 --- a/firebase-ai/src/test/java/com/google/firebase/ai/util/tests.kt +++ b/firebase-ai/src/test/java/com/google/firebase/ai/util/tests.kt @@ -29,11 +29,11 @@ import io.kotest.matchers.collections.shouldNotBeEmpty import io.kotest.matchers.nulls.shouldNotBeNull import io.ktor.client.engine.mock.MockEngine import io.ktor.client.engine.mock.respond +import io.ktor.client.request.HttpRequestData import io.ktor.http.HttpHeaders import io.ktor.http.HttpStatusCode import io.ktor.http.headersOf import io.ktor.utils.io.ByteChannel -import io.ktor.utils.io.close import io.ktor.utils.io.writeFully import java.io.File import kotlinx.coroutines.launch @@ -103,6 +103,7 @@ internal fun commonTest( status: HttpStatusCode = HttpStatusCode.OK, requestOptions: RequestOptions = RequestOptions(), backend: GenerativeBackend = GenerativeBackend.vertexAI(), + requestHandler: (HttpRequestData) -> Unit = {}, block: CommonTest, ) = doBlocking { val channel = ByteChannel(autoFlush = true) @@ -115,6 +116,7 @@ internal fun commonTest( "gemini-pro", requestOptions, MockEngine { + requestHandler(it) respond(channel, status, headersOf(HttpHeaders.ContentType, "application/json")) }, TEST_CLIENT_ID, @@ -144,12 +146,13 @@ internal fun goldenStreamingFile( name: String, httpStatusCode: HttpStatusCode = HttpStatusCode.OK, backend: GenerativeBackend = GenerativeBackend.vertexAI(), + requestHandler: (HttpRequestData) -> Unit, block: CommonTest, ) = doBlocking { val goldenFile = loadGoldenFile(name) val messages = goldenFile.readLines().filter { it.isNotBlank() } - commonTest(httpStatusCode, backend = backend) { + commonTest(httpStatusCode, backend = backend, requestHandler = requestHandler) { launch { for (message in messages) { channel.writeFully("$message$SSE_SEPARATOR".toByteArray()) @@ -175,8 +178,15 @@ internal fun goldenStreamingFile( internal fun goldenVertexStreamingFile( name: String, httpStatusCode: HttpStatusCode = HttpStatusCode.OK, + requestHandler: (HttpRequestData) -> Unit = {}, block: CommonTest, -) = goldenStreamingFile("vertexai/$name", httpStatusCode, block = block) +) = + goldenStreamingFile( + "vertexai/$name", + httpStatusCode, + requestHandler = requestHandler, + block = block + ) /** * A variant of [goldenStreamingFile] for testing the developer api @@ -192,8 +202,16 @@ internal fun goldenVertexStreamingFile( internal fun goldenDevAPIStreamingFile( name: String, httpStatusCode: HttpStatusCode = HttpStatusCode.OK, + requestHandler: (HttpRequestData) -> Unit = {}, block: CommonTest, -) = goldenStreamingFile("googleai/$name", httpStatusCode, GenerativeBackend.googleAI(), block) +) = + goldenStreamingFile( + "googleai/$name", + httpStatusCode, + GenerativeBackend.googleAI(), + requestHandler, + block + ) /** * A variant of [commonTest] for performing snapshot tests. From 032cc0c05f3741e21bca7b86b4bec75174b921bd Mon Sep 17 00:00:00 2001 From: Rodrigo Lazo Paz Date: Fri, 14 Nov 2025 17:50:28 -0500 Subject: [PATCH 3/3] Add changelog entry --- firebase-ai/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/firebase-ai/CHANGELOG.md b/firebase-ai/CHANGELOG.md index fa57a850988..725e0b52ea4 100644 --- a/firebase-ai/CHANGELOG.md +++ b/firebase-ai/CHANGELOG.md @@ -1,5 +1,6 @@ # Unreleased +- [fixed] Fixed an issue causing streaming chat interactions to drop thought signatures. (#7562) - [feature] Added support for server templates via `TemplateGenerativeModel` and `TemplateImagenModel`. (#7503)