feat: add ImageConfig and new FinishReasons for image generation

FirebaseAI/Sources/GenerateContentResponse.swift

@@ -242,6 +242,11 @@ public struct FinishReason: DecodableProtoEnum, Hashable, Sendable {
     case prohibitedContent = "PROHIBITED_CONTENT"
     case spii = "SPII"
     case malformedFunctionCall = "MALFORMED_FUNCTION_CALL"
+    case noImage = "NO_IMAGE"
+    case imageSafety = "IMAGE_SAFETY"
+    case imageProhibitedContent = "IMAGE_PROHIBITED_CONTENT"
+    case imageRecitation = "IMAGE_RECITATION"
+    case imageOther = "IMAGE_OTHER"
   }
 
   /// Natural stop point of the model or provided stop sequence.
@@ -274,6 +279,21 @@ public struct FinishReason: DecodableProtoEnum, Hashable, Sendable {
   /// Token generation was stopped because the function call generated by the model was invalid.
   public static let malformedFunctionCall = FinishReason(kind: .malformedFunctionCall)
 
+  /// The model successfully generated an image, but it was not returned to the user.
+  public static let noImage = FinishReason(kind: .noImage)
+
+  /// Image generation stopped due to safety settings.
+  public static let imageSafety = FinishReason(kind: .imageSafety)
+
+  /// Image generation stopped because generated images has other prohibited content.
+  public static let imageProhibitedContent = FinishReason(kind: .imageProhibitedContent)
+
+  /// Image generation stopped due to recitation.
+  public static let imageRecitation = FinishReason(kind: .imageRecitation)
+
+  /// Image generation stopped because of other miscellaneous issue.
+  public static let imageOther = FinishReason(kind: .imageOther)
+
   /// Returns the raw string representation of the `FinishReason` value.
   ///
   /// > Note: This value directly corresponds to the values in the [REST

FirebaseAI/Sources/GenerationConfig.swift

@@ -54,6 +54,9 @@ public struct GenerationConfig: Sendable {
   /// Configuration for controlling the "thinking" behavior of compatible Gemini models.
   let thinkingConfig: ThinkingConfig?
 
+  /// Image generation parameters.
+  let imageConfig: ImageConfig?
+
   /// Creates a new `GenerationConfig` value.
   ///
   /// See the
@@ -162,7 +165,7 @@ public struct GenerationConfig: Sendable {
               presencePenalty: Float? = nil, frequencyPenalty: Float? = nil,
               stopSequences: [String]? = nil, responseMIMEType: String? = nil,
               responseSchema: Schema? = nil, responseModalities: [ResponseModality]? = nil,
-              thinkingConfig: ThinkingConfig? = nil) {
+              thinkingConfig: ThinkingConfig? = nil, imageConfig: ImageConfig? = nil) {
     // Explicit init because otherwise if we re-arrange the above variables it changes the API
     // surface.
     self.temperature = temperature
@@ -177,6 +180,7 @@ public struct GenerationConfig: Sendable {
     self.responseSchema = responseSchema
     self.responseModalities = responseModalities
     self.thinkingConfig = thinkingConfig
+    self.imageConfig = imageConfig
   }
 }
 
@@ -197,5 +201,6 @@ extension GenerationConfig: Encodable {
     case responseSchema
     case responseModalities
     case thinkingConfig
+    case imageConfig
   }
 }

FirebaseAI/Sources/Types/Public/AspectRatio.swift

@@ -0,0 +1,63 @@
+// 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.
+
+/// An aspect ratio for generated images.
+public struct AspectRatio: Sendable {
+  /// The raw string value of the aspect ratio.
+  let rawValue: String
+
+  /// Creates a new aspect ratio with a raw string value.
+  private init(rawValue: String) {
+    self.rawValue = rawValue
+  }
+
+  /// Square (1:1) aspect ratio.
+  public static let square1x1 = AspectRatio(rawValue: "1:1")
+
+  /// Portrait (2:3) aspect ratio.
+  public static let portrait2x3 = AspectRatio(rawValue: "2:3")
+
+  /// Landscape (3:2) aspect ratio.
+  public static let landscape3x2 = AspectRatio(rawValue: "3:2")
+
+  /// Portrait (3:4) aspect ratio.
+  public static let portrait3x4 = AspectRatio(rawValue: "3:4")
+
+  /// Landscape (4:3) aspect ratio.
+  public static let landscape4x3 = AspectRatio(rawValue: "4:3")
+
+  /// Portrait (4:5) aspect ratio.
+  public static let portrait4x5 = AspectRatio(rawValue: "4:5")
+
+  /// Landscape (5:4) aspect ratio.
+  public static let landscape5x4 = AspectRatio(rawValue: "5:4")
+
+  /// Portrait (9:16) aspect ratio.
+  public static let portrait9x16 = AspectRatio(rawValue: "9:16")
+
+  /// Landscape (16:9) aspect ratio.
+  public static let landscape16x9 = AspectRatio(rawValue: "16:9")
+
+  /// Landscape (21:9) aspect ratio.
+  public static let landscape21x9 = AspectRatio(rawValue: "21:9")
+}
+
+// MARK: - Codable Conformances
+
+extension AspectRatio: Encodable {
+  public func encode(to encoder: Encoder) throws {
+    var container = encoder.singleValueContainer()
+    try container.encode(rawValue)
+  }
+}

FirebaseAI/Sources/Types/Public/ImageConfig.swift

@@ -0,0 +1,31 @@
+// 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 struct defining image generation parameters.
+public struct ImageConfig: Sendable {
+  /// The aspect ratio of the generated image.
+  let aspectRatio: AspectRatio?
+
+  /// Creates a new `ImageConfig` value.
+  ///
+  /// - Parameters:
+  ///   - aspectRatio: The aspect ratio of the generated image.
+  public init(aspectRatio: AspectRatio? = nil) {
+    self.aspectRatio = aspectRatio
+  }
+}
+
+// MARK: - Codable Conformances
+
+extension ImageConfig: Encodable {}

FirebaseAI/Tests/Unit/GenerationConfigTests.swift

@@ -153,4 +153,22 @@
     }
     """)
   }
+
+  func testEncodeGenerationConfig_withImageConfig() throws {
+    let aspectRatio = AspectRatio.square1x1
+    let generationConfig = GenerationConfig(
+      imageConfig: .init(aspectRatio: aspectRatio)
+    )
+
+    let jsonData = try encoder.encode(generationConfig)
+
+    let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
+    XCTAssertEqual(json, """
+    {
+      "imageConfig" : {
+        "aspectRatio" : "\(aspectRatio.rawValue)"
+      }
+    }
+    """)
+  }
 }

FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift

@@ -158,6 +158,76 @@ final class GenerateContentResponseTests: XCTestCase {
     XCTAssertEqual(candidate.finishReason, .stop)
   }
 
+  func testDecodeCandidate_withNoImageFinishReason() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "NO_IMAGE"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertEqual(candidate.finishReason, .noImage)
+  }
+
+  func testDecodeCandidate_withImageSafetyFinishReason() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "IMAGE_SAFETY"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertEqual(candidate.finishReason, .imageSafety)
+  }
+
+  func testDecodeCandidate_withImageProhibitedContentFinishReason() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "IMAGE_PROHIBITED_CONTENT"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertEqual(candidate.finishReason, .imageProhibitedContent)
+  }
+
+  func testDecodeCandidate_withImageRecitationFinishReason() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "IMAGE_RECITATION"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertEqual(candidate.finishReason, .imageRecitation)
+  }
+
+  func testDecodeCandidate_withImageOtherFinishReason() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "IMAGE_OTHER"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertEqual(candidate.finishReason, .imageOther)
+  }
+
   // MARK: - Candidate.isEmpty
 
   func testCandidateIsEmpty_allEmpty_isTrue() throws {