diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift index a7d7da85d67..b8684588a0c 100644 --- a/FirebaseAI/Sources/GenerateContentResponse.swift +++ b/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 diff --git a/FirebaseAI/Sources/GenerationConfig.swift b/FirebaseAI/Sources/GenerationConfig.swift index 27c4310f12d..63c180a86cf 100644 --- a/FirebaseAI/Sources/GenerationConfig.swift +++ b/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 } } diff --git a/FirebaseAI/Sources/Types/Public/AspectRatio.swift b/FirebaseAI/Sources/Types/Public/AspectRatio.swift new file mode 100644 index 00000000000..85608988026 --- /dev/null +++ b/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) + } +} diff --git a/FirebaseAI/Sources/Types/Public/ImageConfig.swift b/FirebaseAI/Sources/Types/Public/ImageConfig.swift new file mode 100644 index 00000000000..ed76256ac64 --- /dev/null +++ b/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 {} diff --git a/FirebaseAI/Tests/Unit/GenerationConfigTests.swift b/FirebaseAI/Tests/Unit/GenerationConfigTests.swift index 22bcd70b035..cf57925c63f 100644 --- a/FirebaseAI/Tests/Unit/GenerationConfigTests.swift +++ b/FirebaseAI/Tests/Unit/GenerationConfigTests.swift @@ -153,4 +153,22 @@ final class GenerationConfigTests: XCTestCase { } """) } + + 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)" + } + } + """) + } } diff --git a/FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift b/FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift index dfc393e2d29..668b016a947 100644 --- a/FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift +++ b/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 {