Skip to content

Commit

Permalink
[Vertex AI] Add EncodableProtoEnum protocol and fix encoding (#13862)
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewheard authored Oct 9, 2024
1 parent 20c9413 commit 67502af
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 18 deletions.
65 changes: 47 additions & 18 deletions FirebaseVertexAI/Sources/Protocols/Internal/CodableProtoEnum.swift
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,8 @@
// 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 {
/// A type that represents a Protocol Buffer raw enum value.
protocol ProtoEnum {
/// The type representing the valid values for the protobuf enum.
///
/// > Important: This type must conform to `RawRepresentable` with the `RawValue == String`.
Expand All @@ -32,8 +29,8 @@ protocol DecodableProtoEnum: Decodable {
/// ```
associatedtype Kind: RawRepresentable<String>

/// Returns the ``VertexLog/MessageCode`` associated with unrecognized (unknown) enum values.
var unrecognizedValueMessageCode: VertexLog.MessageCode { get }
/// Returns the raw string value of the enum.
var rawValue: String { get }

/// Create a new instance of the specified type from a raw enum value.
init(rawValue: String)
Expand All @@ -42,14 +39,48 @@ protocol DecodableProtoEnum: Decodable {
///
/// > Important: A default implementation is provided.
init(kind: Kind)
}

/// 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: ProtoEnum, Decodable {
/// Returns the ``VertexLog/MessageCode`` associated with unrecognized (unknown) enum values.
var unrecognizedValueMessageCode: VertexLog.MessageCode { get }

/// Creates a new instance by decoding from the given decoder.
///
/// > Important: A default implementation is provided.
init(from decoder: Decoder) throws
init(from decoder: any Decoder) throws
}

/// Default `Decodable` implementation for types conforming to `DecodableProtoEnum`.
/// A type that can be encoded as a Protocol Buffer enum value.
///
/// Protobuf enums are represented as strings in JSON. A default `Encodable` implementation is
/// provided when conforming to this type.
protocol EncodableProtoEnum: ProtoEnum, Encodable {
/// Encodes this value into the given encoder.
///
/// > Important: A default implementation is provided.
func encode(to encoder: any Encoder) throws
}

/// A type that can be decoded and encoded from a Protocol Buffer raw enum value.
///
/// See ``ProtoEnum``, ``DecodableProtoEnum`` and ``EncodableProtoEnum`` for more details.
protocol CodableProtoEnum: DecodableProtoEnum, EncodableProtoEnum {}

// MARK: - Default Implementations

// Default implementation of `init(kind: Kind)` for types conforming to `ProtoEnum`.
extension ProtoEnum {
init(kind: Kind) {
self = Self(rawValue: kind.rawValue)
}
}

// 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'.
Expand All @@ -73,14 +104,12 @@ extension DecodableProtoEnum {
}
}

/// Default implementation of `init(kind: Kind)` for types conforming to `DecodableProtoEnum`.
extension DecodableProtoEnum {
init(kind: Kind) {
self = Self(rawValue: kind.rawValue)
// Default `Encodable` implementation for types conforming to `EncodableProtoEnum`.
extension EncodableProtoEnum {
// Note: Method 'encode(to:)' must be declared public because it matches a requirement in public
// protocol 'Encodable'.
public func encode(to encoder: any Encoder) throws {
var container = encoder.singleValueContainer()
try container.encode(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 {}
8 changes: 8 additions & 0 deletions FirebaseVertexAI/Tests/Integration/IntegrationTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,13 @@ final class IntegrationTests: XCTestCase {
role: "system",
parts: "You are a friendly and helpful assistant."
)
let safetySettings = [
SafetySetting(harmCategory: .harassment, threshold: .blockLowAndAbove),
SafetySetting(harmCategory: .hateSpeech, threshold: .blockLowAndAbove),
SafetySetting(harmCategory: .sexuallyExplicit, threshold: .blockLowAndAbove),
SafetySetting(harmCategory: .dangerousContent, threshold: .blockLowAndAbove),
SafetySetting(harmCategory: .civicIntegrity, threshold: .blockLowAndAbove),
]

var vertex: VertexAI!
var model: GenerativeModel!
Expand All @@ -50,6 +57,7 @@ final class IntegrationTests: XCTestCase {
model = vertex.generativeModel(
modelName: "gemini-1.5-flash",
generationConfig: generationConfig,
safetySettings: safetySettings,
tools: [],
systemInstruction: systemInstruction
)
Expand Down

0 comments on commit 67502af

Please sign in to comment.