File tree Expand file tree Collapse file tree 5 files changed +25
-9
lines changed
main/kotlin/com/google/firebase/vertexai
test/java/com/google/firebase/vertexai/common Expand file tree Collapse file tree 5 files changed +25
-9
lines changed Original file line number Diff line number Diff line change @@ -47,7 +47,10 @@ internal data class ToolConfig(
47
47
)
48
48
49
49
@Serializable
50
- internal data class FunctionCallingConfig (val mode : Mode ) {
50
+ internal data class FunctionCallingConfig (
51
+ val mode : Mode ,
52
+ @SerialName(" allowed_function_names" ) val allowedFunctionNames : List <String >? = null
53
+ ) {
51
54
@Serializable
52
55
enum class Mode {
53
56
@SerialName(" MODE_UNSPECIFIED" ) UNSPECIFIED ,
Original file line number Diff line number Diff line change @@ -137,7 +137,8 @@ internal fun ToolConfig.toInternal() =
137
137
com.google.firebase.vertexai.common.client.FunctionCallingConfig .Mode .AUTO
138
138
FunctionCallingConfig .Mode .NONE ->
139
139
com.google.firebase.vertexai.common.client.FunctionCallingConfig .Mode .NONE
140
- }
140
+ },
141
+ functionCallingConfig.allowedFunctionNames
141
142
)
142
143
)
143
144
Original file line number Diff line number Diff line change @@ -21,8 +21,11 @@ package com.google.firebase.vertexai.type
21
21
* calling predictions or disable them.
22
22
*
23
23
* @param mode The function calling mode of the model
24
+ * @param allowedFunctionNames Function names to call. Only set when the [Mode.ANY]. Function names
25
+ * should match [FunctionDeclaration.name]. With [Mode.ANY], model will predict a function call from
26
+ * the set of function names provided.
24
27
*/
25
- class FunctionCallingConfig (val mode : Mode ) {
28
+ class FunctionCallingConfig (val mode : Mode , val allowedFunctionNames : List < String > ? = null ) {
26
29
27
30
/* * Configuration for dictating when the model should call the attached function. */
28
31
enum class Mode {
Original file line number Diff line number Diff line change @@ -27,7 +27,12 @@ class ToolConfig(val functionCallingConfig: FunctionCallingConfig) {
27
27
companion object {
28
28
/* * Shorthand to construct a ToolConfig that restricts the model from calling any functions */
29
29
fun never (): ToolConfig = ToolConfig (FunctionCallingConfig (FunctionCallingConfig .Mode .NONE ))
30
- /* * Shorthand to construct a ToolConfig that restricts the model to always call some function */
31
- fun always (): ToolConfig = ToolConfig (FunctionCallingConfig (FunctionCallingConfig .Mode .ANY ))
30
+ /* *
31
+ * Shorthand to construct a ToolConfig that restricts the model to always call some function.
32
+ * You can optionally [allowedFunctionNames] to restrict the model to only call these functions.
33
+ * See [FunctionCallingConfig] for more information.
34
+ */
35
+ fun always (allowedFunctionNames : List <String >? = null): ToolConfig =
36
+ ToolConfig (FunctionCallingConfig (FunctionCallingConfig .Mode .ANY , allowedFunctionNames))
32
37
}
33
38
}
Original file line number Diff line number Diff line change @@ -185,17 +185,21 @@ internal class RequestFormatTests {
185
185
contents = listOf (Content (parts = listOf (TextPart (" Arbitrary" )))),
186
186
toolConfig =
187
187
ToolConfig (
188
- functionCallingConfig =
189
- FunctionCallingConfig (mode = FunctionCallingConfig .Mode .AUTO )
190
- ),
191
- )
188
+ FunctionCallingConfig (
189
+ mode = FunctionCallingConfig .Mode .ANY ,
190
+ allowedFunctionNames = listOf (" allowedFunctionName" )
191
+ )
192
+ )
193
+ ),
192
194
)
193
195
.collect { channel.close() }
194
196
}
195
197
196
198
val requestBodyAsText = (mockEngine.requestHistory.first().body as TextContent ).text
197
199
198
200
requestBodyAsText shouldContainJsonKey " tool_config.function_calling_config.mode"
201
+ requestBodyAsText shouldContainJsonKey
202
+ " tool_config.function_calling_config.allowed_function_names"
199
203
}
200
204
201
205
@Test
You can’t perform that action at this time.
0 commit comments