Add ToolConfig to support constrained decoding

generativeai/src/main/java/com/google/ai/client/generativeai/GenerativeModel.kt

@@ -40,6 +40,7 @@ import com.google.ai.client.generativeai.type.SafetySetting
 import com.google.ai.client.generativeai.type.SerializationException
 import com.google.ai.client.generativeai.type.ThreeParameterFunction
 import com.google.ai.client.generativeai.type.Tool
+import com.google.ai.client.generativeai.type.ToolConfig
 import com.google.ai.client.generativeai.type.TwoParameterFunction
 import com.google.ai.client.generativeai.type.content
 import kotlinx.coroutines.flow.Flow
@@ -65,6 +66,7 @@ internal constructor(
   val generationConfig: GenerationConfig? = null,
   val safetySettings: List<SafetySetting>? = null,
   val tools: List<Tool>? = null,
+  val toolConfig: ToolConfig? = null,
   val requestOptions: RequestOptions = RequestOptions(),
   private val controller: APIController
 ) {
@@ -76,13 +78,15 @@ internal constructor(
     generationConfig: GenerationConfig? = null,
     safetySettings: List<SafetySetting>? = null,
     tools: List<Tool>? = null,
+    toolConfig: ToolConfig? = null,
     requestOptions: RequestOptions = RequestOptions(),
   ) : this(
     modelName,
     apiKey,
     generationConfig,
     safetySettings,
     tools,
+    toolConfig,
     requestOptions,
     APIController(apiKey, modelName, requestOptions)
   )
@@ -222,7 +226,8 @@ internal constructor(
       prompt.map { it.toInternal() },
       safetySettings?.map { it.toInternal() },
       generationConfig?.toInternal(),
-      tools?.map { it.toInternal() }
+      tools?.map { it.toInternal() },
+      toolConfig?.toInternal()
     )
 
   private fun constructCountTokensRequest(vararg prompt: Content) =

generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/Request.kt

@@ -18,6 +18,7 @@ package com.google.ai.client.generativeai.internal.api
 
 import com.google.ai.client.generativeai.internal.api.client.GenerationConfig
 import com.google.ai.client.generativeai.internal.api.client.Tool
+import com.google.ai.client.generativeai.internal.api.client.ToolConfig
 import com.google.ai.client.generativeai.internal.api.shared.Content
 import com.google.ai.client.generativeai.internal.api.shared.SafetySetting
 import kotlinx.serialization.SerialName
@@ -31,7 +32,8 @@ internal data class GenerateContentRequest(
   val contents: List<Content>,
   @SerialName("safety_settings") val safetySettings: List<SafetySetting>? = null,
   @SerialName("generation_config") val generationConfig: GenerationConfig? = null,
-  val tools: List<Tool>? = null
+  val tools: List<Tool>? = null,
+  @SerialName("tool_config") var toolConfig: ToolConfig?
 ) : Request
 
 @Serializable

generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/client/Types.kt

@@ -31,6 +31,22 @@ internal data class GenerationConfig(
 
 @Serializable internal data class Tool(val functionDeclarations: List<FunctionDeclaration>)
 
+@Serializable
+internal data class ToolConfig(
+  @SerialName("function_calling_config") val functionCallingConfig: FunctionCallingConfig
+)
+
+@Serializable
+internal data class FunctionCallingConfig(val mode: Mode) {
+  @Serializable
+  internal enum class Mode {
+    @SerialName("MODE_UNSPECIFIED") UNSPECIFIED,
+    AUTO,
+    ANY,
+    NONE
+  }
+}
+
 @Serializable
 internal data class FunctionDeclaration(
   val name: String,

generativeai/src/main/java/com/google/ai/client/generativeai/internal/util/conversions.kt

@@ -44,12 +44,14 @@ import com.google.ai.client.generativeai.internal.api.shared.SafetySetting
 import com.google.ai.client.generativeai.internal.api.shared.TextPart
 import com.google.ai.client.generativeai.type.BlockThreshold
 import com.google.ai.client.generativeai.type.CitationMetadata
+import com.google.ai.client.generativeai.type.FunctionCallingConfig
 import com.google.ai.client.generativeai.type.FunctionDeclaration
 import com.google.ai.client.generativeai.type.GenerativeBeta
 import com.google.ai.client.generativeai.type.ImagePart
 import com.google.ai.client.generativeai.type.ParameterDeclaration
 import com.google.ai.client.generativeai.type.SerializationException
 import com.google.ai.client.generativeai.type.Tool
+import com.google.ai.client.generativeai.type.ToolConfig
 import com.google.ai.client.generativeai.type.content
 import java.io.ByteArrayOutputStream
 import kotlinx.serialization.decodeFromString
@@ -118,6 +120,21 @@ internal fun Tool.toInternal() =
     functionDeclarations.map { it.toInternal() }
   )
 
+@GenerativeBeta
+internal fun ToolConfig.toInternal() =
+  com.google.ai.client.generativeai.internal.api.client.ToolConfig(
+    com.google.ai.client.generativeai.internal.api.client.FunctionCallingConfig(
+      when (functionCallingConfig.mode) {
+        FunctionCallingConfig.Mode.ANY ->
+          com.google.ai.client.generativeai.internal.api.client.FunctionCallingConfig.Mode.ANY
+        FunctionCallingConfig.Mode.AUTO ->
+          com.google.ai.client.generativeai.internal.api.client.FunctionCallingConfig.Mode.AUTO
+        FunctionCallingConfig.Mode.NONE ->
+          com.google.ai.client.generativeai.internal.api.client.FunctionCallingConfig.Mode.NONE
+      }
+    )
+  )
+
 @GenerativeBeta
 internal fun FunctionDeclaration.toInternal() =
   com.google.ai.client.generativeai.internal.api.client.FunctionDeclaration(

generativeai/src/main/java/com/google/ai/client/generativeai/type/FunctionCallingConfig.kt

@@ -0,0 +1,43 @@
+/*
+ * 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.ai.client.generativeai.type
+
+/**
+ * Contains configuration for function calling from the model. This can be used to force function
+ * calling predictions or disable them.
+ *
+ * @param mode The function calling mode of the model
+ */
+@GenerativeBeta
+class FunctionCallingConfig(val mode: Mode) {
+  enum class Mode {
+    /**
+     * The default behavior for function calling. The model calls functions to answer queries at its
+     * discretion
+     */
+    AUTO,
+
+    /** The model always predicts a provided function call to answer every query. */
+    ANY,
+
+    /**
+     * 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.
+     */
+    NONE
+  }
+}

generativeai/src/main/java/com/google/ai/client/generativeai/type/ToolConfig.kt

@@ -0,0 +1,34 @@
+/*
+ * 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.ai.client.generativeai.type
+
+/**
+ * Contains configuration for the function calling tools of the model. This can be used to change
+ * when the model can predict function calls.
+ *
+ * @param functionCallingConfig The config for function calling
+ */
+@OptIn(GenerativeBeta::class)
+class ToolConfig(val functionCallingConfig: FunctionCallingConfig) {
+
+  companion object {
+    /** Shorthand to construct a ToolConfig that restricts the model from calling any functions */
+    fun never(): ToolConfig = ToolConfig(FunctionCallingConfig(FunctionCallingConfig.Mode.NONE))
+    /** Shorthand to construct a ToolConfig that restricts the model to always call some function */
+    fun always(): ToolConfig = ToolConfig(FunctionCallingConfig(FunctionCallingConfig.Mode.ANY))
+  }
+}