Improvements to vertexAI types (#6309)

A collection of improvements to the VertexAI SDK:

- Make properties of `GenerativeModel` private
- Rename all `blob.*` to `inlineData.*`
- Improvements to `FunctionCallingConfig`
- Add support for `frequencyPenalty` and `presencePenalty`
- Add support for `HarmBlockMethod`
- Add support for `title` and `publicationDate` in `Citation
- Make error handling more robust when details are missing

---------

Co-authored-by: Daymon <[email protected]>

Author: rlazo

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

@@ -17,10 +17,10 @@
 package com.google.firebase.vertexai
 
 import android.graphics.Bitmap
-import com.google.firebase.vertexai.type.BlobPart
 import com.google.firebase.vertexai.type.Content
 import com.google.firebase.vertexai.type.GenerateContentResponse
 import com.google.firebase.vertexai.type.ImagePart
+import com.google.firebase.vertexai.type.InlineDataPart
 import com.google.firebase.vertexai.type.InvalidStateException
 import com.google.firebase.vertexai.type.TextPart
 import com.google.firebase.vertexai.type.content
@@ -102,21 +102,21 @@ class Chat(private val model: GenerativeModel, val history: MutableList<Content>
 
     val flow = model.generateContentStream(*history.toTypedArray(), prompt)
     val bitmaps = LinkedList<Bitmap>()
-    val blobs = LinkedList<BlobPart>()
+    val inlineDataParts = LinkedList<InlineDataPart>()
     val text = StringBuilder()
 
     /**
-     * TODO: revisit when images and blobs are returned. This will cause issues with how things are
-     * structured in the response. eg; a text/image/text response will be (incorrectly) represented
-     * as image/text
+     * TODO: revisit when images and inline data are returned. This will cause issues with how
+     * things are structured in the response. eg; a text/image/text response will be (incorrectly)
+     * represented as image/text
      */
     return flow
       .onEach {
         for (part in it.candidates.first().content.parts) {
           when (part) {
             is TextPart -> text.append(part.text)
             is ImagePart -> bitmaps.add(part.image)
-            is BlobPart -> blobs.add(part)
+            is InlineDataPart -> inlineDataParts.add(part)
           }
         }
       }
@@ -128,8 +128,8 @@ class Chat(private val model: GenerativeModel, val history: MutableList<Content>
               for (bitmap in bitmaps) {
                 image(bitmap)
               }
-              for (blob in blobs) {
-                blob(blob.mimeType, blob.blob)
+              for (inlineDataPart in inlineDataParts) {
+                inlineData(inlineDataPart.mimeType, inlineDataPart.inlineData)
               }
               if (text.isNotBlank()) {
                 text(text.toString())

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

@@ -52,12 +52,12 @@ import kotlinx.coroutines.tasks.await
  */
 class GenerativeModel
 internal constructor(
-  val modelName: String,
-  val generationConfig: GenerationConfig? = null,
-  val safetySettings: List<SafetySetting>? = null,
-  val tools: List<Tool>? = null,
-  val toolConfig: ToolConfig? = null,
-  val systemInstruction: Content? = null,
+  private val modelName: String,
+  private val generationConfig: GenerationConfig? = null,
+  private val safetySettings: List<SafetySetting>? = null,
+  private val tools: List<Tool>? = null,
+  private val toolConfig: ToolConfig? = null,
+  private val systemInstruction: Content? = null,
   private val controller: APIController
 ) {
 

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

@@ -68,10 +68,25 @@ internal constructor(@JsonNames("citations") val citationSources: List<CitationS
 
 @Serializable
 internal data class CitationSources(
+  val title: String? = null,
   val startIndex: Int = 0,
   val endIndex: Int,
   val uri: String? = null,
   val license: String? = null,
+  val publicationDate: Date? = null,
+)
+
+@Serializable
+internal data class Date(
+  /** Year of the date. Must be between 1 and 9999, or 0 for no year. */
+  val year: Int? = null,
+  /** 1-based index for month. Must be from 1 to 12, or 0 to specify a year without a month. */
+  val month: Int? = null,
+  /**
+   * Day of a month. Must be from 1 to 31 and valid for the year and month, or 0 to specify a year
+   * by itself or a year and month where the day isn't significant.
+   */
+  val day: Int? = null,
 )
 
 @Serializable
@@ -145,7 +160,7 @@ internal enum class FinishReason {
 internal data class GRpcError(
   val code: Int,
   val message: String,
-  val details: List<GRpcErrorDetails>
+  val details: List<GRpcErrorDetails>? = null
 )
 
 @Serializable internal data class GRpcErrorDetails(val reason: String? = null)

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

@@ -51,7 +51,8 @@ internal data class Content(@EncodeDefault val role: String? = "user", val parts
 
 @Serializable internal data class TextPart(val text: String) : Part
 
-@Serializable internal data class BlobPart(@SerialName("inline_data") val inlineData: Blob) : Part
+@Serializable
+internal data class InlineDataPart(@SerialName("inline_data") val inlineData: InlineData) : Part
 
 @Serializable internal data class FunctionCallPart(val functionCall: FunctionCall) : Part
 
@@ -73,7 +74,7 @@ internal data class FileData(
 )
 
 @Serializable
-internal data class Blob(@SerialName("mime_type") val mimeType: String, val data: Base64)
+internal data class InlineData(@SerialName("mime_type") val mimeType: String, val data: Base64)
 
 @Serializable
 internal data class SafetySetting(
@@ -105,7 +106,7 @@ internal object PartSerializer : JsonContentPolymorphicSerializer<Part>(Part::cl
       "text" in jsonObject -> TextPart.serializer()
       "functionCall" in jsonObject -> FunctionCallPart.serializer()
       "functionResponse" in jsonObject -> FunctionResponsePart.serializer()
-      "inlineData" in jsonObject -> BlobPart.serializer()
+      "inlineData" in jsonObject -> InlineDataPart.serializer()
       "fileData" in jsonObject -> FileDataPart.serializer()
       else -> throw SerializationException("Unknown Part type")
     }

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

@@ -20,13 +20,12 @@ import android.graphics.Bitmap
 import android.graphics.BitmapFactory
 import android.util.Base64
 import com.google.firebase.vertexai.common.client.Schema
-import com.google.firebase.vertexai.common.shared.Blob
 import com.google.firebase.vertexai.common.shared.FileData
 import com.google.firebase.vertexai.common.shared.FunctionCall
 import com.google.firebase.vertexai.common.shared.FunctionCallPart
 import com.google.firebase.vertexai.common.shared.FunctionResponse
 import com.google.firebase.vertexai.common.shared.FunctionResponsePart
-import com.google.firebase.vertexai.type.BlobPart
+import com.google.firebase.vertexai.common.shared.InlineData
 import com.google.firebase.vertexai.type.BlockReason
 import com.google.firebase.vertexai.type.Candidate
 import com.google.firebase.vertexai.type.Citation
@@ -39,11 +38,13 @@ import com.google.firebase.vertexai.type.FunctionCallingConfig
 import com.google.firebase.vertexai.type.FunctionDeclaration
 import com.google.firebase.vertexai.type.GenerateContentResponse
 import com.google.firebase.vertexai.type.GenerationConfig
+import com.google.firebase.vertexai.type.HarmBlockMethod
 import com.google.firebase.vertexai.type.HarmBlockThreshold
 import com.google.firebase.vertexai.type.HarmCategory
 import com.google.firebase.vertexai.type.HarmProbability
 import com.google.firebase.vertexai.type.HarmSeverity
 import com.google.firebase.vertexai.type.ImagePart
+import com.google.firebase.vertexai.type.InlineDataPart
 import com.google.firebase.vertexai.type.Part
 import com.google.firebase.vertexai.type.PromptFeedback
 import com.google.firebase.vertexai.type.SafetyRating
@@ -55,6 +56,7 @@ import com.google.firebase.vertexai.type.ToolConfig
 import com.google.firebase.vertexai.type.UsageMetadata
 import com.google.firebase.vertexai.type.content
 import java.io.ByteArrayOutputStream
+import java.util.Calendar
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.JsonObject
 import org.json.JSONObject
@@ -71,12 +73,12 @@ internal fun Part.toInternal(): com.google.firebase.vertexai.common.shared.Part
   return when (this) {
     is TextPart -> com.google.firebase.vertexai.common.shared.TextPart(text)
     is ImagePart ->
-      com.google.firebase.vertexai.common.shared.BlobPart(
-        Blob("image/jpeg", encodeBitmapToBase64Png(image))
+      com.google.firebase.vertexai.common.shared.InlineDataPart(
+        InlineData("image/jpeg", encodeBitmapToBase64Png(image))
       )
-    is BlobPart ->
-      com.google.firebase.vertexai.common.shared.BlobPart(
-        Blob(mimeType, Base64.encodeToString(blob, BASE_64_FLAGS))
+    is InlineDataPart ->
+      com.google.firebase.vertexai.common.shared.InlineDataPart(
+        InlineData(mimeType, Base64.encodeToString(inlineData, BASE_64_FLAGS))
       )
     is com.google.firebase.vertexai.type.FunctionCallPart ->
       FunctionCallPart(FunctionCall(name, args.orEmpty()))
@@ -96,7 +98,8 @@ internal fun Part.toInternal(): com.google.firebase.vertexai.common.shared.Part
 internal fun SafetySetting.toInternal() =
   com.google.firebase.vertexai.common.shared.SafetySetting(
     harmCategory.toInternal(),
-    threshold.toInternal()
+    threshold.toInternal(),
+    method.toInternal()
   )
 
 internal fun GenerationConfig.toInternal() =
@@ -107,11 +110,13 @@ internal fun GenerationConfig.toInternal() =
     candidateCount = candidateCount,
     maxOutputTokens = maxOutputTokens,
     stopSequences = stopSequences,
+    frequencyPenalty = frequencyPenalty,
+    presencePenalty = presencePenalty,
     responseMimeType = responseMimeType,
     responseSchema = responseSchema?.toInternal()
   )
 
-internal fun com.google.firebase.vertexai.type.HarmCategory.toInternal() =
+internal fun HarmCategory.toInternal() =
   when (this) {
     HarmCategory.HARASSMENT -> com.google.firebase.vertexai.common.shared.HarmCategory.HARASSMENT
     HarmCategory.HATE_SPEECH -> com.google.firebase.vertexai.common.shared.HarmCategory.HATE_SPEECH
@@ -122,6 +127,13 @@ internal fun com.google.firebase.vertexai.type.HarmCategory.toInternal() =
     HarmCategory.UNKNOWN -> com.google.firebase.vertexai.common.shared.HarmCategory.UNKNOWN
   }
 
+internal fun HarmBlockMethod.toInternal() =
+  when (this) {
+    HarmBlockMethod.SEVERITY -> com.google.firebase.vertexai.common.shared.HarmBlockMethod.SEVERITY
+    HarmBlockMethod.PROBABILITY ->
+      com.google.firebase.vertexai.common.shared.HarmBlockMethod.PROBABILITY
+  }
+
 internal fun ToolConfig.toInternal() =
   com.google.firebase.vertexai.common.client.ToolConfig(
     com.google.firebase.vertexai.common.client.FunctionCallingConfig(
@@ -150,7 +162,9 @@ internal fun HarmBlockThreshold.toInternal() =
   }
 
 internal fun Tool.toInternal() =
-  com.google.firebase.vertexai.common.client.Tool(functionDeclarations.map { it.toInternal() })
+  com.google.firebase.vertexai.common.client.Tool(
+    functionDeclarations?.map { it.toInternal() } ?: emptyList()
+  )
 
 internal fun FunctionDeclaration.toInternal() =
   com.google.firebase.vertexai.common.client.FunctionDeclaration(name, "", schema.toInternal())
@@ -191,12 +205,12 @@ internal fun com.google.firebase.vertexai.common.shared.Content.toPublic(): Cont
 internal fun com.google.firebase.vertexai.common.shared.Part.toPublic(): Part {
   return when (this) {
     is com.google.firebase.vertexai.common.shared.TextPart -> TextPart(text)
-    is com.google.firebase.vertexai.common.shared.BlobPart -> {
+    is com.google.firebase.vertexai.common.shared.InlineDataPart -> {
       val data = Base64.decode(inlineData.data, BASE_64_FLAGS)
       if (inlineData.mimeType.contains("image")) {
         ImagePart(decodeBitmapFromImage(data))
       } else {
-        BlobPart(inlineData.mimeType, data)
+        InlineDataPart(inlineData.mimeType, data)
       }
     }
     is FunctionCallPart ->
@@ -218,8 +232,29 @@ internal fun com.google.firebase.vertexai.common.shared.Part.toPublic(): Part {
   }
 }
 
-internal fun com.google.firebase.vertexai.common.server.CitationSources.toPublic() =
-  Citation(startIndex = startIndex, endIndex = endIndex, uri = uri, license = license)
+internal fun com.google.firebase.vertexai.common.server.CitationSources.toPublic(): Citation {
+  val publicationDateAsCalendar =
+    publicationDate?.let {
+      val calendar = Calendar.getInstance()
+      // Internal `Date.year` uses 0 to represent not specified. We use 1 as default.
+      val year = if (it.year == null || it.year < 1) 1 else it.year
+      // Internal `Date.month` uses 0 to represent not specified, or is 1-12 as months. The month as
+      // expected by [Calendar] is 0-based, so we subtract 1 or use 0 as default.
+      val month = if (it.month == null || it.month < 1) 0 else it.month - 1
+      // Internal `Date.day` uses 0 to represent not specified. We use 1 as default.
+      val day = if (it.day == null || it.day < 1) 1 else it.day
+      calendar.set(year, month, day)
+      calendar
+    }
+  return Citation(
+    title = title,
+    startIndex = startIndex,
+    endIndex = endIndex,
+    uri = uri,
+    license = license,
+    publicationDate = publicationDateAsCalendar
+  )
+}
 
 internal fun com.google.firebase.vertexai.common.server.CitationMetadata.toPublic() =
   CitationMetadata(citationSources.map { it.toPublic() })

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

@@ -16,6 +16,8 @@
 
 package com.google.firebase.vertexai.type
 
+import java.util.Calendar
+
 /** A response generated by the model. */
 class Candidate
 internal constructor(
@@ -48,19 +50,23 @@ class CitationMetadata internal constructor(val citations: List<Citation>)
  * Provides citation information for sourcing of content provided by the model between a given
  * [startIndex] and [endIndex].
  *
+ * @property title Title of the attribution.
  * @property startIndex The inclusive beginning of a sequence in a model response that derives from
  * a cited source.
  * @property endIndex The exclusive end of a sequence in a model response that derives from a cited
  * source.
  * @property uri A link to the cited source, if available.
  * @property license The license the cited source work is distributed under, if specified.
+ * @property publicationDate Publication date of the attribution, if available.
  */
 class Citation
 internal constructor(
+  val title: String? = null,
   val startIndex: Int = 0,
   val endIndex: Int,
   val uri: String? = null,
-  val license: String? = null
+  val license: String? = null,
+  val publicationDate: Calendar? = null
 )
 
 /** The reason for content finishing. */

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

@@ -48,9 +48,11 @@ class Content @JvmOverloads constructor(val role: String? = "user", val parts: L
     @JvmName("addText") fun text(text: String) = part(TextPart(text))
 
     /**
-     * Wraps the provided [blob] and [mimeType] inside a [BlobPart] and adds it to the [parts] list.
+     * Wraps the provided [bytes] and [mimeType] inside a [InlineDataPart] and adds it to the
+     * [parts] list.
      */
-    @JvmName("addBlob") fun blob(mimeType: String, blob: ByteArray) = part(BlobPart(mimeType, blob))
+    @JvmName("addInlineData")
+    fun inlineData(mimeType: String, bytes: ByteArray) = part(InlineDataPart(mimeType, bytes))
 
     /** Wraps the provided [image] inside an [ImagePart] and adds it to the [parts] list. */
     @JvmName("addImage") fun image(image: Bitmap) = part(ImagePart(image))

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

@@ -16,6 +16,8 @@
 
 package com.google.firebase.vertexai.type
 
+import com.google.firebase.vertexai.common.client.FunctionCallingConfig
+
 /**
  * Contains configuration for function calling from the model. This can be used to force function
  * calling predictions or disable them.
@@ -25,10 +27,14 @@ package com.google.firebase.vertexai.type
  * should match [FunctionDeclaration.name]. With [Mode.ANY], model will predict a function call from
  * the set of function names provided.
  */
-class FunctionCallingConfig(val mode: Mode, val allowedFunctionNames: List<String>? = null) {
+class FunctionCallingConfig
+internal constructor(
+  internal val mode: Mode,
+  internal val allowedFunctionNames: List<String>? = null
+) {
 
   /** Configuration for dictating when the model should call the attached function. */
-  enum class Mode {
+  internal enum class Mode {
     /**
      * The default behavior for function calling. The model calls functions to answer queries at its
      * discretion
@@ -44,4 +50,24 @@ class FunctionCallingConfig(val mode: Mode, val allowedFunctionNames: List<Strin
      */
     NONE
   }
+
+  companion object {
+    /**
+     * The default behavior for function calling. The model calls functions to answer queries at its
+     * discretion
+     */
+    @JvmStatic fun auto() = FunctionCallingConfig(Mode.AUTO)
+
+    /** The model always predicts a provided function call to answer every query. */
+    @JvmStatic
+    @JvmOverloads
+    fun any(allowedFunctionNames: List<String>? = null) =
+      FunctionCallingConfig(Mode.ANY, allowedFunctionNames)
+
+    /**
+     * The model will never predict a function call to answer a query. This can also be achieved by
+     * not passing any tools to the model.
+     */
+    @JvmStatic fun none() = FunctionCallingConfig(Mode.NONE)
+  }
 }