[Vertex AI] Add SourceImage enum to ImageConversionError (#13575)

firebase · Sep 4, 2024 · a6fd721 · a6fd721
1 parent 52b9c4b
commit a6fd721
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 16 deletions.
diff --git a/FirebaseVertexAI/CHANGELOG.md b/FirebaseVertexAI/CHANGELOG.md
@@ -6,6 +6,9 @@
   asynchronous and must be called with `try await`. (#13545, #13573)
 - [changed] **Breaking Change**: Creating a chat instance (`startChat`) is now
   asynchronous and must be called with `await`. (#13545)
+- [changed] **Breaking Change**: The source image in the
+  `ImageConversionError.couldNotConvertToJPEG` error case is now an enum value
+  instead of the `Any` type. (#13575)
 
 # 10.29.0
 - [feature] Added community support for watchOS. (#13215)

diff --git a/FirebaseVertexAI/Sources/PartsRepresentable+Image.swift b/FirebaseVertexAI/Sources/PartsRepresentable+Image.swift
@@ -25,15 +25,28 @@ private let imageCompressionQuality: CGFloat = 0.8
 /// For some image types like `CIImage`, creating valid model content requires creating a JPEG
 /// representation of the image that may not yet exist, which may be computationally expensive.
 public enum ImageConversionError: Error {
+  /// The image that could not be converted.
+  public enum SourceImage {
+    #if canImport(UIKit)
+      case uiImage(UIImage)
+    #elseif canImport(AppKit)
+      case nsImage(NSImage)
+    #endif // canImport(UIKit)
+    case cgImage(CGImage)
+    #if canImport(CoreImage)
+      case ciImage(CIImage)
+    #endif // canImport(CoreImage)
+  }
+
   /// The image (the receiver of the call `toModelContentParts()`) was invalid.
   case invalidUnderlyingImage
 
   /// A valid image destination could not be allocated.
   case couldNotAllocateDestination
 
   /// JPEG image data conversion failed, accompanied by the original image, which may be an
-  /// instance of `NSImageRep`, `UIImage`, `CGImage`, or `CIImage`.
-  case couldNotConvertToJPEG(Any)
+  /// instance of `NSImage`, `UIImage`, `CGImage`, or `CIImage`.
+  case couldNotConvertToJPEG(SourceImage)
 }
 
 #if canImport(UIKit)
@@ -42,7 +55,7 @@ public enum ImageConversionError: Error {
   extension UIImage: ThrowingPartsRepresentable {
     public func tryPartsValue() throws -> [ModelContent.Part] {
       guard let data = jpegData(compressionQuality: imageCompressionQuality) else {
-        throw ImageConversionError.couldNotConvertToJPEG(self)
+        throw ImageConversionError.couldNotConvertToJPEG(.uiImage(self))
       }
       return [ModelContent.Part.data(mimetype: "image/jpeg", data)]
     }
@@ -59,7 +72,7 @@ public enum ImageConversionError: Error {
       let bmp = NSBitmapImageRep(cgImage: cgImage)
       guard let data = bmp.representation(using: .jpeg, properties: [.compressionFactor: 0.8])
       else {
-        throw ImageConversionError.couldNotConvertToJPEG(bmp)
+        throw ImageConversionError.couldNotConvertToJPEG(.nsImage(self))
       }
       return [ModelContent.Part.data(mimetype: "image/jpeg", data)]
     }
@@ -84,7 +97,7 @@ public enum ImageConversionError: Error {
       if CGImageDestinationFinalize(imageDestination) {
         return [.data(mimetype: "image/jpeg", output as Data)]
       }
-      throw ImageConversionError.couldNotConvertToJPEG(self)
+      throw ImageConversionError.couldNotConvertToJPEG(.cgImage(self))
     }
   }
 #endif // !os(watchOS)
@@ -105,7 +118,7 @@ public enum ImageConversionError: Error {
       if let jpegData = jpegData {
         return [.data(mimetype: "image/jpeg", jpegData)]
       }
-      throw ImageConversionError.couldNotConvertToJPEG(self)
+      throw ImageConversionError.couldNotConvertToJPEG(.ciImage(self))
     }
   }
 #endif // canImport(CoreImage)
diff --git a/FirebaseVertexAI/Tests/Unit/PartsRepresentableTests.swift b/FirebaseVertexAI/Tests/Unit/PartsRepresentableTests.swift
@@ -60,22 +60,24 @@ final class PartsRepresentableTests: XCTestCase {
       let image = CIImage.empty()
       do {
         _ = try image.tryPartsValue()
+        XCTFail("Expected model content from invalid image to error")
       } catch {
         guard let imageError = (error as? ImageConversionError) else {
           XCTFail("Got unexpected error type: \(error)")
           return
         }
         switch imageError {
         case let .couldNotConvertToJPEG(source):
-          // String(describing:) works around a type error.
-          XCTAssertEqual(String(describing: source), String(describing: image))
-          return
-        case _:
+          guard case let .ciImage(ciImage) = source else {
+            XCTFail("Unexpected image source: \(source)")
+            return
+          }
+          XCTAssertEqual(ciImage, image)
+        default:
           XCTFail("Expected image conversion error, got \(imageError) instead")
           return
         }
       }
-      XCTFail("Expected model content from invalid image to error")
     }
   #endif // canImport(CoreImage)
 
@@ -84,22 +86,24 @@ final class PartsRepresentableTests: XCTestCase {
       let image = UIImage()
       do {
         _ = try image.tryPartsValue()
+        XCTFail("Expected model content from invalid image to error")
       } catch {
         guard let imageError = (error as? ImageConversionError) else {
           XCTFail("Got unexpected error type: \(error)")
           return
         }
         switch imageError {
         case let .couldNotConvertToJPEG(source):
-          // String(describing:) works around a type error.
-          XCTAssertEqual(String(describing: source), String(describing: image))
-          return
-        case _:
+          guard case let .uiImage(uiImage) = source else {
+            XCTFail("Unexpected image source: \(source)")
+            return
+          }
+          XCTAssertEqual(uiImage, image)
+        default:
           XCTFail("Expected image conversion error, got \(imageError) instead")
           return
         }
       }
-      XCTFail("Expected model content from invalid image to error")
     }
 
     func testModelContentFromUIImageIsNotEmpty() throws {