[Vertex AI] Fix countTokens integration tests (#14673)

Author: andrewheard

FirebaseVertexAI/Tests/TestApp/Tests/Integration/CountTokensIntegrationTests.swift

@@ -121,4 +121,46 @@ struct CountTokensIntegrationTests {
       }
     )
   }
+
+  @Test(arguments: [
+    InstanceConfig.vertexV1,
+    InstanceConfig.vertexV1Staging,
+    InstanceConfig.vertexV1Beta,
+    InstanceConfig.vertexV1BetaStaging,
+    /* System instructions are not supported on the v1 Developer API. */
+    InstanceConfig.developerV1Beta,
+  ])
+  func countTokens_jsonSchema(_ config: InstanceConfig) async throws {
+    let model = VertexAI.componentInstance(config).generativeModel(
+      modelName: ModelNames.gemini2Flash,
+      generationConfig: GenerationConfig(
+        responseMIMEType: "application/json",
+        responseSchema: Schema.object(properties: [
+          "startDate": .string(format: .custom("date-time")),
+          "yearsSince": .integer(format: .custom("int32")),
+          "hoursSince": .integer(format: .int32),
+          "minutesSince": .integer(format: .int64),
+        ])
+      ),
+      safetySettings: safetySettings
+    )
+    let prompt = "It is 2050-01-01, how many years, hours and minutes since 2000-01-01?"
+
+    let response = try await model.countTokens(prompt)
+
+    switch config.apiConfig.service {
+    case .vertexAI:
+      #expect(response.totalTokens == 65)
+      #expect(response.totalBillableCharacters == 165)
+    case .developer:
+      // The Developer API erroneously ignores the `responseSchema` when counting tokens, resulting
+      // in a lower total count than Vertex AI.
+      #expect(response.totalTokens == 34)
+      #expect(response.totalBillableCharacters == nil)
+    }
+    #expect(response.promptTokensDetails.count == 1)
+    let promptTokensDetails = try #require(response.promptTokensDetails.first)
+    #expect(promptTokensDetails.modality == .text)
+    #expect(promptTokensDetails.tokenCount == response.totalTokens)
+  }
 }

FirebaseVertexAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift

@@ -105,9 +105,10 @@ struct GenerateContentIntegrationTests {
     #expect(text == "Blue")
 
     let usageMetadata = try #require(response.usageMetadata)
-    #expect(usageMetadata.promptTokenCount == 14)
+    #expect(usageMetadata.promptTokenCount.isEqual(to: 15, accuracy: tokenCountAccuracy))
     #expect(usageMetadata.candidatesTokenCount.isEqual(to: 1, accuracy: tokenCountAccuracy))
-    #expect(usageMetadata.totalTokenCount.isEqual(to: 15, accuracy: tokenCountAccuracy))
+    #expect(usageMetadata.totalTokenCount
+      == usageMetadata.promptTokenCount + usageMetadata.candidatesTokenCount)
     #expect(usageMetadata.promptTokensDetails.count == 1)
     let promptTokensDetails = try #require(usageMetadata.promptTokensDetails.first)
     #expect(promptTokensDetails.modality == .text)

FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift

@@ -203,31 +203,6 @@ final class IntegrationTests: XCTestCase {
     XCTAssertEqual(promptTokensDetails.tokenCount, 24)
   }
 
-  func testCountTokens_jsonSchema() async throws {
-    model = vertex.generativeModel(
-      modelName: "gemini-2.0-flash",
-      generationConfig: GenerationConfig(
-        responseMIMEType: "application/json",
-        responseSchema: Schema.object(properties: [
-          "startDate": .string(format: .custom("date")),
-          "yearsSince": .integer(format: .custom("int16")),
-          "hoursSince": .integer(format: .int32),
-          "minutesSince": .integer(format: .int64),
-        ])
-      )
-    )
-    let prompt = "It is 2050-01-01, how many years, hours and minutes since 2000-01-01?"
-
-    let response = try await model.countTokens(prompt)
-
-    XCTAssertEqual(response.totalTokens, 58)
-    XCTAssertEqual(response.totalBillableCharacters, 160)
-    XCTAssertEqual(response.promptTokensDetails.count, 1)
-    let promptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first)
-    XCTAssertEqual(promptTokensDetails.modality, .text)
-    XCTAssertEqual(promptTokensDetails.tokenCount, 58)
-  }
-
   func testCountTokens_appCheckNotConfigured_shouldFail() async throws {
     let app = try XCTUnwrap(FirebaseApp.app(name: FirebaseAppNames.appCheckNotConfigured))
     let vertex = VertexAI.vertexAI(app: app)

FirebaseVertexAI/Tests/TestApp/Tests/Utilities/IntegrationTestUtils.swift

@@ -40,6 +40,6 @@ enum IntegrationTestUtils {
 
 extension Numeric where Self: Strideable, Self.Stride.Magnitude: Comparable {
   func isEqual(to other: Self, accuracy: Self.Stride) -> Bool {
-    return distance(to: other).magnitude < accuracy.magnitude
+    return distance(to: other).magnitude <= accuracy.magnitude
   }
 }