Merge branch 'dconeybe/dataconnect/MutableStateFlowUseUpdateInsteadOfCompareAndSet' into dconeybe/dataconnect/FirebaseDataConnectImplState

Author: dconeybe

firebase-crashlytics/src/main/java/com/google/firebase/crashlytics/internal/common/SessionReportingCoordinator.java

@@ -36,6 +36,7 @@
 import com.google.firebase.crashlytics.internal.send.DataTransportCrashlyticsReportSender;
 import com.google.firebase.crashlytics.internal.settings.SettingsProvider;
 import com.google.firebase.crashlytics.internal.stacktrace.StackTraceTrimmingStrategy;
+import java.io.BufferedInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.IOException;
@@ -427,13 +428,15 @@ private static CrashlyticsReport.ApplicationExitInfo convertApplicationExitInfo(
   @VisibleForTesting
   @RequiresApi(api = Build.VERSION_CODES.KITKAT)
   public static String convertInputStreamToString(InputStream inputStream) throws IOException {
-    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
-    byte[] bytes = new byte[DEFAULT_BUFFER_SIZE];
-    int length;
-    while ((length = inputStream.read(bytes)) != -1) {
-      byteArrayOutputStream.write(bytes, 0, length);
+    try (BufferedInputStream bufferedInputStream = new BufferedInputStream(inputStream);
+        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream()) {
+      byte[] bytes = new byte[DEFAULT_BUFFER_SIZE];
+      int length;
+      while ((length = bufferedInputStream.read(bytes)) != -1) {
+        byteArrayOutputStream.write(bytes, 0, length);
+      }
+      return byteArrayOutputStream.toString(StandardCharsets.UTF_8.name());
     }
-    return byteArrayOutputStream.toString(StandardCharsets.UTF_8.name());
   }
 
   /** Finds the first ANR ApplicationExitInfo within the session. */

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/CompositeIndexTestHelper.java

@@ -125,7 +125,11 @@ public void assertOnlineAndOfflineResultsMatch(
       @NonNull CollectionReference collection,
       @NonNull Query query,
       @NonNull String... expectedDocs) {
-    checkOnlineAndOfflineResultsMatch(collection, query, toHashedIds(expectedDocs));
+    // `checkOnlineAndOfflineResultsMatch` first makes sure all documents needed for
+    // `query` are in the cache. It does so making a `get` on the first argument.
+    // Since *all* composite index tests use the same collection, this is very inefficient to do.
+    // Therefore, we should only do so for tests where `TEST_ID_FIELD` matches the current test.
+    checkOnlineAndOfflineResultsMatch(this.query(collection), query, toHashedIds(expectedDocs));
   }
 
   // Asserts that the IDs in the query snapshot matches the expected Ids. The expected document

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/IntegrationTestUtil.java

@@ -535,7 +535,7 @@ public static List<Object> nullList() {
    * @param expectedDocs Ordered list of document keys that are expected to match the query
    */
   public static void checkOnlineAndOfflineResultsMatch(
-      CollectionReference collection, Query query, String... expectedDocs) {
+      Query collection, Query query, String... expectedDocs) {
     // Note: Order matters. The following has to be done in the specific order:
 
     // 1- Pre-populate the cache with the entire collection.

firebase-vertexai/CHANGELOG.md

@@ -1,7 +1,19 @@
 # Unreleased
+* [changed] **Breaking Change**: `LiveModelFutures.connect` now returns `ListenableFuture<LiveSessionFutures>` instead of `ListenableFuture<LiveSession>`.
+    * **Action Required:** Remove any transformations from LiveSession object to LiveSessionFutures object. 
+    * **Action Required:** Change type of variable handling `LiveModelFutures.connect` to `ListenableFuture<LiveSessionsFutures>`
+* [changed] **Breaking Change**: Removed `UNSPECIFIED` value for enum class `ResponseModality`
+    * **Action Required:** Remove all references to `ResponseModality.UNSPECIFIED`
+* [changed] **Breaking Change**: Renamed `LiveGenerationConfig.setResponseModalities` to `LiveGenerationConfig.setResponseModality`
+    * **Action Required:** Replace all references of `LiveGenerationConfig.setResponseModalities` with `LiveGenerationConfig.setResponseModality`
 * [feature] Added support for `HarmBlockThreshold.OFF`. See the
   [model documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters#how_to_configure_content_filters){: .external}
   for more information.
+* [fixed] Improved thread usage when using a `LiveGenerativeModel`. (#6870)
+* [fixed] Fixed an issue with `LiveContentResponse` audio data not being present when the model was
+  interrupted or the turn completed. (#6870)
+* [fixed] Fixed an issue with `LiveSession` not converting exceptions to `FirebaseVertexAIException`. (#6870)
+
 
 # 16.3.0
 * [feature] Emits a warning when attempting to use an incompatible model with

firebase-vertexai/api.txt

@@ -115,7 +115,7 @@ package com.google.firebase.vertexai.java {
   }
 
   @com.google.firebase.vertexai.type.PublicPreviewAPI public abstract class LiveModelFutures {
-    method public abstract com.google.common.util.concurrent.ListenableFuture<com.google.firebase.vertexai.type.LiveSession> connect();
+    method public abstract com.google.common.util.concurrent.ListenableFuture<com.google.firebase.vertexai.java.LiveSessionFutures> connect();
     method public static final com.google.firebase.vertexai.java.LiveModelFutures from(com.google.firebase.vertexai.LiveGenerativeModel model);
     field public static final com.google.firebase.vertexai.java.LiveModelFutures.Companion Companion;
   }
@@ -132,6 +132,7 @@ package com.google.firebase.vertexai.java {
     method public abstract com.google.common.util.concurrent.ListenableFuture<kotlin.Unit> send(String text);
     method public abstract com.google.common.util.concurrent.ListenableFuture<kotlin.Unit> sendFunctionResponse(java.util.List<com.google.firebase.vertexai.type.FunctionResponsePart> functionList);
     method public abstract com.google.common.util.concurrent.ListenableFuture<kotlin.Unit> sendMediaStream(java.util.List<com.google.firebase.vertexai.type.MediaData> mediaChunks);
+    method public abstract com.google.common.util.concurrent.ListenableFuture<kotlin.Unit> startAudioConversation();
     method public abstract com.google.common.util.concurrent.ListenableFuture<kotlin.Unit> startAudioConversation(kotlin.jvm.functions.Function1<? super com.google.firebase.vertexai.type.FunctionCallPart,com.google.firebase.vertexai.type.FunctionResponsePart>? functionCallHandler);
     method public abstract com.google.common.util.concurrent.ListenableFuture<kotlin.Unit> stopAudioConversation();
     method public abstract void stopReceiving();
@@ -597,7 +598,7 @@ package com.google.firebase.vertexai.type {
     method public com.google.firebase.vertexai.type.LiveGenerationConfig.Builder setFrequencyPenalty(Float? frequencyPenalty);
     method public com.google.firebase.vertexai.type.LiveGenerationConfig.Builder setMaxOutputTokens(Integer? maxOutputTokens);
     method public com.google.firebase.vertexai.type.LiveGenerationConfig.Builder setPresencePenalty(Float? presencePenalty);
-    method public com.google.firebase.vertexai.type.LiveGenerationConfig.Builder setResponseModalities(com.google.firebase.vertexai.type.ResponseModality? responseModalities);
+    method public com.google.firebase.vertexai.type.LiveGenerationConfig.Builder setResponseModality(com.google.firebase.vertexai.type.ResponseModality? responseModality);
     method public com.google.firebase.vertexai.type.LiveGenerationConfig.Builder setSpeechConfig(com.google.firebase.vertexai.type.SpeechConfig? speechConfig);
     method public com.google.firebase.vertexai.type.LiveGenerationConfig.Builder setTemperature(Float? temperature);
     method public com.google.firebase.vertexai.type.LiveGenerationConfig.Builder setTopK(Integer? topK);
@@ -628,7 +629,7 @@ package com.google.firebase.vertexai.type {
     method public suspend Object? send(String text, kotlin.coroutines.Continuation<? super kotlin.Unit>);
     method public suspend Object? sendFunctionResponse(java.util.List<com.google.firebase.vertexai.type.FunctionResponsePart> functionList, kotlin.coroutines.Continuation<? super kotlin.Unit>);
     method public suspend Object? sendMediaStream(java.util.List<com.google.firebase.vertexai.type.MediaData> mediaChunks, kotlin.coroutines.Continuation<? super kotlin.Unit>);
-    method public suspend Object? startAudioConversation(kotlin.jvm.functions.Function1<? super com.google.firebase.vertexai.type.FunctionCallPart,com.google.firebase.vertexai.type.FunctionResponsePart>? functionCallHandler = null, kotlin.coroutines.Continuation<? super kotlin.Unit>);
+    method @RequiresPermission(android.Manifest.permission.RECORD_AUDIO) public suspend Object? startAudioConversation(kotlin.jvm.functions.Function1<? super com.google.firebase.vertexai.type.FunctionCallPart,com.google.firebase.vertexai.type.FunctionResponsePart>? functionCallHandler = null, kotlin.coroutines.Continuation<? super kotlin.Unit>);
     method public void stopAudioConversation();
     method public void stopReceiving();
   }
@@ -696,7 +697,6 @@ package com.google.firebase.vertexai.type {
     field public static final com.google.firebase.vertexai.type.ResponseModality.Companion Companion;
     field public static final com.google.firebase.vertexai.type.ResponseModality IMAGE;
     field public static final com.google.firebase.vertexai.type.ResponseModality TEXT;
-    field public static final com.google.firebase.vertexai.type.ResponseModality UNSPECIFIED;
   }
 
   public static final class ResponseModality.Companion {

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/FirebaseVertexAI.kt

@@ -19,7 +19,7 @@ package com.google.firebase.vertexai
 import android.util.Log
 import com.google.firebase.Firebase
 import com.google.firebase.FirebaseApp
-import com.google.firebase.annotations.concurrent.Background
+import com.google.firebase.annotations.concurrent.Blocking
 import com.google.firebase.app
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
@@ -41,7 +41,7 @@ import kotlin.coroutines.CoroutineContext
 public class FirebaseVertexAI
 internal constructor(
   private val firebaseApp: FirebaseApp,
-  @Background private val backgroundDispatcher: CoroutineContext,
+  @Blocking private val blockingDispatcher: CoroutineContext,
   private val location: String,
   private val appCheckProvider: Provider<InteropAppCheckTokenProvider>,
   private val internalAuthProvider: Provider<InternalAuthProvider>,
@@ -133,7 +133,7 @@ internal constructor(
       "projects/${firebaseApp.options.projectId}/locations/${location}/publishers/google/models/${modelName}",
       firebaseApp.options.apiKey,
       firebaseApp,
-      backgroundDispatcher,
+      blockingDispatcher,
       generationConfig,
       tools,
       systemInstruction,

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/FirebaseVertexAIMultiResourceComponent.kt

@@ -18,7 +18,7 @@ package com.google.firebase.vertexai
 
 import androidx.annotation.GuardedBy
 import com.google.firebase.FirebaseApp
-import com.google.firebase.annotations.concurrent.Background
+import com.google.firebase.annotations.concurrent.Blocking
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.inject.Provider
@@ -31,7 +31,7 @@ import kotlin.coroutines.CoroutineContext
  */
 internal class FirebaseVertexAIMultiResourceComponent(
   private val app: FirebaseApp,
-  @Background val backgroundDispatcher: CoroutineContext,
+  @Blocking val blockingDispatcher: CoroutineContext,
   private val appCheckProvider: Provider<InteropAppCheckTokenProvider>,
   private val internalAuthProvider: Provider<InternalAuthProvider>,
 ) {
@@ -43,7 +43,7 @@ internal class FirebaseVertexAIMultiResourceComponent(
       instances[location]
         ?: FirebaseVertexAI(
             app,
-            backgroundDispatcher,
+            blockingDispatcher,
             location,
             appCheckProvider,
             internalAuthProvider

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/FirebaseVertexAIRegistrar.kt

@@ -18,7 +18,7 @@ package com.google.firebase.vertexai
 
 import androidx.annotation.Keep
 import com.google.firebase.FirebaseApp
-import com.google.firebase.annotations.concurrent.Background
+import com.google.firebase.annotations.concurrent.Blocking
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.components.Component
@@ -41,13 +41,13 @@ internal class FirebaseVertexAIRegistrar : ComponentRegistrar {
       Component.builder(FirebaseVertexAIMultiResourceComponent::class.java)
         .name(LIBRARY_NAME)
         .add(Dependency.required(firebaseApp))
-        .add(Dependency.required(backgroundDispatcher))
+        .add(Dependency.required(blockingDispatcher))
         .add(Dependency.optionalProvider(appCheckInterop))
         .add(Dependency.optionalProvider(internalAuthProvider))
         .factory { container ->
           FirebaseVertexAIMultiResourceComponent(
             container[firebaseApp],
-            container.get(backgroundDispatcher),
+            container.get(blockingDispatcher),
             container.getProvider(appCheckInterop),
             container.getProvider(internalAuthProvider)
           )
@@ -62,7 +62,7 @@ internal class FirebaseVertexAIRegistrar : ComponentRegistrar {
     private val firebaseApp = unqualified(FirebaseApp::class.java)
     private val appCheckInterop = unqualified(InteropAppCheckTokenProvider::class.java)
     private val internalAuthProvider = unqualified(InternalAuthProvider::class.java)
-    private val backgroundDispatcher =
-      Qualified.qualified(Background::class.java, CoroutineDispatcher::class.java)
+    private val blockingDispatcher =
+      Qualified.qualified(Blocking::class.java, CoroutineDispatcher::class.java)
   }
 }

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/LiveGenerativeModel.kt

@@ -17,13 +17,14 @@
 package com.google.firebase.vertexai
 
 import com.google.firebase.FirebaseApp
-import com.google.firebase.annotations.concurrent.Background
+import com.google.firebase.annotations.concurrent.Blocking
 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.AppCheckHeaderProvider
-import com.google.firebase.vertexai.type.BidiGenerateContentClientMessage
+import com.google.firebase.vertexai.common.JSON
 import com.google.firebase.vertexai.type.Content
+import com.google.firebase.vertexai.type.LiveClientSetupMessage
 import com.google.firebase.vertexai.type.LiveGenerationConfig
 import com.google.firebase.vertexai.type.LiveSession
 import com.google.firebase.vertexai.type.PublicPreviewAPI
@@ -38,6 +39,7 @@ import kotlinx.coroutines.channels.ClosedReceiveChannelException
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonObject
 
 /**
  * Represents a multimodal model (like Gemini) capable of real-time content generation based on
@@ -47,7 +49,7 @@ import kotlinx.serialization.json.Json
 public class LiveGenerativeModel
 internal constructor(
   private val modelName: String,
-  @Background private val backgroundDispatcher: CoroutineContext,
+  @Blocking private val blockingDispatcher: CoroutineContext,
   private val config: LiveGenerationConfig? = null,
   private val tools: List<Tool>? = null,
   private val systemInstruction: Content? = null,
@@ -58,7 +60,7 @@ internal constructor(
     modelName: String,
     apiKey: String,
     firebaseApp: FirebaseApp,
-    backgroundDispatcher: CoroutineContext,
+    blockingDispatcher: CoroutineContext,
     config: LiveGenerationConfig? = null,
     tools: List<Tool>? = null,
     systemInstruction: Content? = null,
@@ -68,7 +70,7 @@ internal constructor(
     internalAuthProvider: InternalAuthProvider? = null,
   ) : this(
     modelName,
-    backgroundDispatcher,
+    blockingDispatcher,
     config,
     tools,
     systemInstruction,
@@ -93,7 +95,7 @@ internal constructor(
   @OptIn(ExperimentalSerializationApi::class)
   public suspend fun connect(): LiveSession {
     val clientMessage =
-      BidiGenerateContentClientMessage(
+      LiveClientSetupMessage(
           modelName,
           config?.toInternal(),
           tools?.map { it.toInternal() },
@@ -104,10 +106,11 @@ internal constructor(
     try {
       val webSession = controller.getWebSocketSession(location)
       webSession.send(Frame.Text(data))
-      val receivedJson = webSession.incoming.receive().readBytes().toString(Charsets.UTF_8)
-      // TODO: Try to decode the json instead of string matching.
-      return if (receivedJson.contains("setupComplete")) {
-        LiveSession(session = webSession, backgroundDispatcher = backgroundDispatcher)
+      val receivedJsonStr = webSession.incoming.receive().readBytes().toString(Charsets.UTF_8)
+      val receivedJson = JSON.parseToJsonElement(receivedJsonStr)
+
+      return if (receivedJson is JsonObject && "setupComplete" in receivedJson) {
+        LiveSession(session = webSession, blockingDispatcher = blockingDispatcher)
       } else {
         webSession.close()
         throw ServiceConnectionHandshakeFailedException("Unable to connect to the server")

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/APIController.kt

@@ -165,6 +165,7 @@ internal constructor(
 
   suspend fun getWebSocketSession(location: String): ClientWebSocketSession =
     client.webSocketSession(getBidiEndpoint(location))
+
   fun generateContentStream(
     request: GenerateContentRequest
   ): Flow<GenerateContentResponse.Internal> =