[Firebase AI] Add includeSafetyAttributes for Imagen requests (#15191)

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: andrewheard

.github/workflows/firebaseai.yml

@@ -82,10 +82,17 @@ jobs:
       setup_command: scripts/update_vertexai_responses.sh
 
   quickstart:
-    runs-on: macos-15
+    strategy:
+      matrix:
+        include:
+          - os: macos-15
+            xcode: Xcode_16.4
+    runs-on: ${{ matrix.os }}
     env:
       BRANCH_NAME: ${{ github.head_ref || github.ref_name || 'main' }}
     steps:
     - uses: actions/checkout@v4
+    - name: Xcode
+      run: sudo xcode-select -s /Applications/${{ matrix.xcode }}.app/Contents/Developer
     - name: Build Quickstart
       run: scripts/quickstart_build_spm.sh FirebaseAI

FirebaseAI/Sources/Types/Internal/Imagen/ImageGenerationParameters.swift

@@ -23,6 +23,7 @@ struct ImageGenerationParameters {
   let outputOptions: ImageGenerationOutputOptions?
   let addWatermark: Bool?
   let includeResponsibleAIFilterReason: Bool?
+  let includeSafetyAttributes: Bool?
 }
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
@@ -42,6 +43,7 @@ extension ImageGenerationParameters: Encodable {
     case outputOptions
     case addWatermark
     case includeResponsibleAIFilterReason = "includeRaiReason"
+    case includeSafetyAttributes
   }
 
   func encode(to encoder: any Encoder) throws {
@@ -58,5 +60,6 @@ extension ImageGenerationParameters: Encodable {
       includeResponsibleAIFilterReason,
       forKey: .includeResponsibleAIFilterReason
     )
+    try container.encodeIfPresent(includeSafetyAttributes, forKey: .includeSafetyAttributes)
   }
 }

FirebaseAI/Sources/Types/Internal/Imagen/ImagenSafetyAttributes.swift

@@ -0,0 +1,24 @@
+// Copyright 2025 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.
+
+import Foundation
+
+/// A `safetyAttributes` "prediction" from Imagen.
+///
+/// This prediction is currently unused by the SDK and is only checked to be valid JSON. This type
+/// is currently only used to avoid logging unsupported prediction types.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+struct ImagenSafetyAttributes: Decodable {
+  let safetyAttributes: JSONObject
+}

FirebaseAI/Sources/Types/Public/Imagen/ImagenGenerationResponse.swift

@@ -60,6 +60,8 @@ extension ImagenGenerationResponse: Decodable where T: Decodable {
         images.append(image)
       } else if let filteredReason = try? predictionsContainer.decode(RAIFilteredReason.self) {
         filteredReasons.append(filteredReason.raiFilteredReason)
+      } else if let _ = try? predictionsContainer.decode(ImagenSafetyAttributes.self) {
+        // Ignore SafetyAttributes "prediction" to avoid logging in `unsupportedPrediction` below.
       } else if let unsupportedPrediction = try? predictionsContainer.decode(JSONObject.self) {
         AILog.warning(
           code: .decodedUnsupportedImagenPredictionType,

FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift

@@ -159,7 +159,8 @@ public final class ImagenModel {
         )
       },
       addWatermark: generationConfig?.addWatermark,
-      includeResponsibleAIFilterReason: true
+      includeResponsibleAIFilterReason: true,
+      includeSafetyAttributes: true
     )
   }
 }

FirebaseAI/Tests/Unit/Types/Imagen/ImageGenerationParametersTests.swift

@@ -34,7 +34,8 @@ final class ImageGenerationParametersTests: XCTestCase {
       personGeneration: nil,
       outputOptions: nil,
       addWatermark: nil,
-      includeResponsibleAIFilterReason: true
+      includeResponsibleAIFilterReason: true,
+      includeSafetyAttributes: true
     )
 
     let parameters = ImagenModel.imageGenerationParameters(
@@ -57,7 +58,8 @@ final class ImageGenerationParametersTests: XCTestCase {
       personGeneration: nil,
       outputOptions: nil,
       addWatermark: nil,
-      includeResponsibleAIFilterReason: true
+      includeResponsibleAIFilterReason: true,
+      includeSafetyAttributes: true
     )
 
     let parameters = ImagenModel.imageGenerationParameters(
@@ -95,7 +97,8 @@ final class ImageGenerationParametersTests: XCTestCase {
         compressionQuality: imageFormat.compressionQuality
       ),
       addWatermark: addWatermark,
-      includeResponsibleAIFilterReason: true
+      includeResponsibleAIFilterReason: true,
+      includeSafetyAttributes: true
     )
 
     let parameters = ImagenModel.imageGenerationParameters(
@@ -124,7 +127,8 @@ final class ImageGenerationParametersTests: XCTestCase {
       personGeneration: personFilterLevel.rawValue,
       outputOptions: nil,
       addWatermark: nil,
-      includeResponsibleAIFilterReason: true
+      includeResponsibleAIFilterReason: true,
+      includeSafetyAttributes: true
     )
 
     let parameters = ImagenModel.imageGenerationParameters(
@@ -170,7 +174,8 @@ final class ImageGenerationParametersTests: XCTestCase {
         compressionQuality: imageFormat.compressionQuality
       ),
       addWatermark: addWatermark,
-      includeResponsibleAIFilterReason: true
+      includeResponsibleAIFilterReason: true,
+      includeSafetyAttributes: true
     )
 
     let parameters = ImagenModel.imageGenerationParameters(
@@ -200,6 +205,7 @@ final class ImageGenerationParametersTests: XCTestCase {
     let outputOptions = ImageGenerationOutputOptions(mimeType: mimeType, compressionQuality: nil)
     let addWatermark = false
     let includeRAIReason = true
+    let includeSafetyAttributes = true
     let parameters = ImageGenerationParameters(
       sampleCount: sampleCount,
       storageURI: storageURI,
@@ -209,7 +215,8 @@ final class ImageGenerationParametersTests: XCTestCase {
       personGeneration: personGeneration,
       outputOptions: outputOptions,
       addWatermark: addWatermark,
-      includeResponsibleAIFilterReason: includeRAIReason
+      includeResponsibleAIFilterReason: includeRAIReason,
+      includeSafetyAttributes: includeSafetyAttributes
     )
 
     let jsonData = try encoder.encode(parameters)
@@ -220,6 +227,7 @@ final class ImageGenerationParametersTests: XCTestCase {
       "addWatermark" : \(addWatermark),
       "aspectRatio" : "\(aspectRatio)",
       "includeRaiReason" : \(includeRAIReason),
+      "includeSafetyAttributes" : \(includeSafetyAttributes),
       "negativePrompt" : "\(negativePrompt)",
       "outputOptions" : {
         "mimeType" : "\(mimeType)"
@@ -246,7 +254,8 @@ final class ImageGenerationParametersTests: XCTestCase {
       personGeneration: nil,
       outputOptions: nil,
       addWatermark: addWatermark,
-      includeResponsibleAIFilterReason: nil
+      includeResponsibleAIFilterReason: nil,
+      includeSafetyAttributes: nil
     )
 
     let jsonData = try encoder.encode(parameters)
@@ -272,7 +281,8 @@ final class ImageGenerationParametersTests: XCTestCase {
       personGeneration: nil,
       outputOptions: nil,
       addWatermark: nil,
-      includeResponsibleAIFilterReason: nil
+      includeResponsibleAIFilterReason: nil,
+      includeSafetyAttributes: nil
     )
 
     let jsonData = try encoder.encode(parameters)

FirebaseAI/Tests/Unit/Types/Imagen/ImagenGenerationRequestTests.swift

@@ -25,6 +25,7 @@ final class ImagenGenerationRequestTests: XCTestCase {
   let aspectRatio = "16:9"
   let safetyFilterLevel = "block_low_and_above"
   let includeResponsibleAIFilterReason = true
+  let includeSafetyAttributes = true
   lazy var parameters = ImageGenerationParameters(
     sampleCount: sampleCount,
     storageURI: nil,
@@ -34,7 +35,8 @@ final class ImagenGenerationRequestTests: XCTestCase {
     personGeneration: nil,
     outputOptions: nil,
     addWatermark: nil,
-    includeResponsibleAIFilterReason: includeResponsibleAIFilterReason
+    includeResponsibleAIFilterReason: includeResponsibleAIFilterReason,
+    includeSafetyAttributes: includeSafetyAttributes
   )
   let apiConfig = FirebaseAI.defaultVertexAIAPIConfig
 
@@ -108,6 +110,7 @@ final class ImagenGenerationRequestTests: XCTestCase {
       "parameters" : {
         "aspectRatio" : "\(aspectRatio)",
         "includeRaiReason" : \(includeResponsibleAIFilterReason),
+        "includeSafetyAttributes" : \(includeSafetyAttributes),
         "safetySetting" : "\(safetyFilterLevel)",
         "sampleCount" : \(sampleCount)
       }
@@ -137,6 +140,7 @@ final class ImagenGenerationRequestTests: XCTestCase {
       "parameters" : {
         "aspectRatio" : "\(aspectRatio)",
         "includeRaiReason" : \(includeResponsibleAIFilterReason),
+        "includeSafetyAttributes" : \(includeSafetyAttributes),
         "safetySetting" : "\(safetyFilterLevel)",
         "sampleCount" : \(sampleCount)
       }