Skip to content

Commit 1b28489

Browse files
feat(FirebaseAI): Add support for imageConfig and NO_IMAGE FinishReason
1 parent 61ca043 commit 1b28489

File tree

6 files changed

+136
-1
lines changed

6 files changed

+136
-1
lines changed

FirebaseAI/Sources/GenerateContentResponse.swift

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,7 @@ public struct FinishReason: DecodableProtoEnum, Hashable, Sendable {
242242
case prohibitedContent = "PROHIBITED_CONTENT"
243243
case spii = "SPII"
244244
case malformedFunctionCall = "MALFORMED_FUNCTION_CALL"
245+
case noImage = "NO_IMAGE"
245246
}
246247

247248
/// Natural stop point of the model or provided stop sequence.
@@ -274,6 +275,9 @@ public struct FinishReason: DecodableProtoEnum, Hashable, Sendable {
274275
/// Token generation was stopped because the function call generated by the model was invalid.
275276
public static let malformedFunctionCall = FinishReason(kind: .malformedFunctionCall)
276277

278+
/// The model successfully generated an image, but it was not returned to the user.
279+
public static let noImage = FinishReason(kind: .noImage)
280+
277281
/// Returns the raw string representation of the `FinishReason` value.
278282
///
279283
/// > Note: This value directly corresponds to the values in the [REST

FirebaseAI/Sources/GenerationConfig.swift

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ public struct GenerationConfig: Sendable {
5454
/// Configuration for controlling the "thinking" behavior of compatible Gemini models.
5555
let thinkingConfig: ThinkingConfig?
5656

57+
/// Image generation parameters.
58+
let imageConfig: ImageConfig?
59+
5760
/// Creates a new `GenerationConfig` value.
5861
///
5962
/// See the
@@ -162,7 +165,7 @@ public struct GenerationConfig: Sendable {
162165
presencePenalty: Float? = nil, frequencyPenalty: Float? = nil,
163166
stopSequences: [String]? = nil, responseMIMEType: String? = nil,
164167
responseSchema: Schema? = nil, responseModalities: [ResponseModality]? = nil,
165-
thinkingConfig: ThinkingConfig? = nil) {
168+
thinkingConfig: ThinkingConfig? = nil, imageConfig: ImageConfig? = nil) {
166169
// Explicit init because otherwise if we re-arrange the above variables it changes the API
167170
// surface.
168171
self.temperature = temperature
@@ -177,6 +180,7 @@ public struct GenerationConfig: Sendable {
177180
self.responseSchema = responseSchema
178181
self.responseModalities = responseModalities
179182
self.thinkingConfig = thinkingConfig
183+
self.imageConfig = imageConfig
180184
}
181185
}
182186

@@ -197,5 +201,6 @@ extension GenerationConfig: Encodable {
197201
case responseSchema
198202
case responseModalities
199203
case thinkingConfig
204+
case imageConfig
200205
}
201206
}
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
// Copyright 2024 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
/// An aspect ratio for generated images.
16+
public struct AspectRatio: Sendable {
17+
/// The raw string value of the aspect ratio.
18+
let rawValue: String
19+
20+
/// Creates a new aspect ratio with a raw string value.
21+
private init(rawValue: String) {
22+
self.rawValue = rawValue
23+
}
24+
25+
/// Square (1:1) aspect ratio.
26+
public static let square1x1 = AspectRatio(rawValue: "1:1")
27+
28+
/// Portrait (2:3) aspect ratio.
29+
public static let portrait2x3 = AspectRatio(rawValue: "2:3")
30+
31+
/// Landscape (3:2) aspect ratio.
32+
public static let landscape3x2 = AspectRatio(rawValue: "3:2")
33+
34+
/// Portrait (3:4) aspect ratio.
35+
public static let portrait3x4 = AspectRatio(rawValue: "3:4")
36+
37+
/// Landscape (4:3) aspect ratio.
38+
public static let landscape4x3 = AspectRatio(rawValue: "4:3")
39+
40+
/// Portrait (4:5) aspect ratio.
41+
public static let portrait4x5 = AspectRatio(rawValue: "4:5")
42+
43+
/// Landscape (5:4) aspect ratio.
44+
public static let landscape5x4 = AspectRatio(rawValue: "5:4")
45+
46+
/// Portrait (9:16) aspect ratio.
47+
public static let portrait9x16 = AspectRatio(rawValue: "9:16")
48+
49+
/// Landscape (16:9) aspect ratio.
50+
public static let landscape16x9 = AspectRatio(rawValue: "16:9")
51+
52+
/// Landscape (21:9) aspect ratio.
53+
public static let landscape21x9 = AspectRatio(rawValue: "21:9")
54+
}
55+
56+
// MARK: - Codable Conformances
57+
58+
extension AspectRatio: Encodable {
59+
public func encode(to encoder: Encoder) throws {
60+
var container = encoder.singleValueContainer()
61+
try container.encode(rawValue)
62+
}
63+
}
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
// Copyright 2024 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
/// A struct defining image generation parameters.
16+
public struct ImageConfig: Sendable {
17+
/// The aspect ratio of the generated image.
18+
let aspectRatio: AspectRatio?
19+
20+
/// Creates a new `ImageConfig` value.
21+
///
22+
/// - Parameters:
23+
/// - aspectRatio: The aspect ratio of the generated image.
24+
public init(aspectRatio: AspectRatio? = nil) {
25+
self.aspectRatio = aspectRatio
26+
}
27+
}
28+
29+
// MARK: - Codable Conformances
30+
31+
extension ImageConfig: Encodable {}

FirebaseAI/Tests/Unit/GenerationConfigTests.swift

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,4 +153,22 @@ final class GenerationConfigTests: XCTestCase {
153153
}
154154
""")
155155
}
156+
157+
func testEncodeGenerationConfig_withImageConfig() throws {
158+
let aspectRatio = AspectRatio.square1x1
159+
let generationConfig = GenerationConfig(
160+
imageConfig: .init(aspectRatio: aspectRatio)
161+
)
162+
163+
let jsonData = try encoder.encode(generationConfig)
164+
165+
let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
166+
XCTAssertEqual(json, """
167+
{
168+
"imageConfig" : {
169+
"aspectRatio" : "\(aspectRatio.rawValue)"
170+
}
171+
}
172+
""")
173+
}
156174
}

FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,20 @@ final class GenerateContentResponseTests: XCTestCase {
158158
XCTAssertEqual(candidate.finishReason, .stop)
159159
}
160160

161+
func testDecodeCandidate_withNoImageFinishReason() throws {
162+
let json = """
163+
{
164+
"content": { "role": "model", "parts": [ { "text": "Some text." } ] },
165+
"finishReason": "NO_IMAGE"
166+
}
167+
"""
168+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
169+
170+
let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
171+
172+
XCTAssertEqual(candidate.finishReason, .noImage)
173+
}
174+
161175
// MARK: - Candidate.isEmpty
162176

163177
func testCandidateIsEmpty_allEmpty_isTrue() throws {

0 commit comments

Comments
 (0)