add background dispatcher

Author: VinayGuthal

firebase-vertexai/api.txt

@@ -818,14 +818,6 @@ package com.google.firebase.vertexai.type {
     property public final com.google.firebase.vertexai.type.Voices voice;
   }
 
-  public final class StreamOutput {
-    ctor public StreamOutput(boolean interrupted, com.google.firebase.vertexai.type.Content? data);
-    method public com.google.firebase.vertexai.type.Content? getData();
-    method public boolean getInterrupted();
-    property public final com.google.firebase.vertexai.type.Content? data;
-    property public final boolean interrupted;
-  }
-
   public abstract class StringFormat {
   }
 

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

@@ -19,6 +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.app
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
@@ -34,11 +35,13 @@ import com.google.firebase.vertexai.type.RequestOptions
 import com.google.firebase.vertexai.type.SafetySetting
 import com.google.firebase.vertexai.type.Tool
 import com.google.firebase.vertexai.type.ToolConfig
+import kotlin.coroutines.CoroutineContext
 
 /** Entry point for all _Vertex AI for Firebase_ functionality. */
 public class FirebaseVertexAI
 internal constructor(
   private val firebaseApp: FirebaseApp,
+  @Background private val backgroundDispatcher: CoroutineContext,
   private val location: String,
   private val appCheckProvider: Provider<InteropAppCheckTokenProvider>,
   private val internalAuthProvider: Provider<InternalAuthProvider>,
@@ -130,6 +133,7 @@ internal constructor(
       "projects/${firebaseApp.options.projectId}/locations/${location}/publishers/google/models/${modelName}",
       firebaseApp.options.apiKey,
       firebaseApp,
+      backgroundDispatcher,
       generationConfig,
       tools,
       systemInstruction,

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

@@ -18,9 +18,11 @@ 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.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.inject.Provider
+import kotlin.coroutines.CoroutineContext
 
 /**
  * Multi-resource container for Firebase Vertex AI.
@@ -29,17 +31,23 @@ import com.google.firebase.inject.Provider
  */
 internal class FirebaseVertexAIMultiResourceComponent(
   private val app: FirebaseApp,
+  @Background val backgroundDispatcher: CoroutineContext,
   private val appCheckProvider: Provider<InteropAppCheckTokenProvider>,
-  private val internalAuthProvider: Provider<InternalAuthProvider>
+  private val internalAuthProvider: Provider<InternalAuthProvider>,
 ) {
 
   @GuardedBy("this") private val instances: MutableMap<String, FirebaseVertexAI> = mutableMapOf()
 
   fun get(location: String): FirebaseVertexAI =
     synchronized(this) {
       instances[location]
-        ?: FirebaseVertexAI(app, location, appCheckProvider, internalAuthProvider).also {
-          instances[location] = it
-        }
+        ?: FirebaseVertexAI(
+            app,
+            backgroundDispatcher,
+            location,
+            appCheckProvider,
+            internalAuthProvider
+          )
+          .also { instances[location] = it }
     }
 }

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

@@ -18,13 +18,16 @@ 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.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.components.Component
 import com.google.firebase.components.ComponentRegistrar
 import com.google.firebase.components.Dependency
+import com.google.firebase.components.Qualified
 import com.google.firebase.components.Qualified.unqualified
 import com.google.firebase.platforminfo.LibraryVersionComponent
+import kotlinx.coroutines.CoroutineDispatcher
 
 /**
  * [ComponentRegistrar] for setting up [FirebaseVertexAI] and its internal dependencies.
@@ -38,11 +41,13 @@ internal class FirebaseVertexAIRegistrar : ComponentRegistrar {
       Component.builder(FirebaseVertexAIMultiResourceComponent::class.java)
         .name(LIBRARY_NAME)
         .add(Dependency.required(firebaseApp))
+        .add(Dependency.required(backgroundDispatcher))
         .add(Dependency.optionalProvider(appCheckInterop))
         .add(Dependency.optionalProvider(internalAuthProvider))
         .factory { container ->
           FirebaseVertexAIMultiResourceComponent(
             container[firebaseApp],
+            container.get(backgroundDispatcher),
             container.getProvider(appCheckInterop),
             container.getProvider(internalAuthProvider)
           )
@@ -57,5 +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)
   }
 }

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

@@ -17,6 +17,7 @@
 package com.google.firebase.vertexai
 
 import com.google.firebase.FirebaseApp
+import com.google.firebase.annotations.concurrent.Background
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.vertexai.common.APIController
@@ -31,6 +32,7 @@ import com.google.firebase.vertexai.type.Tool
 import io.ktor.websocket.Frame
 import io.ktor.websocket.close
 import io.ktor.websocket.readBytes
+import kotlin.coroutines.CoroutineContext
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
 
@@ -41,6 +43,7 @@ import kotlinx.serialization.json.Json
 public class LiveGenerativeModel
 internal constructor(
   private val modelName: String,
+  @Background private val backgroundDispatcher: CoroutineContext,
   private val config: LiveGenerationConfig? = null,
   private val tools: List<Tool>? = null,
   private val systemInstruction: Content? = null,
@@ -51,6 +54,7 @@ internal constructor(
     modelName: String,
     apiKey: String,
     firebaseApp: FirebaseApp,
+    backgroundDispatcher: CoroutineContext,
     config: LiveGenerationConfig? = null,
     tools: List<Tool>? = null,
     systemInstruction: Content? = null,
@@ -60,6 +64,7 @@ internal constructor(
     internalAuthProvider: InternalAuthProvider? = null,
   ) : this(
     modelName,
+    backgroundDispatcher,
     config,
     tools,
     systemInstruction,
@@ -94,7 +99,11 @@ internal constructor(
     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, isRecording = false)
+      LiveSession(
+        session = webSession,
+        backgroundDispatcher = backgroundDispatcher,
+        isRecording = false
+      )
     } else {
       webSession.close()
       throw GeminiConnectionHandshakeFailed()

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/LiveSession.kt

@@ -18,11 +18,13 @@ package com.google.firebase.vertexai.type
 
 import android.media.AudioFormat
 import android.media.AudioTrack
+import com.google.firebase.annotations.concurrent.Background
 import io.ktor.client.plugins.websocket.ClientWebSocketSession
 import io.ktor.websocket.Frame
 import io.ktor.websocket.close
 import io.ktor.websocket.readBytes
 import java.util.concurrent.ConcurrentLinkedQueue
+import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.cancel
@@ -42,6 +44,7 @@ import kotlinx.serialization.json.JsonNull
 public class LiveSession
 internal constructor(
   private val session: ClientWebSocketSession?,
+  @Background private val backgroundDispatcher: CoroutineContext,
   private var isRecording: Boolean,
   private var audioHelper: AudioHelper? = null
 ) {
@@ -133,7 +136,7 @@ internal constructor(
   }
 
   private fun fillRecordedAudioQueue() {
-    CoroutineScope(Dispatchers.IO).launch {
+    CoroutineScope(backgroundDispatcher).launch {
       audioHelper!!.startRecording().collect {
         if (!isRecording) {
           cancel()
@@ -159,7 +162,7 @@ internal constructor(
   }
 
   private fun fillServerResponseAudioQueue() {
-    CoroutineScope(Dispatchers.Default).launch {
+    CoroutineScope(backgroundDispatcher).launch {
       receive(listOf(ContentModality.AUDIO)).collect {
         if (!isRecording) {
           cancel()
@@ -182,7 +185,7 @@ internal constructor(
   }
 
   private fun playServerResponseAudio() {
-    CoroutineScope(Dispatchers.IO).launch {
+    CoroutineScope(backgroundDispatcher).launch {
       while (isRecording) {
         val x = playBackQueue.poll() ?: continue
         audioHelper?.playAudio(x)
@@ -202,7 +205,6 @@ internal constructor(
     isRecording = true
     audioHelper = AudioHelper()
     audioHelper!!.setupAudioTrack()
-    val scope = CoroutineScope(Dispatchers.Default)
     fillRecordedAudioQueue()
     CoroutineScope(Dispatchers.Default).launch { sendAudioDataToServer() }
     fillServerResponseAudioQueue()

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/MediaData.kt

@@ -20,7 +20,7 @@ import android.util.Base64
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 
-/* Class to represent the media data that needs to be sent to the server. */
+/* Represent the media data that needs to be sent to the server. */
 public class MediaData(public val mimeType: String, public val data: ByteArray) {
   @Serializable
   internal class Internal(@SerialName("mimeType") val mimeType: String, val data: String)