Simplify top-level schema declaration in FunctionDeclaration (#117)

google-gemini · Mar 14, 2024 · ae60f57 · ae60f57
1 parent 44edbb1
commit ae60f57
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 24 deletions.
diff --git a/Examples/GenerativeAICLI/Sources/GenerateContent.swift b/Examples/GenerativeAICLI/Sources/GenerateContent.swift
@@ -81,24 +81,21 @@ struct GenerateContent: AsyncParsableCommand {
           FunctionDeclaration(
             name: "get_exchange_rate",
             description: "Get the exchange rate for currencies between countries",
-            parameters: Schema(
-              type: .object,
-              properties: [
-                "currency_from": Schema(
-                  type: .string,
-                  format: "enum",
-                  description: "The currency to convert from in ISO 4217 format",
-                  enumValues: ["USD", "EUR", "JPY", "GBP", "AUD", "CAD"]
-                ),
-                "currency_to": Schema(
-                  type: .string,
-                  format: "enum",
-                  description: "The currency to convert to in ISO 4217 format",
-                  enumValues: ["USD", "EUR", "JPY", "GBP", "AUD", "CAD"]
-                ),
-              ],
-              required: ["currency_from", "currency_to"]
-            )
+            parameters: [
+              "currency_from": Schema(
+                type: .string,
+                format: "enum",
+                description: "The currency to convert from in ISO 4217 format",
+                enumValues: ["USD", "EUR", "JPY", "GBP", "AUD", "CAD"]
+              ),
+              "currency_to": Schema(
+                type: .string,
+                format: "enum",
+                description: "The currency to convert to in ISO 4217 format",
+                enumValues: ["USD", "EUR", "JPY", "GBP", "AUD", "CAD"]
+              ),
+            ],
+            requiredParameters: ["currency_from", "currency_to"]
           ),
         ])],
         requestOptions: RequestOptions(apiVersion: "v1beta")

diff --git a/Sources/GoogleAI/FunctionCalling.swift b/Sources/GoogleAI/FunctionCalling.swift
@@ -41,7 +41,7 @@ public class Schema: Encodable {
 
   let properties: [String: Schema]?
 
-  let required: [String]?
+  let requiredProperties: [String]?
 
   enum CodingKeys: String, CodingKey {
     case type
@@ -51,22 +51,22 @@ public class Schema: Encodable {
     case enumValues = "enum"
     case items
     case properties
-    case required
+    case requiredProperties = "required"
   }
 
   public init(type: DataType, format: String? = nil, description: String? = nil,
               nullable: Bool? = nil,
               enumValues: [String]? = nil, items: Schema? = nil,
               properties: [String: Schema]? = nil,
-              required: [String]? = nil) {
+              requiredProperties: [String]? = nil) {
     self.type = type
     self.format = format
     self.description = description
     self.nullable = nullable
     self.enumValues = enumValues
     self.items = items
     self.properties = properties
-    self.required = required
+    self.requiredProperties = requiredProperties
   }
 }
 
@@ -88,10 +88,15 @@ public struct FunctionDeclaration {
 
   let parameters: Schema?
 
-  public init(name: String, description: String, parameters: Schema?) {
+  public init(name: String, description: String, parameters: [String: Schema]?,
+              requiredParameters: [String]?) {
     self.name = name
     self.description = description
-    self.parameters = parameters
+    self.parameters = Schema(
+      type: .object,
+      properties: parameters,
+      requiredProperties: requiredParameters
+    )
   }
 }