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

firebase · Nov 7, 2024 · 32a6788 · 32a6788
1 parent 869c7f5
commit 32a6788
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
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 {