From f84cfac2f8ca4222efd7157e5bd67fc4412b9306 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 4 Oct 2024 17:35:57 -0400
Subject: [PATCH 1/2] [Vertex AI] Add image / function call count tokens
 integration tests

---
 .../Tests/Integration/IntegrationTests.swift  | 53 ++++++++++++++++++-
 1 file changed, 51 insertions(+), 2 deletions(-)

diff --git a/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift
index 0ccbb98e83a..00e54cbbc9e 100644
--- a/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift
+++ b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift
@@ -13,9 +13,10 @@
 // limitations under the License.
 
 import FirebaseCore
-import FirebaseVertexAI
 import XCTest
 
+@testable import FirebaseVertexAI
+
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 final class IntegrationTests: XCTestCase {
   // Set temperature, topP and topK to lowest allowed values to make responses more deterministic.
@@ -74,7 +75,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 +83,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)
+  }
 }

From 7e60b2eba042530666e7eaa39ec2dd596b0cb17f Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 4 Oct 2024 17:56:54 -0400
Subject: [PATCH 2/2] Make `FunctionCall` constructor public

---
 FirebaseVertexAI/Sources/FunctionCalling.swift       | 12 ++++++++++++
 .../Tests/Integration/IntegrationTests.swift         |  3 +--
 2 files changed, 13 insertions(+), 2 deletions(-)

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 00e54cbbc9e..8eddf79a648 100644
--- a/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift
+++ b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift
@@ -13,10 +13,9 @@
 // limitations under the License.
 
 import FirebaseCore
+import FirebaseVertexAI
 import XCTest
 
-@testable import FirebaseVertexAI
-
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 final class IntegrationTests: XCTestCase {
   // Set temperature, topP and topK to lowest allowed values to make responses more deterministic.