Skip to content

Commit 8565a55

Browse files
authored
Add functionCalls accessor to GenerateContentResponse (#123)
1 parent f4a8808 commit 8565a55

File tree

2 files changed

+18
-0
lines changed

2 files changed

+18
-0
lines changed

Sources/GoogleAI/GenerateContentResponse.swift

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,19 @@ public struct GenerateContentResponse {
3737
return text
3838
}
3939

40+
/// Returns function calls found in any `Part`s of the first candidate of the response, if any.
41+
public var functionCalls: [FunctionCall] {
42+
guard let candidate = candidates.first else {
43+
return []
44+
}
45+
return candidate.content.parts.compactMap { part in
46+
guard case let .functionCall(functionCall) = part else {
47+
return nil
48+
}
49+
return functionCall
50+
}
51+
}
52+
4053
/// Initializer for SwiftUI previews or tests.
4154
public init(candidates: [CandidateResponse], promptFeedback: PromptFeedback?) {
4255
self.candidates = candidates

Tests/GoogleAITests/GenerativeModelTests.swift

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ final class GenerativeModelTests: XCTestCase {
6363
let promptFeedback = try XCTUnwrap(response.promptFeedback)
6464
XCTAssertNil(promptFeedback.blockReason)
6565
XCTAssertEqual(promptFeedback.safetyRatings, safetyRatingsNegligible)
66+
XCTAssertEqual(response.functionCalls, [])
6667
}
6768

6869
func testGenerateContent_success_basicReplyShort() async throws {
@@ -86,6 +87,7 @@ final class GenerativeModelTests: XCTestCase {
8687
let promptFeedback = try XCTUnwrap(response.promptFeedback)
8788
XCTAssertNil(promptFeedback.blockReason)
8889
XCTAssertEqual(promptFeedback.safetyRatings, safetyRatingsNegligible)
90+
XCTAssertEqual(response.functionCalls, [])
8991
}
9092

9193
func testGenerateContent_success_citations() async throws {
@@ -188,6 +190,7 @@ final class GenerativeModelTests: XCTestCase {
188190
}
189191
XCTAssertEqual(functionCall.name, "current_time")
190192
XCTAssertTrue(functionCall.args.isEmpty)
193+
XCTAssertEqual(response.functionCalls, [functionCall])
191194
}
192195

193196
func testGenerateContent_success_functionCall_noArguments() async throws {
@@ -209,6 +212,7 @@ final class GenerativeModelTests: XCTestCase {
209212
}
210213
XCTAssertEqual(functionCall.name, "current_time")
211214
XCTAssertTrue(functionCall.args.isEmpty)
215+
XCTAssertEqual(response.functionCalls, [functionCall])
212216
}
213217

214218
func testGenerateContent_success_functionCall_withArguments() async throws {
@@ -234,6 +238,7 @@ final class GenerativeModelTests: XCTestCase {
234238
XCTAssertEqual(argX, .number(4))
235239
let argY = try XCTUnwrap(functionCall.args["y"])
236240
XCTAssertEqual(argY, .number(5))
241+
XCTAssertEqual(response.functionCalls, [functionCall])
237242
}
238243

239244
func testGenerateContent_failure_invalidAPIKey() async throws {

0 commit comments

Comments
 (0)