From 6070645135308d638c8afe8dc8f505e92f73a1f0 Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Fri, 5 Apr 2024 16:26:01 -0400 Subject: [PATCH 1/3] Make `totalBillableCharacters` optional in Vertex AI --- FirebaseVertexAI/Sources/CountTokensRequest.swift | 2 +- .../success-no-billable-characters.json | 3 +++ .../Tests/Unit/GenerativeModelTests.swift | 15 +++++++++++++++ 3 files changed, 19 insertions(+), 1 deletion(-) create mode 100644 FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-no-billable-characters.json diff --git a/FirebaseVertexAI/Sources/CountTokensRequest.swift b/FirebaseVertexAI/Sources/CountTokensRequest.swift index d294212658c..ca8d6f1743c 100644 --- a/FirebaseVertexAI/Sources/CountTokensRequest.swift +++ b/FirebaseVertexAI/Sources/CountTokensRequest.swift @@ -44,5 +44,5 @@ public struct CountTokensResponse: Decodable { public let totalTokens: Int /// The total number of billable characters in the input given to the model as a prompt. - public let totalBillableCharacters: Int + public let totalBillableCharacters: Int? } diff --git a/FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-no-billable-characters.json b/FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-no-billable-characters.json new file mode 100644 index 00000000000..03425d51db0 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-no-billable-characters.json @@ -0,0 +1,3 @@ +{ + "totalTokens": 258 +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift index 13cbb6df94e..4bda29ddb66 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -946,6 +946,21 @@ final class GenerativeModelTests: XCTestCase { XCTAssertEqual(response.totalBillableCharacters, 16) } + func testCountTokens_succeeds_noBillableCharacters() async throws { + MockURLProtocol.requestHandler = try httpRequestHandler( + forResource: "success-no-billable-characters", + withExtension: "json" + ) + + let response = try await model.countTokens(ModelContent.Part.data( + mimetype: "image/jpeg", + Data() + )) + + XCTAssertEqual(response.totalTokens, 258) + XCTAssertNil(response.totalBillableCharacters) + } + func testCountTokens_modelNotFound() async throws { MockURLProtocol.requestHandler = try httpRequestHandler( forResource: "failure-model-not-found", withExtension: "json", From 43eddac7c979e7eca907ab35a18225426dd67653 Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Fri, 5 Apr 2024 16:56:20 -0400 Subject: [PATCH 2/3] Make `totalBillableCharacters` 0 if omitted --- .../Sources/CountTokensRequest.swift | 21 +++++++++++++++++-- .../Tests/Unit/GenerativeModelTests.swift | 2 +- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/FirebaseVertexAI/Sources/CountTokensRequest.swift b/FirebaseVertexAI/Sources/CountTokensRequest.swift index ca8d6f1743c..feb20aedc3e 100644 --- a/FirebaseVertexAI/Sources/CountTokensRequest.swift +++ b/FirebaseVertexAI/Sources/CountTokensRequest.swift @@ -39,10 +39,27 @@ extension CountTokensRequest: GenerativeAIRequest { /// The model's response to a count tokens request. @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *) -public struct CountTokensResponse: Decodable { +public struct CountTokensResponse { /// 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? + public let totalBillableCharacters: Int +} + +@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *) +extension CountTokensResponse: Decodable { + enum CodingKeys: CodingKey { + case totalTokens + case totalBillableCharacters + } + + public init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + totalTokens = try container.decode(Int.self, forKey: .totalTokens) + totalBillableCharacters = try container.decodeIfPresent( + Int.self, + forKey: .totalBillableCharacters + ) ?? 0 + } } diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift index 4bda29ddb66..92585266891 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -958,7 +958,7 @@ final class GenerativeModelTests: XCTestCase { )) XCTAssertEqual(response.totalTokens, 258) - XCTAssertNil(response.totalBillableCharacters) + XCTAssertEqual(response.totalBillableCharacters, 0) } func testCountTokens_modelNotFound() async throws { From 6f89d08132f79d875af7a7c6448c86fe490e62ca Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Wed, 10 Apr 2024 11:54:57 -0400 Subject: [PATCH 3/3] Add details to `totalBillableCharacters` docs --- FirebaseVertexAI/Sources/CountTokensRequest.swift | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/FirebaseVertexAI/Sources/CountTokensRequest.swift b/FirebaseVertexAI/Sources/CountTokensRequest.swift index feb20aedc3e..5072412cf9b 100644 --- a/FirebaseVertexAI/Sources/CountTokensRequest.swift +++ b/FirebaseVertexAI/Sources/CountTokensRequest.swift @@ -43,7 +43,10 @@ public struct CountTokensResponse { /// 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. + /// The total number of billable characters in the text input given to the model as a prompt. + /// + /// > Important: This does not include billable image, video or other non-text input. See + /// [Vertex AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing) for details. public let totalBillableCharacters: Int }