Handle totalBillableCharacters optional decoding in Vertex AI (#12715)

Author: andrewheard

FirebaseVertexAI/Sources/CountTokensRequest.swift

@@ -39,10 +39,30 @@ 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.
+  /// 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
 }
+
+@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
+  }
+}

FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-no-billable-characters.json

@@ -0,0 +1,3 @@
+{
+  "totalTokens": 258
+}

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)
+    XCTAssertEqual(response.totalBillableCharacters, 0)
+  }
+
   func testCountTokens_modelNotFound() async throws {
     MockURLProtocol.requestHandler = try httpRequestHandler(
       forResource: "failure-model-not-found", withExtension: "json",