[Firebase AI] Handle empty or unknown Part data

FirebaseAI/Sources/AILog.swift

@@ -62,6 +62,7 @@ enum AILog {
     case decodedInvalidCitationPublicationDate = 3011
     case generateContentResponseUnrecognizedContentModality = 3012
     case decodedUnsupportedImagenPredictionType = 3013
+    case decodedUnsupportedPartData = 3014
 
     // SDK State Errors
     case generateContentResponseNoCandidates = 4000

FirebaseAI/Sources/ModelContent.swift

@@ -39,9 +39,17 @@ struct InternalPart: Equatable, Sendable {
     case fileData(FileData)
     case functionCall(FunctionCall)
     case functionResponse(FunctionResponse)
+
+    struct UnsupportedDataError: Error {
+      let decodingError: DecodingError
+
+      var localizedDescription: String {
+        decodingError.localizedDescription
+      }
+    }
   }
 
-  let data: OneOfData
+  let data: OneOfData?
 
   let isThought: Bool?
 
@@ -65,7 +73,7 @@ public struct ModelContent: Equatable, Sendable {
 
   /// The data parts comprising this ``ModelContent`` value.
   public var parts: [any Part] {
-    return internalParts.map { part -> any Part in
+    return internalParts.compactMap { part -> (any Part)? in
       switch part.data {
       case let .text(text):
         return TextPart(text, isThought: part.isThought, thoughtSignature: part.thoughtSignature)
@@ -85,6 +93,9 @@ public struct ModelContent: Equatable, Sendable {
         return FunctionResponsePart(
           functionResponse, isThought: part.isThought, thoughtSignature: part.thoughtSignature
         )
+      case .none:
+        // Filter out parts that contain missing or unrecognized data
+        return nil
       }
     }
   }
@@ -179,7 +190,14 @@ extension InternalPart: Codable {
   }
 
   public init(from decoder: Decoder) throws {
-    data = try OneOfData(from: decoder)
+    do {
+      data = try OneOfData(from: decoder)
+    } catch let error as OneOfData.UnsupportedDataError {
+      AILog.error(code: .decodedUnsupportedPartData, error.localizedDescription)
+      data = nil
+    } catch { // Re-throw any other error types
+      throw error
+    }
     let container = try decoder.container(keyedBy: CodingKeys.self)
     isThought = try container.decodeIfPresent(Bool.self, forKey: .isThought)
     thoughtSignature = try container.decodeIfPresent(String.self, forKey: .thoughtSignature)
@@ -226,9 +244,11 @@ extension InternalPart.OneOfData: Codable {
       self = try .functionResponse(values.decode(FunctionResponse.self, forKey: .functionResponse))
     } else {
       let unexpectedKeys = values.allKeys.map { $0.stringValue }
-      throw DecodingError.dataCorrupted(DecodingError.Context(
-        codingPath: values.codingPath,
-        debugDescription: "Unexpected Part type(s): \(unexpectedKeys)"
+      throw UnsupportedDataError(decodingError: DecodingError.dataCorrupted(
+        DecodingError.Context(
+          codingPath: values.codingPath,
+          debugDescription: "Unexpected Part type(s): \(unexpectedKeys)"
+        )
       ))
     }
   }

FirebaseAI/Tests/TestApp/Sources/Constants.swift

@@ -24,6 +24,7 @@ public enum ModelNames {
   public static let gemini2Flash = "gemini-2.0-flash-001"
   public static let gemini2FlashLite = "gemini-2.0-flash-lite-001"
   public static let gemini2FlashPreviewImageGeneration = "gemini-2.0-flash-preview-image-generation"
+  public static let gemini2_5_FlashImagePreview = "gemini-2.5-flash-image-preview"
   public static let gemini2_5_Flash = "gemini-2.5-flash"
   public static let gemini2_5_Pro = "gemini-2.5-pro"
   public static let gemma3_4B = "gemma-3-4b-it"

FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift

@@ -322,14 +322,20 @@ struct GenerateContentIntegrationTests {
   }
 
   @Test(arguments: [
-    InstanceConfig.vertexAI_v1beta,
-    InstanceConfig.vertexAI_v1beta_global,
-    InstanceConfig.googleAI_v1beta,
+    (InstanceConfig.vertexAI_v1beta, ModelNames.gemini2FlashPreviewImageGeneration),
+    (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2FlashPreviewImageGeneration),
+    (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2_5_FlashImagePreview),
+    (InstanceConfig.googleAI_v1beta, ModelNames.gemini2FlashPreviewImageGeneration),
+    (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_FlashImagePreview),
     // Note: The following configs are commented out for easy one-off manual testing.
-    // InstanceConfig.googleAI_v1beta_staging,
-    // InstanceConfig.googleAI_v1beta_freeTier_bypassProxy,
+    // (InstanceConfig.googleAI_v1beta_staging, ModelNames.gemini2FlashPreviewImageGeneration)
+    // (InstanceConfig.googleAI_v1beta_freeTier, ModelNames.gemini2FlashPreviewImageGeneration),
+    // (
+    //  InstanceConfig.googleAI_v1beta_freeTier_bypassProxy,
+    //  ModelNames.gemini2FlashPreviewImageGeneration
+    // ),
   ])
-  func generateImage(_ config: InstanceConfig) async throws {
+  func generateImage(_ config: InstanceConfig, modelName: String) async throws {
     let generationConfig = GenerationConfig(
       temperature: 0.0,
       topP: 0.0,
@@ -342,7 +348,7 @@ struct GenerateContentIntegrationTests {
       $0.harmCategory != .civicIntegrity
     }
     let model = FirebaseAI.componentInstance(config).generativeModel(
-      modelName: ModelNames.gemini2FlashPreviewImageGeneration,
+      modelName: modelName,
       generationConfig: generationConfig,
       safetySettings: safetySettings
     )