Skip to content

Commit 1e954e9

Browse files
authored
[Vertex AI] Refactored FunctionCallingConfig.Mode as a struct (#13864)
1 parent 4b263b6 commit 1e954e9

File tree

2 files changed

+43
-20
lines changed

2 files changed

+43
-20
lines changed

FirebaseVertexAI/Sources/FunctionCalling.swift

Lines changed: 39 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -73,33 +73,53 @@ public struct Tool {
7373

7474
/// Configuration for specifying function calling behavior.
7575
public struct FunctionCallingConfig {
76-
/// Defines the execution behavior for function calling by defining the
77-
/// execution mode.
78-
public enum Mode: String {
79-
/// The default behavior for function calling. The model calls functions to answer queries at
80-
/// its discretion.
81-
case auto = "AUTO"
76+
/// Defines the execution behavior for function calling by defining the execution mode.
77+
public struct Mode: EncodableProtoEnum {
78+
enum Kind: String {
79+
case auto = "AUTO"
80+
case any = "ANY"
81+
case none = "NONE"
82+
}
83+
84+
/// The default behavior for function calling.
85+
///
86+
/// The model calls functions to answer queries at its discretion.
87+
public static var auto: Mode {
88+
return self.init(kind: .auto)
89+
}
8290

8391
/// The model always predicts a provided function call to answer every query.
84-
case any = "ANY"
85-
86-
/// The model will never predict a function call to answer a query. This can also be achieved by
87-
/// not passing any tools to the model.
88-
case none = "NONE"
92+
public static var any: Mode {
93+
return self.init(kind: .any)
94+
}
95+
96+
/// The model will never predict a function call to answer a query.
97+
///
98+
/// > Note: This can also be achieved by not passing any ``FunctionDeclaration`` tools
99+
/// > when instantiating the model.
100+
public static var none: Mode {
101+
return self.init(kind: .none)
102+
}
103+
104+
let rawValue: String
89105
}
90106

91-
/// Specifies the mode in which function calling should execute. If
92-
/// unspecified, the default value will be set to AUTO.
107+
/// Specifies the mode in which function calling should execute.
93108
let mode: Mode?
94109

95-
/// A set of function names that, when provided, limits the functions the model
96-
/// will call.
97-
///
98-
/// This should only be set when the Mode is ANY. Function names
99-
/// should match [FunctionDeclaration.name]. With mode set to ANY, model will
100-
/// predict a function call from the set of function names provided.
110+
/// A set of function names that, when provided, limits the functions the model will call.
101111
let allowedFunctionNames: [String]?
102112

113+
/// Creates a new `FunctionCallingConfig`.
114+
///
115+
/// - Parameters:
116+
/// - mode: Specifies the mode in which function calling should execute; if unspecified, the
117+
/// default behavior will be ``Mode/auto``.
118+
/// - allowedFunctionNames: A set of function names that, when provided, limits the functions
119+
/// the model will call.
120+
/// Note: This should only be set when the ``Mode`` is ``Mode/any``. Function names should match
121+
/// `[FunctionDeclaration.name]`. With mode set to ``Mode/any``, the model will predict a
122+
/// function call from the set of function names provided.
103123
public init(mode: FunctionCallingConfig.Mode? = nil, allowedFunctionNames: [String]? = nil) {
104124
self.mode = mode
105125
self.allowedFunctionNames = allowedFunctionNames

FirebaseVertexAI/Tests/Integration/IntegrationTests.swift

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ final class IntegrationTests: XCTestCase {
5959
generationConfig: generationConfig,
6060
safetySettings: safetySettings,
6161
tools: [],
62+
toolConfig: .init(functionCallingConfig: .init(mode: FunctionCallingConfig.Mode.none)),
6263
systemInstruction: systemInstruction
6364
)
6465
}
@@ -94,6 +95,7 @@ final class IntegrationTests: XCTestCase {
9495
SafetySetting(harmCategory: .dangerousContent, threshold: .blockNone),
9596
SafetySetting(harmCategory: .civicIntegrity, threshold: .off),
9697
],
98+
toolConfig: .init(functionCallingConfig: .init(mode: .auto)),
9799
systemInstruction: systemInstruction
98100
)
99101

@@ -135,7 +137,8 @@ final class IntegrationTests: XCTestCase {
135137
)
136138
model = vertex.generativeModel(
137139
modelName: "gemini-1.5-flash",
138-
tools: [Tool(functionDeclarations: [sumDeclaration])]
140+
tools: [Tool(functionDeclarations: [sumDeclaration])],
141+
toolConfig: .init(functionCallingConfig: .init(mode: .any, allowedFunctionNames: ["sum"]))
139142
)
140143
let prompt = "What is 10 + 32?"
141144
let sumCall = FunctionCallPart(name: "sum", args: ["x": .number(10), "y": .number(32)])

0 commit comments

Comments
 (0)