Merge GenAI Common into Vertex

build.gradle

@@ -41,6 +41,7 @@ buildscript {
         classpath 'com.google.firebase:firebase-appdistribution-gradle:5.0.0'
         classpath 'com.google.firebase:firebase-crashlytics-gradle:2.9.5'
         classpath "com.diffplug.spotless:spotless-plugin-gradle:7.0.0.BETA1"
+        classpath "org.jetbrains.kotlin:kotlin-serialization:1.8.22"
     }
 }
 

firebase-vertexai/CHANGELOG.md

@@ -1,4 +1,5 @@
 # Unreleased
+* [changed] Merged core networking code into VertexAI from a separate library
 * [feature] added support for `responseSchema` in `GenerationConfig`.
 
 # 16.0.0-beta03

firebase-vertexai/firebase-vertexai.gradle.kts

@@ -19,6 +19,7 @@
 plugins {
   id("firebase-library")
   id("kotlin-android")
+  kotlin("plugin.serialization")
 }
 
 firebaseLibrary {
@@ -56,12 +57,19 @@ android {
 }
 
 dependencies {
-  api("com.google.firebase:firebase-common:21.0.0")
+  val ktorVersion = "2.3.2"
+
+  implementation("io.ktor:ktor-client-okhttp:$ktorVersion")
+  implementation("io.ktor:ktor-client-core:$ktorVersion")
+  implementation("io.ktor:ktor-client-content-negotiation:$ktorVersion")
+  implementation("io.ktor:ktor-serialization-kotlinx-json:$ktorVersion")
+  implementation("io.ktor:ktor-client-logging:$ktorVersion")
+  compileOnly("io.ktor:ktor-client-mock:$ktorVersion")
 
+  implementation("com.google.firebase:firebase-common:21.0.0")
   implementation("com.google.firebase:firebase-components:18.0.0")
   implementation("com.google.firebase:firebase-annotations:16.2.0")
   implementation("com.google.firebase:firebase-appcheck-interop:17.1.0")
-  implementation("com.google.ai.client.generativeai:common:0.10.0")
   implementation(libs.androidx.annotation)
   implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.5.1")
   implementation("androidx.core:core-ktx:1.12.0")
@@ -74,9 +82,9 @@ dependencies {
   implementation("androidx.concurrent:concurrent-futures-ktx:1.2.0-alpha03")
   implementation("com.google.firebase:firebase-auth-interop:18.0.0")
 
-  val ktorVersion = "2.3.2"
   testImplementation("io.kotest:kotest-assertions-core:5.5.5")
   testImplementation("io.kotest:kotest-assertions-core-jvm:5.5.5")
+  testImplementation("io.kotest:kotest-assertions-json:5.5.5")
   testImplementation("io.ktor:ktor-client-okhttp:$ktorVersion")
   testImplementation("io.ktor:ktor-client-mock:$ktorVersion")
   testImplementation("org.json:json:20240303")

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

@@ -18,12 +18,12 @@ package com.google.firebase.vertexai
 
 import android.graphics.Bitmap
 import android.util.Log
-import com.google.ai.client.generativeai.common.APIController
-import com.google.ai.client.generativeai.common.CountTokensRequest
-import com.google.ai.client.generativeai.common.GenerateContentRequest
-import com.google.ai.client.generativeai.common.HeaderProvider
 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.CountTokensRequest
+import com.google.firebase.vertexai.common.GenerateContentRequest
+import com.google.firebase.vertexai.common.HeaderProvider
 import com.google.firebase.vertexai.internal.util.toInternal
 import com.google.firebase.vertexai.internal.util.toPublic
 import com.google.firebase.vertexai.type.Content

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

@@ -0,0 +1,277 @@
+/*
+ * Copyright 2024 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
+
+import android.util.Log
+import androidx.annotation.VisibleForTesting
+import com.google.firebase.vertexai.common.server.FinishReason
+import com.google.firebase.vertexai.common.util.decodeToFlow
+import com.google.firebase.vertexai.common.util.fullModelName
+import io.ktor.client.HttpClient
+import io.ktor.client.call.body
+import io.ktor.client.engine.HttpClientEngine
+import io.ktor.client.engine.mock.MockEngine
+import io.ktor.client.engine.mock.respond
+import io.ktor.client.engine.okhttp.OkHttp
+import io.ktor.client.plugins.HttpTimeout
+import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
+import io.ktor.client.request.HttpRequestBuilder
+import io.ktor.client.request.header
+import io.ktor.client.request.post
+import io.ktor.client.request.preparePost
+import io.ktor.client.request.setBody
+import io.ktor.client.statement.HttpResponse
+import io.ktor.client.statement.bodyAsChannel
+import io.ktor.client.statement.bodyAsText
+import io.ktor.http.ContentType
+import io.ktor.http.HttpHeaders
+import io.ktor.http.HttpStatusCode
+import io.ktor.http.contentType
+import io.ktor.http.headersOf
+import io.ktor.serialization.kotlinx.json.json
+import io.ktor.utils.io.ByteChannel
+import kotlin.time.Duration
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.catch
+import kotlinx.coroutines.flow.channelFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withTimeout
+import kotlinx.serialization.json.Json
+
+internal val JSON = Json {
+  ignoreUnknownKeys = true
+  prettyPrint = false
+  isLenient = true
+}
+
+/**
+ * Backend class for interfacing with the Gemini API.
+ *
+ * This class handles making HTTP requests to the API and streaming the responses back.
+ *
+ * @param httpEngine The HTTP client engine to be used for making requests. Defaults to CIO engine.
+ * Exposed primarily for DI in tests.
+ * @property key The API key used for authentication.
+ * @property model The model to use for generation.
+ * @property apiClient The value to pass in the `x-goog-api-client` header.
+ * @property headerProvider A provider that generates extra headers to include in all HTTP requests.
+ */
+internal class APIController
+internal constructor(
+  private val key: String,
+  model: String,
+  private val requestOptions: RequestOptions,
+  httpEngine: HttpClientEngine,
+  private val apiClient: String,
+  private val headerProvider: HeaderProvider?,
+) {
+
+  constructor(
+    key: String,
+    model: String,
+    requestOptions: RequestOptions,
+    apiClient: String,
+    headerProvider: HeaderProvider? = null,
+  ) : this(key, model, requestOptions, OkHttp.create(), apiClient, headerProvider)
+
+  @VisibleForTesting(otherwise = VisibleForTesting.NONE)
+  constructor(
+    key: String,
+    model: String,
+    requestOptions: RequestOptions,
+    apiClient: String,
+    headerProvider: HeaderProvider?,
+    channel: ByteChannel,
+    status: HttpStatusCode,
+  ) : this(
+    key,
+    model,
+    requestOptions,
+    MockEngine { respond(channel, status, headersOf(HttpHeaders.ContentType, "application/json")) },
+    apiClient,
+    headerProvider,
+  )
+
+  private val model = fullModelName(model)
+
+  private val client =
+    HttpClient(httpEngine) {
+      install(HttpTimeout) {
+        requestTimeoutMillis = requestOptions.timeout.inWholeMilliseconds
+        socketTimeoutMillis = 80_000
+      }
+      install(ContentNegotiation) { json(JSON) }
+    }
+
+  suspend fun generateContent(request: GenerateContentRequest): GenerateContentResponse =
+    try {
+      client
+        .post("${requestOptions.endpoint}/${requestOptions.apiVersion}/$model:generateContent") {
+          applyCommonConfiguration(request)
+          applyHeaderProvider()
+        }
+        .also { validateResponse(it) }
+        .body<GenerateContentResponse>()
+        .validate()
+    } catch (e: Throwable) {
+      throw FirebaseCommonAIException.from(e)
+    }
+
+  fun generateContentStream(request: GenerateContentRequest): Flow<GenerateContentResponse> =
+    client
+      .postStream<GenerateContentResponse>(
+        "${requestOptions.endpoint}/${requestOptions.apiVersion}/$model:streamGenerateContent?alt=sse"
+      ) {
+        applyCommonConfiguration(request)
+      }
+      .map { it.validate() }
+      .catch { throw FirebaseCommonAIException.from(it) }
+
+  suspend fun countTokens(request: CountTokensRequest): CountTokensResponse =
+    try {
+      client
+        .post("${requestOptions.endpoint}/${requestOptions.apiVersion}/$model:countTokens") {
+          applyCommonConfiguration(request)
+          applyHeaderProvider()
+        }
+        .also { validateResponse(it) }
+        .body()
+    } catch (e: Throwable) {
+      throw FirebaseCommonAIException.from(e)
+    }
+
+  private fun HttpRequestBuilder.applyCommonConfiguration(request: Request) {
+    when (request) {
+      is GenerateContentRequest -> setBody<GenerateContentRequest>(request)
+      is CountTokensRequest -> setBody<CountTokensRequest>(request)
+    }
+    contentType(ContentType.Application.Json)
+    header("x-goog-api-key", key)
+    header("x-goog-api-client", apiClient)
+  }
+
+  private suspend fun HttpRequestBuilder.applyHeaderProvider() {
+    if (headerProvider != null) {
+      try {
+        withTimeout(headerProvider.timeout) {
+          for ((tag, value) in headerProvider.generateHeaders()) {
+            header(tag, value)
+          }
+        }
+      } catch (e: TimeoutCancellationException) {
+        Log.w(TAG, "HeaderProvided timed out without generating headers, ignoring")
+      }
+    }
+  }
+
+  /**
+   * Makes a POST request to the specified [url] and returns a [Flow] of deserialized response
+   * objects of type [R]. The response is expected to be a stream of JSON objects that are parsed in
+   * real-time as they are received from the server.
+   *
+   * This function is intended for internal use within the client that handles streaming responses.
+   *
+   * Example usage:
+   * ```
+   * val client: HttpClient = HttpClient(CIO)
+   * val request: Request = GenerateContentRequest(...)
+   * val url: String = "http://example.com/stream"
+   *
+   * val responses: GenerateContentResponse = client.postStream(url) {
+   *   setBody(request)
+   *   contentType(ContentType.Application.Json)
+   * }
+   * responses.collect {
+   *   println("Got a response: $it")
+   * }
+   * ```
+   *
+   * @param R The type of the response object.
+   * @param url The URL to which the POST request will be made.
+   * @param config An optional [HttpRequestBuilder] callback for request configuration.
+   * @return A [Flow] of response objects of type [R].
+   */
+  private inline fun <reified R : Response> HttpClient.postStream(
+    url: String,
+    crossinline config: HttpRequestBuilder.() -> Unit = {},
+  ): Flow<R> = channelFlow {
+    launch(CoroutineName("postStream")) {
+      preparePost(url) {
+          applyHeaderProvider()
+          config()
+        }
+        .execute {
+          validateResponse(it)
+
+          val channel = it.bodyAsChannel()
+          val flow = JSON.decodeToFlow<R>(channel)
+
+          flow.collect { send(it) }
+        }
+    }
+  }
+
+  companion object {
+    private val TAG = APIController::class.java.simpleName
+  }
+}
+
+internal interface HeaderProvider {
+  val timeout: Duration
+
+  suspend fun generateHeaders(): Map<String, String>
+}
+
+private suspend fun validateResponse(response: HttpResponse) {
+  if (response.status == HttpStatusCode.OK) return
+  val text = response.bodyAsText()
+  val error =
+    try {
+      JSON.decodeFromString<GRpcErrorResponse>(text).error
+    } catch (e: Throwable) {
+      throw ServerException("Unexpected Response:\n$text $e")
+    }
+  val message = error.message
+  if (message.contains("API key not valid")) {
+    throw InvalidAPIKeyException(message)
+  }
+  // TODO (b/325117891): Use a better method than string matching.
+  if (message == "User location is not supported for the API use.") {
+    throw UnsupportedUserLocationException()
+  }
+  if (message.contains("quota")) {
+    throw QuotaExceededException(message)
+  }
+  if (error.details?.any { "SERVICE_DISABLED" == it.reason } == true) {
+    throw ServiceDisabledException(message)
+  }
+  throw ServerException(message)
+}
+
+private fun GenerateContentResponse.validate() = apply {
+  if ((candidates?.isEmpty() != false) && promptFeedback == null) {
+    throw SerializationException("Error deserializing response, found no valid fields")
+  }
+  promptFeedback?.blockReason?.let { throw PromptBlockedException(this) }
+  candidates
+    ?.mapNotNull { it.finishReason }
+    ?.firstOrNull { it != FinishReason.STOP }
+    ?.let { throw ResponseStoppedException(this) }
+}