[Vertex AI] Fix testCountTokens_jsonSchema integration test (#14049)

andrewheard · web-flow · commit 32a6788e806b · 2024-11-07T17:30:54.000-05:00
diff --git a/FirebaseVertexAI/CHANGELOG.md b/FirebaseVertexAI/CHANGELOG.md
@@ -1,3 +1,8 @@
+# Unreleased
+- [changed] The token counts from `GenerativeModel.countTokens(...)` now include
+  tokens from the schema for JSON output and function calling; reported token
+  counts will now be higher if using these features.
+
 # 11.5.0
 - [fixed] Fixed an issue where `VertexAI.vertexAI(app: app1)` and
   `VertexAI.vertexAI(app: app2)` would return the same instance if their
diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift
@@ -217,8 +217,8 @@ final class IntegrationTests: XCTestCase {
 
     let response = try await model.countTokens(prompt)
 
-    XCTAssertEqual(response.totalTokens, 34)
-    XCTAssertEqual(response.totalBillableCharacters, 59)
+    XCTAssertEqual(response.totalTokens, 58)
+    XCTAssertEqual(response.totalBillableCharacters, 160)
   }
 
   func testCountTokens_appCheckNotConfigured_shouldFail() async throws {

Original file line number	Diff line number	Diff line change
`@@ -217,8 +217,8 @@ final class IntegrationTests: XCTestCase {`
`217`	`217`
`218`	`218`	`let response = try await model.countTokens(prompt)`
`219`	`219`
`220`		`- XCTAssertEqual(response.totalTokens, 34)`
`221`		`- XCTAssertEqual(response.totalBillableCharacters, 59)`
	`220`	`+ XCTAssertEqual(response.totalTokens, 58)`
	`221`	`+ XCTAssertEqual(response.totalBillableCharacters, 160)`
`222`	`222`	`}`
`223`	`223`
`224`	`224`	`func testCountTokens_appCheckNotConfigured_shouldFail() async throws {`