Add totalBillableCharacters to CountTokensResponse in Vertex AI (#12646)

andrewheard · andrewheard · commit c9536b825428 · 2024-05-13T12:14:00.000-04:00
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 {

Original file line number	Diff line number	Diff line change
`@@ -42,4 +42,7 @@ extension CountTokensRequest: GenerativeAIRequest {`
`42`	`42`	`public struct CountTokensResponse: Decodable {`
`43`	`43`	`/// The total number of tokens in the input given to the model as a prompt.`
`44`	`44`	`public let totalTokens: Int`
	`45`	`+`
	`46`	`+ /// The total number of billable characters in the input given to the model as a prompt.`
	`47`	`+ public let totalBillableCharacters: Int`
`45`	`48`	`}`
Original file line number	Diff line number	Diff line change
`@@ -935,7 +935,9 @@ final class GenerativeModelTests: XCTestCase {`
`935`	`935`	`)`
`936`	`936`
`937`	`937`	`let response = try await model.countTokens("Why is the sky blue?")`
	`938`	`+`
`938`	`939`	`XCTAssertEqual(response.totalTokens, 6)`
	`940`	`+ XCTAssertEqual(response.totalBillableCharacters, 16)`
`939`	`941`	`}`
`940`	`942`
`941`	`943`	`func testCountTokens_modelNotFound() async throws {`