Add support for GoogleAI to firebase-vertexai

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

@@ -0,0 +1,35 @@
+/*
+ * 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
+
+import com.google.firebase.Firebase
+import com.google.firebase.FirebaseApp
+import com.google.firebase.app
+
+public class FirebaseGenAI {
+  public fun vertexAI(app: FirebaseApp = Firebase.app, location: String = "us-central1"): FirebaseVertexAI =
+    FirebaseVertexAI.getInstance(app, location)
+
+  public fun googleAI(app: FirebaseApp = Firebase.app): FirebaseGoogleAI = FirebaseGoogleAI.getInstance(app)
+}
+
+/** Returns the [FirebaseGenAI] instance. */
+public val Firebase.genAI: FirebaseGenAI
+  get() = FirebaseGenAI()
+
+/** Returns the [FirebaseGenAI] instance. */
+public fun Firebase.genAI(): FirebaseGenAI = FirebaseGenAI()

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

@@ -0,0 +1,99 @@
+/*
+ * 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
+
+import com.google.firebase.Firebase
+import com.google.firebase.FirebaseApp
+import com.google.firebase.app
+import com.google.firebase.vertexai.type.Content
+import com.google.firebase.vertexai.type.GenerationConfig
+import com.google.firebase.vertexai.type.GenerativeBackend
+import com.google.firebase.vertexai.type.ImagenGenerationConfig
+import com.google.firebase.vertexai.type.ImagenSafetySettings
+import com.google.firebase.vertexai.type.PublicPreviewAPI
+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
+
+public class FirebaseGoogleAI internal constructor(private val proxy: FirebaseVertexAI) {
+
+  /**
+   * Instantiates a new [GenerativeModel] given the provided parameters.
+   *
+   * @param modelName The name of the model to use, for example `"gemini-1.5-pro"`.
+   * @param generationConfig The configuration parameters to use for content generation.
+   * @param safetySettings The safety bounds the model will abide to during content generation.
+   * @param tools A list of [Tool]s the model may use to generate content.
+   * @param toolConfig The [ToolConfig] that defines how the model handles the tools provided.
+   * @param systemInstruction [Content] instructions that direct the model to behave a certain way.
+   * Currently only text content is supported.
+   * @param requestOptions Configuration options for sending requests to the backend.
+   * @return The initialized [GenerativeModel] instance.
+   */
+  @JvmOverloads
+  public fun generativeModel(
+    modelName: String,
+    generationConfig: GenerationConfig? = null,
+    safetySettings: List<SafetySetting>? = null,
+    tools: List<Tool>? = null,
+    toolConfig: ToolConfig? = null,
+    systemInstruction: Content? = null,
+    requestOptions: RequestOptions = RequestOptions(),
+  ): GenerativeModel =
+    proxy.generativeModel(
+      modelName,
+      generationConfig,
+      safetySettings,
+      tools,
+      toolConfig,
+      systemInstruction,
+      requestOptions,
+    )
+
+  /**
+   * Instantiates a new [ImagenModel] given the provided parameters.
+   *
+   * @param modelName The name of the model to use, for example `"imagen-3.0-generate-001"`.
+   * @param generationConfig The configuration parameters to use for image generation.
+   * @param safetySettings The safety bounds the model will abide by during image generation.
+   * @param requestOptions Configuration options for sending requests to the backend.
+   * @return The initialized [ImagenModel] instance.
+   */
+  @JvmOverloads
+  @PublicPreviewAPI
+  public fun imagenModel(
+    modelName: String,
+    generationConfig: ImagenGenerationConfig? = null,
+    safetySettings: ImagenSafetySettings? = null,
+    requestOptions: RequestOptions = RequestOptions(),
+  ): ImagenModel = proxy.imagenModel(modelName, generationConfig, safetySettings, requestOptions)
+
+  public companion object {
+    /** The [FirebaseGoogleAI] instance for the default [FirebaseApp] */
+    @JvmStatic
+    public val instance: FirebaseGoogleAI
+      get() = getInstance()
+
+    /** Returns the [FirebaseGoogleAI] instance for the provided [FirebaseApp]. */
+    @JvmStatic
+    public fun getInstance(app: FirebaseApp = Firebase.app): FirebaseGoogleAI {
+      val multiResourceComponent = app[FirebaseVertexAIMultiResourceComponent::class.java]
+      return FirebaseGoogleAI(multiResourceComponent.get(GenerativeBackend.GOOGLE_AI, "UNUSED"))
+    }
+  }
+}

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

@@ -24,6 +24,7 @@ import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.inject.Provider
 import com.google.firebase.vertexai.type.Content
 import com.google.firebase.vertexai.type.GenerationConfig
+import com.google.firebase.vertexai.type.GenerativeBackend
 import com.google.firebase.vertexai.type.ImagenGenerationConfig
 import com.google.firebase.vertexai.type.ImagenSafetySettings
 import com.google.firebase.vertexai.type.InvalidLocationException
@@ -37,6 +38,7 @@ import com.google.firebase.vertexai.type.ToolConfig
 public class FirebaseVertexAI
 internal constructor(
   private val firebaseApp: FirebaseApp,
+  private val backend: GenerativeBackend,
   private val location: String,
   private val appCheckProvider: Provider<InteropAppCheckTokenProvider>,
   private val internalAuthProvider: Provider<InternalAuthProvider>,
@@ -68,15 +70,23 @@ internal constructor(
     if (location.trim().isEmpty() || location.contains("/")) {
       throw InvalidLocationException(location)
     }
+    val modelUri =
+      when (backend) {
+        GenerativeBackend.VERTEX_AI ->
+          "projects/${firebaseApp.options.projectId}/locations/${location}/publishers/google/models/${modelName}"
+        GenerativeBackend.GOOGLE_AI ->
+          "projects/${firebaseApp.options.projectId}/models/${modelName}"
+      }
     return GenerativeModel(
-      "projects/${firebaseApp.options.projectId}/locations/${location}/publishers/google/models/${modelName}",
+      modelUri,
       firebaseApp.options.apiKey,
       generationConfig,
       safetySettings,
       tools,
       toolConfig,
       systemInstruction,
       requestOptions,
+      backend,
       appCheckProvider.get(),
       internalAuthProvider.get(),
     )
@@ -117,9 +127,7 @@ internal constructor(
     /** The [FirebaseVertexAI] instance for the default [FirebaseApp] */
     @JvmStatic
     public val instance: FirebaseVertexAI
-      get() = getInstance(location = "us-central1")
-
-    @JvmStatic public fun getInstance(app: FirebaseApp): FirebaseVertexAI = getInstance(app)
+      get() = getInstance(Firebase.app)
 
     /**
      * Returns the [FirebaseVertexAI] instance for the provided [FirebaseApp] and [location].
@@ -130,9 +138,9 @@ internal constructor(
      */
     @JvmStatic
     @JvmOverloads
-    public fun getInstance(app: FirebaseApp = Firebase.app, location: String): FirebaseVertexAI {
+    public fun getInstance(app: FirebaseApp, location: String = "us-central1"): FirebaseVertexAI {
       val multiResourceComponent = app[FirebaseVertexAIMultiResourceComponent::class.java]
-      return multiResourceComponent.get(location)
+      return multiResourceComponent.get(GenerativeBackend.VERTEX_AI, location)
     }
   }
 }
@@ -144,5 +152,5 @@ public val Firebase.vertexAI: FirebaseVertexAI
 /** Returns the [FirebaseVertexAI] instance of a given [FirebaseApp]. */
 public fun Firebase.vertexAI(
   app: FirebaseApp = Firebase.app,
-  location: String = "us-central1"
+  location: String = "us-central1",
 ): FirebaseVertexAI = FirebaseVertexAI.getInstance(app, location)

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

@@ -21,6 +21,7 @@ import com.google.firebase.FirebaseApp
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.inject.Provider
+import com.google.firebase.vertexai.type.GenerativeBackend
 
 /**
  * Multi-resource container for Firebase Vertex AI.
@@ -35,11 +36,16 @@ internal class FirebaseVertexAIMultiResourceComponent(
 
   @GuardedBy("this") private val instances: MutableMap<String, FirebaseVertexAI> = mutableMapOf()
 
-  fun get(location: String): FirebaseVertexAI =
+  fun get(generativeBackend: GenerativeBackend, location: String): FirebaseVertexAI =
     synchronized(this) {
       instances[location]
-        ?: FirebaseVertexAI(app, location, appCheckProvider, internalAuthProvider).also {
-          instances[location] = it
-        }
+        ?: FirebaseVertexAI(
+            app,
+            generativeBackend,
+            location,
+            appCheckProvider,
+            internalAuthProvider
+          )
+          .also { instances[location] = it }
     }
 }

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

@@ -29,6 +29,7 @@ import com.google.firebase.vertexai.type.FinishReason
 import com.google.firebase.vertexai.type.FirebaseVertexAIException
 import com.google.firebase.vertexai.type.GenerateContentResponse
 import com.google.firebase.vertexai.type.GenerationConfig
+import com.google.firebase.vertexai.type.GenerativeBackend
 import com.google.firebase.vertexai.type.PromptBlockedException
 import com.google.firebase.vertexai.type.RequestOptions
 import com.google.firebase.vertexai.type.ResponseStoppedException
@@ -54,6 +55,7 @@ internal constructor(
   private val tools: List<Tool>? = null,
   private val toolConfig: ToolConfig? = null,
   private val systemInstruction: Content? = null,
+  private val generativeBackend: GenerativeBackend = GenerativeBackend.VERTEX_AI,
   private val controller: APIController,
 ) {
   internal constructor(
@@ -65,6 +67,7 @@ internal constructor(
     toolConfig: ToolConfig? = null,
     systemInstruction: Content? = null,
     requestOptions: RequestOptions = RequestOptions(),
+    generativeBackend: GenerativeBackend,
     appCheckTokenProvider: InteropAppCheckTokenProvider? = null,
     internalAuthProvider: InternalAuthProvider? = null,
   ) : this(
@@ -74,6 +77,7 @@ internal constructor(
     tools,
     toolConfig,
     systemInstruction,
+    generativeBackend,
     APIController(
       apiKey,
       modelName,
@@ -213,7 +217,10 @@ internal constructor(
     )
 
   private fun constructCountTokensRequest(vararg prompt: Content) =
-    CountTokensRequest.forVertexAI(constructRequest(*prompt))
+    when (generativeBackend) {
+      GenerativeBackend.GOOGLE_AI -> CountTokensRequest.forGenAI(constructRequest(*prompt))
+      GenerativeBackend.VERTEX_AI -> CountTokensRequest.forVertexAI(constructRequest(*prompt))
+    }
 
   private fun GenerateContentResponse.validate() = apply {
     if (candidates.isEmpty() && promptFeedback == null) {

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

@@ -53,6 +53,15 @@ internal data class CountTokensRequest(
 ) : Request {
   companion object {
 
+    fun forGenAI(generateContentRequest: GenerateContentRequest) =
+      CountTokensRequest(
+        generateContentRequest =
+          generateContentRequest.model?.let {
+            generateContentRequest.copy(model = fullModelName(it))
+          }
+            ?: generateContentRequest
+      )
+
     fun forVertexAI(generateContentRequest: GenerateContentRequest) =
       CountTokensRequest(
         model = generateContentRequest.model?.let { fullModelName(it) },

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

@@ -0,0 +1,22 @@
+/*
+ * 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.type
+
+internal enum class GenerativeBackend {
+  VERTEX_AI,
+  GOOGLE_AI,
+}

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

@@ -37,6 +37,8 @@ internal constructor(
    */
   @JvmOverloads
   public constructor(
-    timeoutInMillis: Long = 180.seconds.inWholeMilliseconds
-  ) : this(timeout = timeoutInMillis.toDuration(DurationUnit.MILLISECONDS))
+    timeoutInMillis: Long = 180.seconds.inWholeMilliseconds,
+  ) : this(
+    timeout = timeoutInMillis.toDuration(DurationUnit.MILLISECONDS),
+  )
 }