[Vertex AI] Extract common protobuf enum to struct decoding logic (#13859)

Author: andrewheard

FirebaseVertexAI/Sources/Protocols/Internal/CodableProtoEnum.swift

@@ -0,0 +1,86 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// A type that can be decoded from a Protocol Buffer raw enum value.
+///
+/// Protobuf enums are represented as strings in JSON. A default `Decodable` implementation is
+/// provided when conforming to this type.
+protocol DecodableProtoEnum: Decodable {
+  /// The type representing the valid values for the protobuf enum.
+  ///
+  /// > Important: This type must conform to `RawRepresentable` with the `RawValue == String`.
+  ///
+  /// This is typically a Swift enum, e.g.:
+  /// ```
+  /// enum Kind: String {
+  ///   case north = "WIND_DIRECTION_NORTH"
+  ///   case south = "WIND_DIRECTION_SOUTH"
+  ///   case east = "WIND_DIRECTION_EAST"
+  ///   case west = "WIND_DIRECTION_WEST"
+  /// }
+  /// ```
+  associatedtype Kind: RawRepresentable<String>
+
+  /// Returns the ``VertexLog/MessageCode`` associated with unrecognized (unknown) enum values.
+  var unrecognizedValueMessageCode: VertexLog.MessageCode { get }
+
+  /// Create a new instance of the specified type from a raw enum value.
+  init(rawValue: String)
+
+  /// Creates a new instance from the ``Kind``'s raw value.
+  ///
+  /// > Important: A default implementation is provided.
+  init(kind: Kind)
+
+  /// Creates a new instance by decoding from the given decoder.
+  ///
+  /// > Important: A default implementation is provided.
+  init(from decoder: Decoder) throws
+}
+
+/// Default `Decodable` implementation for types conforming to `DecodableProtoEnum`.
+extension DecodableProtoEnum {
+  // Note: Initializer 'init(from:)' must be declared public because it matches a requirement in
+  // public protocol 'Decodable'.
+  public init(from decoder: Decoder) throws {
+    let rawValue = try decoder.singleValueContainer().decode(String.self)
+
+    self = Self(rawValue: rawValue)
+
+    if Kind(rawValue: rawValue) == nil {
+      VertexLog.error(
+        code: unrecognizedValueMessageCode,
+        """
+        Unrecognized \(Self.self) with value "\(rawValue)":
+        - Check for updates to the SDK as support for "\(rawValue)" may have been added; see \
+        release notes at https://firebase.google.com/support/release-notes/ios
+        - Search for "\(rawValue)" in the Firebase Apple SDK Issue Tracker at \
+        https://github.com/firebase/firebase-ios-sdk/issues and file a Bug Report if none found
+        """
+      )
+    }
+  }
+}
+
+/// Default implementation of `init(kind: Kind)` for types conforming to `DecodableProtoEnum`.
+extension DecodableProtoEnum {
+  init(kind: Kind) {
+    self = Self(rawValue: kind.rawValue)
+  }
+}
+
+/// A type that can be decoded and encoded from a Protocol Buffer raw enum value.
+///
+/// See ``DecodableProtoEnum`` for more details.
+protocol CodableProtoEnum: DecodableProtoEnum, Encodable {}

FirebaseVertexAI/Sources/Safety.swift

@@ -41,7 +41,7 @@ public struct SafetyRating: Equatable, Hashable, Sendable {
   /// The probability that a given model output falls under a harmful content category.
   ///
   /// > Note: This does not indicate the severity of harm for a piece of content.
-  public struct HarmProbability: Sendable, Equatable, Hashable {
+  public struct HarmProbability: DecodableProtoEnum, Hashable, Sendable {
     enum Kind: String {
       case negligible = "NEGLIGIBLE"
       case low = "LOW"
@@ -79,24 +79,8 @@ public struct SafetyRating: Equatable, Hashable, Sendable {
     /// > API](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/GenerateContentResponse#SafetyRating).
     public let rawValue: String
 
-    init(kind: Kind) {
-      rawValue = kind.rawValue
-    }
-
-    init(rawValue: String) {
-      if Kind(rawValue: rawValue) == nil {
-        VertexLog.error(
-          code: .generateContentResponseUnrecognizedHarmProbability,
-          """
-          Unrecognized HarmProbability with value "\(rawValue)":
-          - Check for updates to the SDK as support for "\(rawValue)" may have been added; see \
-          release notes at https://firebase.google.com/support/release-notes/ios
-          - Search for "\(rawValue)" in the Firebase Apple SDK Issue Tracker at \
-          https://github.com/firebase/firebase-ios-sdk/issues and file a Bug Report if none found
-          """
-        )
-      }
-      self.rawValue = rawValue
+    var unrecognizedValueMessageCode: VertexLog.MessageCode {
+      .generateContentResponseUnrecognizedHarmProbability
     }
   }
 }
@@ -139,7 +123,7 @@ public struct SafetySetting {
 }
 
 /// Categories describing the potential harm a piece of content may pose.
-public struct HarmCategory: Sendable, Equatable, Hashable {
+public struct HarmCategory: CodableProtoEnum, Hashable, Sendable {
   enum Kind: String {
     case harassment = "HARM_CATEGORY_HARASSMENT"
     case hateSpeech = "HARM_CATEGORY_HATE_SPEECH"
@@ -179,48 +163,16 @@ public struct HarmCategory: Sendable, Equatable, Hashable {
   /// > [REST API](https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/HarmCategory).
   public let rawValue: String
 
-  init(kind: Kind) {
-    rawValue = kind.rawValue
-  }
-
-  init(rawValue: String) {
-    if Kind(rawValue: rawValue) == nil {
-      VertexLog.error(
-        code: .generateContentResponseUnrecognizedHarmCategory,
-        """
-        Unrecognized HarmCategory with value "\(rawValue)":
-        - Check for updates to the SDK as support for "\(rawValue)" may have been added; see \
-        release notes at https://firebase.google.com/support/release-notes/ios
-        - Search for "\(rawValue)" in the Firebase Apple SDK Issue Tracker at \
-        https://github.com/firebase/firebase-ios-sdk/issues and file a Bug Report if none found
-        """
-      )
-    }
-    self.rawValue = rawValue
+  var unrecognizedValueMessageCode: VertexLog.MessageCode {
+    .generateContentResponseUnrecognizedHarmCategory
   }
 }
 
 // MARK: - Codable Conformances
 
-@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-extension SafetyRating.HarmProbability: Decodable {
-  public init(from decoder: Decoder) throws {
-    let rawValue = try decoder.singleValueContainer().decode(String.self)
-    self = SafetyRating.HarmProbability(rawValue: rawValue)
-  }
-}
-
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension SafetyRating: Decodable {}
 
-@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-extension HarmCategory: Codable {
-  public init(from decoder: Decoder) throws {
-    let rawValue = try decoder.singleValueContainer().decode(String.self)
-    self = HarmCategory(rawValue: rawValue)
-  }
-}
-
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension SafetySetting.HarmBlockThreshold: Encodable {}