diff --git a/FirebaseVertexAI/Sources/FunctionCalling.swift b/FirebaseVertexAI/Sources/FunctionCalling.swift index 4641a6f400f..3fb17838d4c 100644 --- a/FirebaseVertexAI/Sources/FunctionCalling.swift +++ b/FirebaseVertexAI/Sources/FunctionCalling.swift @@ -21,6 +21,18 @@ public struct FunctionCall: Equatable, Sendable { /// The function parameters and values. public let args: JSONObject + + /// Constructs a new function call. + /// + /// > Note: A `FunctionCall` is typically received from the model, rather than created manually. + /// + /// - Parameters: + /// - name: The name of the function to call. + /// - args: The function parameters and values. + public init(name: String, args: JSONObject) { + self.name = name + self.args = args + } } /// Structured representation of a function declaration. diff --git a/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift index 0ccbb98e83a..8eddf79a648 100644 --- a/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift +++ b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift @@ -74,7 +74,7 @@ final class IntegrationTests: XCTestCase { // MARK: - Count Tokens - func testCountTokens() async throws { + func testCountTokens_text() async throws { let prompt = "Why is the sky blue?" let response = try await model.countTokens(prompt) @@ -82,4 +82,52 @@ final class IntegrationTests: XCTestCase { XCTAssertEqual(response.totalTokens, 14) XCTAssertEqual(response.totalBillableCharacters, 51) } + + func testCountTokens_image_inlineData() async throws { + guard let image = UIImage(systemName: "cloud") else { + XCTFail("Image not found.") + return + } + + let response = try await model.countTokens(image) + + XCTAssertEqual(response.totalTokens, 266) + XCTAssertEqual(response.totalBillableCharacters, 35) + } + + func testCountTokens_image_fileData() async throws { + let fileData = ModelContent(parts: [.fileData( + mimetype: "image/jpeg", + uri: "gs://ios-opensource-samples.appspot.com/ios/public/blank.jpg" + )]) + + let response = try await model.countTokens([fileData]) + + XCTAssertEqual(response.totalTokens, 266) + XCTAssertEqual(response.totalBillableCharacters, 35) + } + + func testCountTokens_functionCalling() async throws { + let sumDeclaration = FunctionDeclaration( + name: "sum", + description: "Adds two integers.", + parameters: ["x": .integer(), "y": .integer()] + ) + model = vertex.generativeModel( + modelName: "gemini-1.5-flash", + tools: [Tool(functionDeclarations: [sumDeclaration])] + ) + let prompt = "What is 10 + 32?" + let sumCall = FunctionCall(name: "sum", args: ["x": .number(10), "y": .number(32)]) + let sumResponse = FunctionResponse(name: "sum", response: ["result": .number(42)]) + + let response = try await model.countTokens([ + ModelContent(role: "user", parts: [.text(prompt)]), + ModelContent(role: "model", parts: [.functionCall(sumCall)]), + ModelContent(role: "function", parts: [.functionResponse(sumResponse)]), + ]) + + XCTAssertEqual(response.totalTokens, 24) + XCTAssertEqual(response.totalBillableCharacters, 71) + } }