Add FunctionCall decoding (google-gemini#114)

Author: andrewheard

Sources/GoogleAI/Chat.swift

@@ -162,6 +162,10 @@ public class Chat {
           }
 
           parts.append(part)
+
+        case .functionCall:
+          // TODO(andrewheard): Add function call to the chat history when encoding is implemented.
+          fatalError("Function calling not yet implemented in chat.")
         }
       }
     }

Sources/GoogleAI/FunctionCalling.swift

@@ -0,0 +1,41 @@
+// 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.
+
+import Foundation
+
+/// A predicted function call returned from the model.
+public struct FunctionCall: Equatable {
+  /// The name of the function to call.
+  let name: String
+
+  /// The function parameters and values.
+  let args: JSONObject
+}
+
+extension FunctionCall: Decodable {
+  enum CodingKeys: CodingKey {
+    case name
+    case args
+  }
+
+  public init(from decoder: Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+    name = try container.decode(String.self, forKey: .name)
+    if let args = try container.decodeIfPresent(JSONObject.self, forKey: .args) {
+      self.args = args
+    } else {
+      args = JSONObject()
+    }
+  }
+}

Sources/GoogleAI/ModelContent.swift

@@ -25,6 +25,7 @@ public struct ModelContent: Codable, Equatable {
     enum CodingKeys: String, CodingKey {
       case text
       case inlineData
+      case functionCall
     }
 
     enum InlineDataKeys: String, CodingKey {
@@ -38,6 +39,9 @@ public struct ModelContent: Codable, Equatable {
     /// Data with a specified media type. Not all media types may be supported by the AI model.
     case data(mimetype: String, Data)
 
+    /// A predicted function call returned from the model.
+    case functionCall(FunctionCall)
+
     // MARK: Convenience Initializers
 
     /// Convenience function for populating a Part with JPEG data.
@@ -64,6 +68,9 @@ public struct ModelContent: Codable, Equatable {
         )
         try inlineDataContainer.encode(mimetype, forKey: .mimeType)
         try inlineDataContainer.encode(bytes, forKey: .bytes)
+      case .functionCall:
+        // TODO(andrewheard): Encode FunctionCalls when when encoding is implemented.
+        fatalError("FunctionCall encoding not implemented.")
       }
     }
 
@@ -79,10 +86,12 @@ public struct ModelContent: Codable, Equatable {
         let mimetype = try dataContainer.decode(String.self, forKey: .mimeType)
         let bytes = try dataContainer.decode(Data.self, forKey: .bytes)
         self = .data(mimetype: mimetype, bytes)
+      } else if values.contains(.functionCall) {
+        self = try .functionCall(values.decode(FunctionCall.self, forKey: .functionCall))
       } else {
         throw DecodingError.dataCorrupted(.init(
           codingPath: [CodingKeys.text, CodingKeys.inlineData],
-          debugDescription: "Neither text or inline data was found."
+          debugDescription: "No text, inline data or function call was found."
         ))
       }
     }

Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-empty-arguments.json

@@ -0,0 +1,19 @@
+{
+  "candidates": [
+    {
+      "content": {
+        "parts": [
+          {
+            "functionCall": {
+              "name": "current_time"
+            }
+          }
+        ],
+        "role": "model"
+      },
+      "finishReason": "STOP",
+      "index": 0
+    }
+  ]
+}
+

Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-no-arguments.json

@@ -0,0 +1,19 @@
+{
+  "candidates": [
+    {
+      "content": {
+        "parts": [
+          {
+            "functionCall": {
+              "name": "current_time",
+              "args": {}
+            }
+          }
+        ],
+        "role": "model"
+      },
+      "finishReason": "STOP",
+      "index": 0
+    }
+  ]
+}

Tests/GoogleAITests/GenerateContentResponses/unary-success-function-call-with-arguments.json

@@ -0,0 +1,22 @@
+{
+  "candidates": [
+    {
+      "content": {
+        "parts": [
+          {
+            "functionCall": {
+              "name": "sum",
+              "args": {
+                "y": 5,
+                "x": 4
+              }
+            }
+          }
+        ],
+        "role": "model"
+      },
+      "finishReason": "STOP",
+      "index": 0
+    }
+  ]
+}

Tests/GoogleAITests/GenerativeModelTests.swift

@@ -169,6 +169,73 @@ final class GenerativeModelTests: XCTestCase {
     _ = try await model.generateContent(testPrompt)
   }
 
+  func testGenerateContent_success_functionCall_emptyArguments() async throws {
+    MockURLProtocol
+      .requestHandler = try httpRequestHandler(
+        forResource: "unary-success-function-call-empty-arguments",
+        withExtension: "json"
+      )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.content.parts.count, 1)
+    let part = try XCTUnwrap(candidate.content.parts.first)
+    guard case let .functionCall(functionCall) = part else {
+      XCTFail("Part is not a FunctionCall.")
+      return
+    }
+    XCTAssertEqual(functionCall.name, "current_time")
+    XCTAssertTrue(functionCall.args.isEmpty)
+  }
+
+  func testGenerateContent_success_functionCall_noArguments() async throws {
+    MockURLProtocol
+      .requestHandler = try httpRequestHandler(
+        forResource: "unary-success-function-call-no-arguments",
+        withExtension: "json"
+      )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.content.parts.count, 1)
+    let part = try XCTUnwrap(candidate.content.parts.first)
+    guard case let .functionCall(functionCall) = part else {
+      XCTFail("Part is not a FunctionCall.")
+      return
+    }
+    XCTAssertEqual(functionCall.name, "current_time")
+    XCTAssertTrue(functionCall.args.isEmpty)
+  }
+
+  func testGenerateContent_success_functionCall_withArguments() async throws {
+    MockURLProtocol
+      .requestHandler = try httpRequestHandler(
+        forResource: "unary-success-function-call-with-arguments",
+        withExtension: "json"
+      )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.content.parts.count, 1)
+    let part = try XCTUnwrap(candidate.content.parts.first)
+    guard case let .functionCall(functionCall) = part else {
+      XCTFail("Part is not a FunctionCall.")
+      return
+    }
+    XCTAssertEqual(functionCall.name, "sum")
+    XCTAssertEqual(functionCall.args.count, 2)
+    let argX = try XCTUnwrap(functionCall.args["x"])
+    XCTAssertEqual(argX, .number(4))
+    let argY = try XCTUnwrap(functionCall.args["y"])
+    XCTAssertEqual(argY, .number(5))
+  }
+
   func testGenerateContent_failure_invalidAPIKey() async throws {
     let expectedStatusCode = 400
     MockURLProtocol