wire through response schema from common (#6114)

Co-authored-by: David Motsonashvili <davidmotson@google.com>
firebase · Jul 23, 2024 · 9a1d120 · 9a1d120
1 parent 1c03d4b
commit 9a1d120
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/firebase-vertexai/CHANGELOG.md b/firebase-vertexai/CHANGELOG.md
@@ -1,5 +1,5 @@
 # Unreleased
-
+* [feature] added support for `responseSchema` in `GenerationConfig`.
 
 # 16.0.0-beta03
 * [changed] Breaking Change: changed `Schema.int` to return 32 bit integers instead of 64 bit (long).

diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt
@@ -114,7 +114,8 @@ internal fun GenerationConfig.toInternal() =
     candidateCount = candidateCount,
     maxOutputTokens = maxOutputTokens,
     stopSequences = stopSequences,
-    responseMimeType = responseMimeType
+    responseMimeType = responseMimeType,
+    responseSchema = responseSchema?.toInternal()
   )
 
 internal fun com.google.firebase.vertexai.type.HarmCategory.toInternal() =

diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/GenerationConfig.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/GenerationConfig.kt
@@ -25,9 +25,11 @@ package com.google.firebase.vertexai.type
  * @property candidateCount The max *unique* responses to return
  * @property maxOutputTokens The max tokens to generate per response
  * @property stopSequences A list of strings to stop generation on occurrence of
- * * @property responseMimeType Response type for generated candidate text. See the
+ * @property responseMimeType Response type for generated candidate text. See the
  * [vertex docs](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/GenerationConfig)
  * for a list of supported types.
+ * @property responseSchema A schema that the response must adhere to, used with the
+ * `application/json` mimeType.
  */
 class GenerationConfig
 private constructor(
@@ -37,7 +39,8 @@ private constructor(
   val candidateCount: Int?,
   val maxOutputTokens: Int?,
   val stopSequences: List<String>?,
-  val responseMimeType: String?
+  val responseMimeType: String?,
+  val responseSchema: Schema<*>? = null,
 ) {
 
   /**
@@ -63,6 +66,7 @@ private constructor(
     @JvmField var maxOutputTokens: Int? = null
     @JvmField var stopSequences: List<String>? = null
     @JvmField var responseMimeType: String? = null
+    @JvmField var responseSchema: Schema<*>? = null
 
     /** Create a new [GenerationConfig] with the attached arguments. */
     fun build() =
@@ -73,7 +77,8 @@ private constructor(
         candidateCount = candidateCount,
         maxOutputTokens = maxOutputTokens,
         stopSequences = stopSequences,
-        responseMimeType = responseMimeType
+        responseMimeType = responseMimeType,
+        responseSchema = responseSchema
       )
   }