Add functionCalls accessor to GenerateContentResponse (google-gemini#123

)
SsomsakTH · Jun 21, 2024 · ab4f33c · ab4f33c
1 parent 62a13f1
commit ab4f33c
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/Sources/GoogleAI/GenerateContentResponse.swift b/Sources/GoogleAI/GenerateContentResponse.swift
@@ -37,6 +37,19 @@ public struct GenerateContentResponse {
     return text
   }
 
+  /// Returns function calls found in any `Part`s of the first candidate of the response, if any.
+  public var functionCalls: [FunctionCall] {
+    guard let candidate = candidates.first else {
+      return []
+    }
+    return candidate.content.parts.compactMap { part in
+      guard case let .functionCall(functionCall) = part else {
+        return nil
+      }
+      return functionCall
+    }
+  }
+
   /// Initializer for SwiftUI previews or tests.
   public init(candidates: [CandidateResponse], promptFeedback: PromptFeedback?) {
     self.candidates = candidates

diff --git a/Tests/GoogleAITests/GenerativeModelTests.swift b/Tests/GoogleAITests/GenerativeModelTests.swift
@@ -63,6 +63,7 @@ final class GenerativeModelTests: XCTestCase {
     let promptFeedback = try XCTUnwrap(response.promptFeedback)
     XCTAssertNil(promptFeedback.blockReason)
     XCTAssertEqual(promptFeedback.safetyRatings, safetyRatingsNegligible)
+    XCTAssertEqual(response.functionCalls, [])
   }
 
   func testGenerateContent_success_basicReplyShort() async throws {
@@ -86,6 +87,7 @@ final class GenerativeModelTests: XCTestCase {
     let promptFeedback = try XCTUnwrap(response.promptFeedback)
     XCTAssertNil(promptFeedback.blockReason)
     XCTAssertEqual(promptFeedback.safetyRatings, safetyRatingsNegligible)
+    XCTAssertEqual(response.functionCalls, [])
   }
 
   func testGenerateContent_success_citations() async throws {
@@ -188,6 +190,7 @@ final class GenerativeModelTests: XCTestCase {
     }
     XCTAssertEqual(functionCall.name, "current_time")
     XCTAssertTrue(functionCall.args.isEmpty)
+    XCTAssertEqual(response.functionCalls, [functionCall])
   }
 
   func testGenerateContent_success_functionCall_noArguments() async throws {
@@ -209,6 +212,7 @@ final class GenerativeModelTests: XCTestCase {
     }
     XCTAssertEqual(functionCall.name, "current_time")
     XCTAssertTrue(functionCall.args.isEmpty)
+    XCTAssertEqual(response.functionCalls, [functionCall])
   }
 
   func testGenerateContent_success_functionCall_withArguments() async throws {
@@ -234,6 +238,7 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(argX, .number(4))
     let argY = try XCTUnwrap(functionCall.args["y"])
     XCTAssertEqual(argY, .number(5))
+    XCTAssertEqual(response.functionCalls, [functionCall])
   }
 
   func testGenerateContent_failure_invalidAPIKey() async throws {