Add totalBillableCharacters to CountTokensResponse in Vertex AI (#…

…12646)
firebase · Mar 27, 2024 · 0711dc1 · 0711dc1
1 parent 3e32efe
commit 0711dc1
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/FirebaseVertexAI/Sources/CountTokensRequest.swift b/FirebaseVertexAI/Sources/CountTokensRequest.swift
@@ -42,4 +42,7 @@ extension CountTokensRequest: GenerativeAIRequest {
 public struct CountTokensResponse: Decodable {
   /// The total number of tokens in the input given to the model as a prompt.
   public let totalTokens: Int
+
+  /// The total number of billable characters in the input given to the model as a prompt.
+  public let totalBillableCharacters: Int
 }
diff --git a/FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-total-tokens.json b/FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-total-tokens.json
@@ -1,3 +1,4 @@
 {
-  "totalTokens": 6
+  "totalTokens": 6,
+  "totalBillableCharacters": 16
 }
diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift
@@ -935,7 +935,9 @@ final class GenerativeModelTests: XCTestCase {
     )
 
     let response = try await model.countTokens("Why is the sky blue?")
+
     XCTAssertEqual(response.totalTokens, 6)
+    XCTAssertEqual(response.totalBillableCharacters, 16)
   }
 
   func testCountTokens_modelNotFound() async throws {