diff --git a/FirebaseVertexAI/Sources/CountTokensRequest.swift b/FirebaseVertexAI/Sources/CountTokensRequest.swift index 0a58d40acc0..d294212658c 100644 --- 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 index 0bccd9e7c3d..d2ad6e4ff30 100644 --- 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 index 9a1210c123c..123299ac1a7 100644 --- 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 {