Add forced function calling support in Vertex AI (#12748)

Co-authored-by: Paul Beusterien <[email protected]>

Author: andrewheard

FirebaseVertexAI/Sources/FunctionCalling.swift

@@ -177,6 +177,51 @@ public struct Tool: Encodable {
   }
 }
 
+/// Configuration for specifying function calling behavior.
+public struct FunctionCallingConfig: Encodable {
+  /// Defines the execution behavior for function calling by defining the
+  /// execution mode.
+  public enum Mode: String, Encodable {
+    /// The default behavior for function calling. The model calls functions to answer queries at
+    /// its discretion.
+    case auto = "AUTO"
+
+    /// The model always predicts a provided function call to answer every query.
+    case any = "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.
+    case none = "NONE"
+  }
+
+  /// Specifies the mode in which function calling should execute. If
+  /// unspecified, the default value will be set to AUTO.
+  let mode: Mode?
+
+  /// A set of function names that, when provided, limits the functions the model
+  /// will call.
+  ///
+  /// This should only be set when the Mode is ANY. Function names
+  /// should match [FunctionDeclaration.name]. With mode set to ANY, model will
+  /// predict a function call from the set of function names provided.
+  let allowedFunctionNames: [String]?
+
+  public init(mode: FunctionCallingConfig.Mode? = nil, allowedFunctionNames: [String]? = nil) {
+    self.mode = mode
+    self.allowedFunctionNames = allowedFunctionNames
+  }
+}
+
+/// Tool configuration for any `Tool` specified in the request.
+@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
+public struct ToolConfig: Encodable {
+  let functionCallingConfig: FunctionCallingConfig?
+
+  public init(functionCallingConfig: FunctionCallingConfig? = nil) {
+    self.functionCallingConfig = functionCallingConfig
+  }
+}
+
 /// Result output from a ``FunctionCall``.
 ///
 /// Contains a string representing the `FunctionDeclaration.name` and a structured JSON object

FirebaseVertexAI/Sources/GenerateContentRequest.swift

@@ -22,6 +22,7 @@ struct GenerateContentRequest {
   let generationConfig: GenerationConfig?
   let safetySettings: [SafetySetting]?
   let tools: [Tool]?
+  let toolConfig: ToolConfig?
   let isStreaming: Bool
   let options: RequestOptions
 }
@@ -33,6 +34,7 @@ extension GenerateContentRequest: Encodable {
     case generationConfig
     case safetySettings
     case tools
+    case toolConfig
   }
 }
 

FirebaseVertexAI/Sources/GenerativeModel.swift

@@ -37,6 +37,9 @@ public final class GenerativeModel {
   /// A list of tools the model may use to generate the next response.
   let tools: [Tool]?
 
+  /// Tool configuration for any `Tool` specified in the request.
+  let toolConfig: ToolConfig?
+
   /// Configuration parameters for sending requests to the backend.
   let requestOptions: RequestOptions
 
@@ -49,13 +52,15 @@ public final class GenerativeModel {
   ///   - generationConfig: The content generation parameters your model should use.
   ///   - safetySettings: A value describing what types of harmful content your model should allow.
   ///   - tools: A list of ``Tool`` objects that the model may use to generate the next response.
+  ///   - toolConfig: Tool configuration for any `Tool` specified in the request.
   ///   - requestOptions: Configuration parameters for sending requests to the backend.
   ///   - urlSession: The `URLSession` to use for requests; defaults to `URLSession.shared`.
   init(name: String,
        apiKey: String,
        generationConfig: GenerationConfig? = nil,
        safetySettings: [SafetySetting]? = nil,
        tools: [Tool]?,
+       toolConfig: ToolConfig? = nil,
        requestOptions: RequestOptions,
        appCheck: AppCheckInterop?,
        urlSession: URLSession = .shared) {
@@ -68,6 +73,7 @@ public final class GenerativeModel {
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings
     self.tools = tools
+    self.toolConfig = toolConfig
     self.requestOptions = requestOptions
 
     Logging.default.info("""
@@ -114,6 +120,7 @@ public final class GenerativeModel {
                                                               generationConfig: generationConfig,
                                                               safetySettings: safetySettings,
                                                               tools: tools,
+                                                              toolConfig: toolConfig,
                                                               isStreaming: false,
                                                               options: requestOptions)
       response = try await generativeAIService.loadRequest(request: generateContentRequest)
@@ -186,6 +193,7 @@ public final class GenerativeModel {
                                                         generationConfig: generationConfig,
                                                         safetySettings: safetySettings,
                                                         tools: tools,
+                                                        toolConfig: toolConfig,
                                                         isStreaming: true,
                                                         options: requestOptions)
 

FirebaseVertexAI/Sources/VertexAI.swift

@@ -58,11 +58,13 @@ public class VertexAI: NSObject {
   ///   - generationConfig: The content generation parameters your model should use.
   ///   - safetySettings: A value describing what types of harmful content your model should allow.
   ///   - tools: A list of ``Tool`` objects that the model may use to generate the next response.
+  ///   - toolConfig: Tool configuration for any `Tool` specified in the request.
   ///   - requestOptions: Configuration parameters for sending requests to the backend.
   public func generativeModel(modelName: String,
                               generationConfig: GenerationConfig? = nil,
                               safetySettings: [SafetySetting]? = nil,
                               tools: [Tool]? = nil,
+                              toolConfig: ToolConfig? = nil,
                               requestOptions: RequestOptions = RequestOptions())
     -> GenerativeModel {
     let modelResourceName = modelResourceName(modelName: modelName, location: location)
@@ -77,6 +79,7 @@ public class VertexAI: NSObject {
       generationConfig: generationConfig,
       safetySettings: safetySettings,
       tools: tools,
+      toolConfig: toolConfig,
       requestOptions: requestOptions,
       appCheck: appCheck
     )