[Vertex AI] Add image / function call count tokens integration tests

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.

FirebaseVertexAI/Tests/Integration/IntegrationTests.swift

@@ -74,12 +74,60 @@ 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)
 
     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)
+  }
 }