Skip to content

Commit 5b5da99

Browse files
feat(FirebaseAI): Add support for imageConfig and new FinishReason cases
- Defines the new `ImageConfig` struct. - Defines a new `AspectRatio` struct with constants for common aspect ratios, which is used by `ImageConfig`. - Updates `GenerationConfig` to include the optional `imageConfig` field. - Adds a unit test to verify the correct encoding of `GenerationConfig` with `imageConfig`. - Aligns `ImageConfig` with `ThinkingConfig` by removing unnecessary annotations and making properties internal. - Adds `NO_IMAGE`, `IMAGE_SAFETY`, `IMAGE_PROHIBITED_CONTENT`, `IMAGE_RECITATION`, and `IMAGE_OTHER` to the `FinishReason` enum in `GenerateContentResponse`. - Adds unit tests to verify the correct deserialization of the new finish reasons.
1 parent 1b28489 commit 5b5da99

File tree

2 files changed

+72
-0
lines changed

2 files changed

+72
-0
lines changed

FirebaseAI/Sources/GenerateContentResponse.swift

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,10 @@ public struct FinishReason: DecodableProtoEnum, Hashable, Sendable {
243243
case spii = "SPII"
244244
case malformedFunctionCall = "MALFORMED_FUNCTION_CALL"
245245
case noImage = "NO_IMAGE"
246+
case imageSafety = "IMAGE_SAFETY"
247+
case imageProhibitedContent = "IMAGE_PROHIBITED_CONTENT"
248+
case imageRecitation = "IMAGE_RECITATION"
249+
case imageOther = "IMAGE_OTHER"
246250
}
247251

248252
/// Natural stop point of the model or provided stop sequence.
@@ -278,6 +282,18 @@ public struct FinishReason: DecodableProtoEnum, Hashable, Sendable {
278282
/// The model successfully generated an image, but it was not returned to the user.
279283
public static let noImage = FinishReason(kind: .noImage)
280284

285+
/// Image generation stopped due to safety settings.
286+
public static let imageSafety = FinishReason(kind: .imageSafety)
287+
288+
/// Image generation stopped because generated images has other prohibited content.
289+
public static let imageProhibitedContent = FinishReason(kind: .imageProhibitedContent)
290+
291+
/// Image generation stopped due to recitation.
292+
public static let imageRecitation = FinishReason(kind: .imageRecitation)
293+
294+
/// Image generation stopped because of other miscellaneous issue.
295+
public static let imageOther = FinishReason(kind: .imageOther)
296+
281297
/// Returns the raw string representation of the `FinishReason` value.
282298
///
283299
/// > Note: This value directly corresponds to the values in the [REST

FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,62 @@ final class GenerateContentResponseTests: XCTestCase {
172172
XCTAssertEqual(candidate.finishReason, .noImage)
173173
}
174174

175+
func testDecodeCandidate_withImageSafetyFinishReason() throws {
176+
let json = """
177+
{
178+
"content": { "role": "model", "parts": [ { "text": "Some text." } ] },
179+
"finishReason": "IMAGE_SAFETY"
180+
}
181+
"""
182+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
183+
184+
let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
185+
186+
XCTAssertEqual(candidate.finishReason, .imageSafety)
187+
}
188+
189+
func testDecodeCandidate_withImageProhibitedContentFinishReason() throws {
190+
let json = """
191+
{
192+
"content": { "role": "model", "parts": [ { "text": "Some text." } ] },
193+
"finishReason": "IMAGE_PROHIBITED_CONTENT"
194+
}
195+
"""
196+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
197+
198+
let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
199+
200+
XCTAssertEqual(candidate.finishReason, .imageProhibitedContent)
201+
}
202+
203+
func testDecodeCandidate_withImageRecitationFinishReason() throws {
204+
let json = """
205+
{
206+
"content": { "role": "model", "parts": [ { "text": "Some text." } ] },
207+
"finishReason": "IMAGE_RECITATION"
208+
}
209+
"""
210+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
211+
212+
let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
213+
214+
XCTAssertEqual(candidate.finishReason, .imageRecitation)
215+
}
216+
217+
func testDecodeCandidate_withImageOtherFinishReason() throws {
218+
let json = """
219+
{
220+
"content": { "role": "model", "parts": [ { "text": "Some text." } ] },
221+
"finishReason": "IMAGE_OTHER"
222+
}
223+
"""
224+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
225+
226+
let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
227+
228+
XCTAssertEqual(candidate.finishReason, .imageOther)
229+
}
230+
175231
// MARK: - Candidate.isEmpty
176232

177233
func testCandidateIsEmpty_allEmpty_isTrue() throws {

0 commit comments

Comments
 (0)