Skip to content

Commit

Permalink
[Vertex AI] Extract common protobuf enum to struct decoding logic (fi…
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewheard authored and MojtabaHs committed Oct 17, 2024
1 parent 90df54c commit 83dc2cc
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 54 deletions.
86 changes: 86 additions & 0 deletions FirebaseVertexAI/Sources/Protocols/Internal/CodableProtoEnum.swift
Original file line number Diff line number Diff line change
@@ -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 {}
60 changes: 6 additions & 54 deletions FirebaseVertexAI/Sources/Safety.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
}
}
}
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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 {}

Expand Down

0 comments on commit 83dc2cc

Please sign in to comment.