fix: restore Demo moderation and image creation

CHANGELOG.md

@@ -6,6 +6,9 @@ Compatibility promise: the public API is additive-only. Anything `public` is dep
 
 ## [Unreleased]
 
+### Added
+- `Model.gpt_image_1_5` for `gpt-image-1.5` and `Model.omniModerationLatest` for `omni-moderation-latest` ([#485](https://github.com/MacPaw/OpenAI/pull/485)).
+
 ### Changed
 - CI: the *API Breakage* allowlist is emptied, because the 0.6.0 release is now the comparison baseline for the breaks it listed.
 

Demo/DemoChat/Sources/ImageStore.swift

@@ -8,12 +8,25 @@
 import Foundation
 import OpenAI
 import Combine
+import CoreGraphics
+import ImageIO
+
+struct GeneratedImage: Identifiable {
+    let id = UUID()
+    let url: String?
+    /// Decoded from `b64_json`, or `nil` when it is absent or not an image.
+    let cgImage: CGImage?
+    /// `b64_json` was present, so the view can tell a failed decode from a URL image.
+    let hasBase64: Bool
+    let revisedPrompt: String?
+}
 
 public final class ImageStore: ObservableObject {
     public var openAIClient: OpenAIProtocol
 
     @Published var imagesQueryInProgress = false
-    @Published var images: [ImagesResult.Image] = []
+    @Published var images: [GeneratedImage] = []
+    @Published var imagesQueryError: Error?
 
     private var subscription: AnyCancellable?
 
@@ -24,33 +37,63 @@ public final class ImageStore: ObservableObject {
     }
 
     @MainActor
-    func images(query: ImagesQuery) {
+    func images(prompt: String, n: Int, size: ImagesQuery.Size) {
+        cancelImagesQuery()
         imagesQueryInProgress = true
         images.removeAll()
+        imagesQueryError = nil
+        let query = ImagesQuery(prompt: prompt, model: .gpt_image_1_5, n: n, size: size)
 
         subscription = openAIClient
             .images(query: query)
+            .map(Self.makeImages)
             .receive(on: DispatchQueue.main)
-            .handleEvents(receiveCancel: {
-                self.imagesQueryInProgress = false
-            })
-            .sink(receiveCompletion: { completion in
+            .sink(receiveCompletion: { [weak self] completion in
+                guard let self else { return }
                 self.imagesQueryInProgress = false
 
                 switch completion {
                 case .finished:
                     break
                 case .failure(let failure):
-                    // TODO: Better error handling
-                    print(failure)
+                    // `cancelImagesQuery()` cancels the subscription, so its cancellation never reaches here. A cancelled
+                    // error still can when the request is cancelled underneath the subscription, for example when the
+                    // URL session is invalidated, and that isn't a failure to show.
+                    if !(failure is CancellationError), (failure as? URLError)?.code != .cancelled {
+                        self.imagesQueryError = failure
+                    }
                 }
-            }, receiveValue: { imagesResult in
-                self.images = imagesResult.data
+            }, receiveValue: { [weak self] in
+                self?.images = $0
             })
     }
+
+    // A function rather than a closure: a closure written in `images(prompt:n:size:)` would be
+    // isolated to the main actor, but this runs on the URL session's queue.
+    private static func makeImages(_ result: ImagesResult) -> [GeneratedImage] {
+        result.data.map { image in
+            GeneratedImage(
+                url: image.url,
+                cgImage: image.b64Json.flatMap(decode),
+                hasBase64: image.b64Json != nil,
+                revisedPrompt: image.revisedPrompt
+            )
+        }
+    }
+
+    private static func decode(_ base64: String) -> CGImage? {
+        guard let data = Data(base64Encoded: base64),
+              let source = CGImageSourceCreateWithData(data as CFData, nil) else {
+            return nil
+        }
+        // Decode the pixels here, off the main thread, rather than on first draw
+        return CGImageSourceCreateImageAtIndex(source, 0, [kCGImageSourceShouldCacheImmediately: true] as CFDictionary)
+    }
 
+    @MainActor
     func cancelImagesQuery() {
         subscription?.cancel()
         subscription = nil
+        imagesQueryInProgress = false
     }
 }

Demo/DemoChat/Sources/MiscStore.swift

@@ -52,7 +52,7 @@ public final class MiscStore: ObservableObject {
             let response = try await openAIClient.moderations(
                 query: ModerationsQuery(
                     input: .init(message.content),
-                    model: .textModerationLatest
+                    model: .omniModerationLatest
                 )
             )
 

Demo/DemoChat/Sources/UI/Images/ImageCreationView.swift

@@ -15,11 +15,12 @@ public struct ImageCreationView: View {
     @State private var n: Int = 1
     @State private var size = ImagesQuery.Size._1024
 
-    private var sizes = ImagesQuery.Size.allCases.dropLast(2)
+    // The sizes the GPT image models accept. `ImageStore` always requests `gpt_image_1_5`, and the other
+    // `ImagesQuery.Size` cases are for dall-e-2 (256 and 512) or dall-e-3 (1792), which the API rejects for it.
+    private let sizes: [ImagesQuery.Size] = [._1024, ._1536x1024, ._1024x1536, .auto]
 
     public init(store: ImageStore) {
         self.store = store
-        size = sizes[0]
     }
 
     public var body: some View {
@@ -47,17 +48,33 @@ public struct ImageCreationView: View {
                     actionButton
                     .foregroundColor(.accentColor)
                     Spacer()
+                    if store.imagesQueryInProgress {
+                        ProgressView()
+                    }
+                }
+                if let error = store.imagesQueryError {
+                    Text(error.localizedDescription)
+                        .foregroundColor(.red)
+                        .font(.caption)
                 }
             }
-            if !$store.images.isEmpty {
+            if !store.images.isEmpty {
                 Section("Images") {
-                    ForEach($store.images, id: \.url) { image in
-                        let urlString = image.wrappedValue.url ?? ""
-                        if let imageURL = URL(string: urlString), UIApplication.shared.canOpenURL(imageURL) {
+                    ForEach(store.images) { image in
+                        if image.hasBase64 {
+                            if let cgImage = image.cgImage {
+                                Image(cgImage, scale: 1, label: Text(image.revisedPrompt ?? "Generated image"))
+                                    .resizable()
+                                    .aspectRatio(contentMode: .fit)
+                            } else {
+                                Text("Failed to decode image.")
+                                    .foregroundColor(.red)
+                            }
+                        } else if let urlString = image.url, let imageURL = URL(string: urlString), UIApplication.shared.canOpenURL(imageURL) {
                             LinkPreview(previewURL: imageURL)
                                 .aspectRatio(contentMode: .fit)
                         } else {
-                            Text(urlString)
+                            Text(image.url ?? "No image returned.")
                                 .foregroundStyle(.secondary)
                         }
                     }
@@ -75,8 +92,7 @@ public struct ImageCreationView: View {
             }
         } else {
             Button("Create Image" + (n == 1 ? "" : "s")) {
-                let query = ImagesQuery(prompt: prompt, n: n, size: size)
-                store.images(query: query)
+                store.images(prompt: prompt, n: n, size: size)
             }
         }
     }

Demo/DemoChat/Tests/ImageStoreTests.swift

@@ -0,0 +1,87 @@
+import Foundation
+import Testing
+@testable import DemoChat
+@testable import OpenAI
+
+@MainActor
+@Suite(.timeLimit(.minutes(1)))
+struct ImageStoreTests {
+    @Test(arguments: [ImagesQuery.Size._1024, ._1536x1024, ._1024x1536, .auto], [1, 10])
+    func creationSendsExplicitModel(size: ImagesQuery.Size, n: Int) async throws {
+        let mock = MockURL { _ in
+            Data(#"{"created":1,"data":[{"url":"https://example.com/puppy.png"}]}"#.utf8)
+        }
+        defer { mock.remove() }
+        let store = ImageStore(openAIClient: mock.makeClient())
+
+        store.images(prompt: "A puppy", n: n, size: size)
+        try await sleepUntil { !store.imagesQueryInProgress }
+
+        let body = try #require(mock.lastBody)
+        let json = try #require(JSONSerialization.jsonObject(with: body) as? [String: Any])
+        #expect(json["model"] as? String == "gpt-image-1.5")
+        #expect(json["prompt"] as? String == "A puppy")
+        #expect(json["n"] as? Int == n)
+        #expect(json["size"] as? String == size.rawValue)
+        #expect(json["response_format"] == nil)
+        #expect(store.images.first?.url == "https://example.com/puppy.png")
+        #expect(store.imagesQueryError == nil)
+    }
+
+    @Test
+    func failureIsVisible() async throws {
+        let mock = MockURL { _ in
+            Data(#"{"error":{"message":"Image generation rejected","type":"invalid_request_error"}}"#.utf8)
+        }
+        defer { mock.remove() }
+        let store = ImageStore(openAIClient: mock.makeClient())
+
+        store.images(prompt: "A puppy", n: 2, size: ._1024)
+        try await sleepUntil { !store.imagesQueryInProgress }
+        #expect(store.imagesQueryError?.localizedDescription == "Image generation rejected")
+        #expect(store.images.isEmpty)
+    }
+
+    @Test
+    func retryClearsErrorAndPublishesBase64Images() async throws {
+        let png = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII="
+        let mock = MockURL { _ in
+            Data(#"{"error":{"message":"Image generation rejected","type":"invalid_request_error"}}"#.utf8)
+        }
+        defer { mock.remove() }
+        let store = ImageStore(openAIClient: mock.makeClient())
+
+        store.images(prompt: "A puppy", n: 2, size: ._1024)
+        try await sleepUntil { !store.imagesQueryInProgress }
+        #expect(store.imagesQueryError != nil)
+
+        mock.respond = { _ in
+            try JSONSerialization.data(withJSONObject: [
+                "created": 1, "data": [["b64_json": png], ["b64_json": png]]
+            ])
+        }
+
+        store.images(prompt: "A puppy", n: 2, size: ._1024)
+        #expect(store.imagesQueryError == nil)
+        try await sleepUntil { !store.imagesQueryInProgress }
+        #expect(store.imagesQueryError == nil)
+        #expect(store.images.count == 2)
+        #expect(store.images.allSatisfy { $0.hasBase64 && $0.url == nil && $0.cgImage?.width == 1 && $0.cgImage?.height == 1 })
+    }
+
+    @Test
+    func cancellationEndsProgress() async throws {
+        // Returning nil never answers the request, so it stays in flight until cancelled
+        let mock = MockURL { _ in nil }
+        defer { mock.remove() }
+        let store = ImageStore(openAIClient: mock.makeClient())
+
+        store.images(prompt: "A puppy", n: 2, size: ._1024)
+        try await sleepUntil { mock.requestCount == 1 }
+        #expect(store.imagesQueryInProgress)
+        store.cancelImagesQuery()
+        #expect(!store.imagesQueryInProgress)
+        #expect(store.imagesQueryError == nil)
+        #expect(store.images.isEmpty)
+    }
+}

Demo/DemoChat/Tests/MiscStoreTests.swift

@@ -0,0 +1,48 @@
+import Foundation
+import Testing
+@testable import DemoChat
+@testable import OpenAI
+
+@MainActor
+@Suite(.timeLimit(.minutes(1)))
+struct MiscStoreTests {
+    @Test(arguments: [200, 400])
+    func moderationUsesOmniAndPreservesConversationHandling(status: Int) async throws {
+        let mock = MockURL { _ in
+            if status == 400 {
+                return Data(#"{"error":{"message":"Moderation rejected","type":"invalid_request_error"}}"#.utf8)
+            }
+            let keys = ["harassment", "harassment/threatening", "hate", "hate/threatening",
+                        "self-harm", "self-harm/intent", "self-harm/instructions", "sexual",
+                        "sexual/minors", "violence", "violence/graphic", "illicit", "illicit/violent"]
+            return try JSONSerialization.data(withJSONObject: [
+                "id": "modr-test", "model": "omni-moderation-latest",
+                "results": [[
+                    "flagged": false,
+                    "categories": Dictionary(uniqueKeysWithValues: keys.map { ($0, false) }),
+                    "category_scores": Dictionary(uniqueKeysWithValues: keys.map { ($0, 0.0) }),
+                    "category_applied_input_types": ["illicit": ["text"]]
+                ]]
+            ])
+        }
+        defer { mock.remove() }
+        let store = MiscStore(openAIClient: mock.makeClient())
+
+        await store.sendModerationMessage(Message(id: "user", role: .user, content: "I love puppies", createdAt: Date()))
+
+        let body = try #require(mock.lastBody)
+        let query = try JSONDecoder().decode(ModerationsQuery.self, from: body)
+        #expect(query.model == "omni-moderation-latest")
+        #expect(query.input == "I love puppies")
+        #expect(store.moderationConversation.messages.first?.content == "I love puppies")
+        if status == 200 {
+            #expect(store.moderationConversationError == nil)
+            #expect(store.moderationConversation.messages.last?.role == .assistant)
+            #expect(store.moderationConversation.messages.last?.content.contains("🟢 harassment") == true)
+        } else {
+            #expect(store.moderationConversationError?.localizedDescription == "Moderation rejected")
+            #expect(store.moderationConversation.messages.count == 1)
+        }
+    }
+}
+