Refactor live bidi

firebase-vertexai/CHANGELOG.md

@@ -9,6 +9,11 @@
 * [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

@@ -629,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();
   }

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> =

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

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.firebase.vertexai.common.util
+
+import android.media.AudioRecord
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.yield
+
+/**
+ * The minimum buffer size for this instance.
+ *
+ * The same as calling [AudioRecord.getMinBufferSize], except the params are pre-populated.
+ */
+internal val AudioRecord.minBufferSize: Int
+  get() = AudioRecord.getMinBufferSize(sampleRate, channelConfiguration, audioFormat)
+
+/**
+ * Reads from this [AudioRecord] and returns the data in a flow.
+ *
+ * Will yield when this instance is not recording.
+ */
+internal fun AudioRecord.readAsFlow() = flow {
+  val buffer = ByteArray(minBufferSize)
+
+  while (true) {
+    if (recordingState != AudioRecord.RECORDSTATE_RECORDING) {
+      yield()
+      continue
+    }
+
+    val bytesRead = read(buffer, 0, buffer.size)
+    if (bytesRead > 0) {
+      emit(buffer.copyOf(bytesRead))
+    }
+  }
+}

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

@@ -16,7 +16,16 @@
 
 package com.google.firebase.vertexai.common.util
 
+import java.io.ByteArrayOutputStream
 import java.lang.reflect.Field
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.fold
 
 /**
  * Removes the last character from the [StringBuilder].
@@ -39,3 +48,56 @@ internal fun StringBuilder.removeLast(): StringBuilder =
  * ```
  */
 internal inline fun <reified T : Annotation> Field.getAnnotation() = getAnnotation(T::class.java)
+
+/**
+ * Collects bytes from this flow and doesn't emit them back until [minSize] is reached.
+ *
+ * For example:
+ * ```
+ * val byteArr = flowOf(byteArrayOf(1), byteArrayOf(2, 3, 4), byteArrayOf(5, 6, 7, 8))
+ * val expectedResult = listOf(byteArrayOf(1, 2, 3, 4), byteArrayOf( 5, 6, 7, 8))
+ *
+ * byteArr.accumulateUntil(4).toList() shouldContainExactly expectedResult
+ * ```
+ *
+ * @param minSize The minimum about of bytes the array should have before being sent down-stream
+ * @param emitLeftOvers If the flow completes and there are bytes left over that don't meet the
+ * [minSize], send them anyways.
+ */
+internal fun Flow<ByteArray>.accumulateUntil(
+  minSize: Int,
+  emitLeftOvers: Boolean = false
+): Flow<ByteArray> = flow {
+  val remaining =
+    fold(ByteArrayOutputStream()) { buffer, it ->
+      buffer.apply {
+        write(it, 0, it.size)
+        if (size() >= minSize) {
+          emit(toByteArray())
+          reset()
+        }
+      }
+    }
+
+  if (emitLeftOvers && remaining.size() > 0) {
+    emit(remaining.toByteArray())
+  }
+}
+
+/**
+ * Create a [Job] that is a child of the [currentCoroutineContext], if any.
+ *
+ * This is useful when you want a coroutine scope to be canceled when its parent scope is canceled,
+ * and you don't have full control over the parent scope, but you don't want the cancellation of the
+ * child to impact the parent.
+ *
+ * If the parent coroutine context does not have a job, an empty one will be created.
+ */
+internal suspend inline fun childJob() = Job(currentCoroutineContext()[Job] ?: Job())
+
+/**
+ * A constant value pointing to a cancelled [CoroutineScope].
+ *
+ * Useful when you want to initialize a mutable [CoroutineScope] in a canceled state.
+ */
+internal val CancelledCoroutineScope = CoroutineScope(EmptyCoroutineContext).apply { cancel() }