Skip to content

Commit 67502af

Browse files
authored
[Vertex AI] Add EncodableProtoEnum protocol and fix encoding (#13862)
1 parent 20c9413 commit 67502af

File tree

2 files changed

+55
-18
lines changed

2 files changed

+55
-18
lines changed

FirebaseVertexAI/Sources/Protocols/Internal/CodableProtoEnum.swift

Lines changed: 47 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,8 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15-
/// A type that can be decoded from a Protocol Buffer raw enum value.
16-
///
17-
/// Protobuf enums are represented as strings in JSON. A default `Decodable` implementation is
18-
/// provided when conforming to this type.
19-
protocol DecodableProtoEnum: Decodable {
15+
/// A type that represents a Protocol Buffer raw enum value.
16+
protocol ProtoEnum {
2017
/// The type representing the valid values for the protobuf enum.
2118
///
2219
/// > Important: This type must conform to `RawRepresentable` with the `RawValue == String`.
@@ -32,8 +29,8 @@ protocol DecodableProtoEnum: Decodable {
3229
/// ```
3330
associatedtype Kind: RawRepresentable<String>
3431

35-
/// Returns the ``VertexLog/MessageCode`` associated with unrecognized (unknown) enum values.
36-
var unrecognizedValueMessageCode: VertexLog.MessageCode { get }
32+
/// Returns the raw string value of the enum.
33+
var rawValue: String { get }
3734

3835
/// Create a new instance of the specified type from a raw enum value.
3936
init(rawValue: String)
@@ -42,14 +39,48 @@ protocol DecodableProtoEnum: Decodable {
4239
///
4340
/// > Important: A default implementation is provided.
4441
init(kind: Kind)
42+
}
43+
44+
/// A type that can be decoded from a Protocol Buffer raw enum value.
45+
///
46+
/// Protobuf enums are represented as strings in JSON. A default `Decodable` implementation is
47+
/// provided when conforming to this type.
48+
protocol DecodableProtoEnum: ProtoEnum, Decodable {
49+
/// Returns the ``VertexLog/MessageCode`` associated with unrecognized (unknown) enum values.
50+
var unrecognizedValueMessageCode: VertexLog.MessageCode { get }
4551

4652
/// Creates a new instance by decoding from the given decoder.
4753
///
4854
/// > Important: A default implementation is provided.
49-
init(from decoder: Decoder) throws
55+
init(from decoder: any Decoder) throws
5056
}
5157

52-
/// Default `Decodable` implementation for types conforming to `DecodableProtoEnum`.
58+
/// A type that can be encoded as a Protocol Buffer enum value.
59+
///
60+
/// Protobuf enums are represented as strings in JSON. A default `Encodable` implementation is
61+
/// provided when conforming to this type.
62+
protocol EncodableProtoEnum: ProtoEnum, Encodable {
63+
/// Encodes this value into the given encoder.
64+
///
65+
/// > Important: A default implementation is provided.
66+
func encode(to encoder: any Encoder) throws
67+
}
68+
69+
/// A type that can be decoded and encoded from a Protocol Buffer raw enum value.
70+
///
71+
/// See ``ProtoEnum``, ``DecodableProtoEnum`` and ``EncodableProtoEnum`` for more details.
72+
protocol CodableProtoEnum: DecodableProtoEnum, EncodableProtoEnum {}
73+
74+
// MARK: - Default Implementations
75+
76+
// Default implementation of `init(kind: Kind)` for types conforming to `ProtoEnum`.
77+
extension ProtoEnum {
78+
init(kind: Kind) {
79+
self = Self(rawValue: kind.rawValue)
80+
}
81+
}
82+
83+
// Default `Decodable` implementation for types conforming to `DecodableProtoEnum`.
5384
extension DecodableProtoEnum {
5485
// Note: Initializer 'init(from:)' must be declared public because it matches a requirement in
5586
// public protocol 'Decodable'.
@@ -73,14 +104,12 @@ extension DecodableProtoEnum {
73104
}
74105
}
75106

76-
/// Default implementation of `init(kind: Kind)` for types conforming to `DecodableProtoEnum`.
77-
extension DecodableProtoEnum {
78-
init(kind: Kind) {
79-
self = Self(rawValue: kind.rawValue)
107+
// Default `Encodable` implementation for types conforming to `EncodableProtoEnum`.
108+
extension EncodableProtoEnum {
109+
// Note: Method 'encode(to:)' must be declared public because it matches a requirement in public
110+
// protocol 'Encodable'.
111+
public func encode(to encoder: any Encoder) throws {
112+
var container = encoder.singleValueContainer()
113+
try container.encode(rawValue)
80114
}
81115
}
82-
83-
/// A type that can be decoded and encoded from a Protocol Buffer raw enum value.
84-
///
85-
/// See ``DecodableProtoEnum`` for more details.
86-
protocol CodableProtoEnum: DecodableProtoEnum, Encodable {}

FirebaseVertexAI/Tests/Integration/IntegrationTests.swift

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,13 @@ final class IntegrationTests: XCTestCase {
2929
role: "system",
3030
parts: "You are a friendly and helpful assistant."
3131
)
32+
let safetySettings = [
33+
SafetySetting(harmCategory: .harassment, threshold: .blockLowAndAbove),
34+
SafetySetting(harmCategory: .hateSpeech, threshold: .blockLowAndAbove),
35+
SafetySetting(harmCategory: .sexuallyExplicit, threshold: .blockLowAndAbove),
36+
SafetySetting(harmCategory: .dangerousContent, threshold: .blockLowAndAbove),
37+
SafetySetting(harmCategory: .civicIntegrity, threshold: .blockLowAndAbove),
38+
]
3239

3340
var vertex: VertexAI!
3441
var model: GenerativeModel!
@@ -50,6 +57,7 @@ final class IntegrationTests: XCTestCase {
5057
model = vertex.generativeModel(
5158
modelName: "gemini-1.5-flash",
5259
generationConfig: generationConfig,
60+
safetySettings: safetySettings,
5361
tools: [],
5462
systemInstruction: systemInstruction
5563
)

0 commit comments

Comments
 (0)