fix: restore Demo moderation and image creation

Demo/DemoChat/Sources/ImageStore.swift

@@ -14,6 +14,7 @@ public final class ImageStore: ObservableObject {
 
     @Published var imagesQueryInProgress = false
     @Published var images: [ImagesResult.Image] = []
+    @Published var imagesQueryError: Error?
 
     private var subscription: AnyCancellable?
 
@@ -24,33 +25,37 @@ 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, n: n, size: size)
 
         subscription = openAIClient
             .images(query: query)
             .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)
+                    if !(failure is CancellationError), (failure as? URLError)?.code != .cancelled {
+                        self.imagesQueryError = failure
+                    }
                 }
-            }, receiveValue: { imagesResult in
-                self.images = imagesResult.data
+            }, receiveValue: { [weak self] imagesResult in
+                self?.images = imagesResult.data
             })
     }
 
+    @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: "omni-moderation-latest"
                 )
             )
 

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

@@ -15,11 +15,10 @@ 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)
+    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 +46,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(Array(store.images.enumerated()), id: \.offset) { _, image in
+                        if let b64JSON = image.b64Json {
+                            if let data = Data(base64Encoded: b64JSON), let result = UIImage(data: data) {
+                                Image(uiImage: result)
+                                    .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 +90,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,125 @@
+import Foundation
+#if canImport(FoundationNetworking)
+import FoundationNetworking
+#endif
+import Testing
+import os
+@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 host = UUID().uuidString.lowercased() + ".test"
+        let capturedBody = OSAllocatedUnfairLock(initialState: Data?.none)
+        ImageURLProtocol.handlers.withLock { handlers in
+            handlers[host] = { request in
+                let stream = try #require(request.httpBody.map { InputStream(data: $0) } ?? request.httpBodyStream)
+                stream.open()
+                defer { stream.close() }
+                let json = try JSONSerialization.jsonObject(with: stream)
+                let body = try JSONSerialization.data(withJSONObject: json)
+                capturedBody.withLock { $0 = body }
+                return Data(#"{"created":1,"data":[{"url":"https://example.com/puppy.png"}]}"#.utf8)
+            }
+        }
+        defer { ImageURLProtocol.handlers.withLock { $0.removeValue(forKey: host) } }
+        let configuration = URLSessionConfiguration.ephemeral
+        configuration.protocolClasses = [ImageURLProtocol.self]
+        let session = URLSession(configuration: configuration)
+        defer { session.invalidateAndCancel() }
+        let client = OpenAI(configuration: .init(token: "test", host: host), session: session)
+        let store = ImageStore(openAIClient: client)
+
+        store.images(prompt: "A puppy", n: n, size: size)
+        try await sleepUntil { !store.imagesQueryInProgress }
+
+        let body = try #require(capturedBody.withLock { $0 })
+        let json = try #require(JSONSerialization.jsonObject(with: body) as? [String: Any])
+        #expect(json["model"] as? String == "gpt-image-1")
+        #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 failureIsVisibleAndRetryPublishesBase64ImagesThenCancellationEndsProgress() async throws {
+        let host = UUID().uuidString.lowercased() + ".test"
+        let requests = OSAllocatedUnfairLock(initialState: 0)
+        let png = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII="
+        ImageURLProtocol.handlers.withLock { handlers in
+            handlers[host] = { _ in
+                let number = requests.withLock { count in
+                    count += 1
+                    return count
+                }
+                switch number {
+                case 1:
+                    return Data(#"{"error":{"message":"Image generation rejected","type":"invalid_request_error"}}"#.utf8)
+                case 2:
+                    return try JSONSerialization.data(withJSONObject: [
+                        "created": 1, "data": [["b64_json": png], ["b64_json": png]]
+                    ])
+                default:
+                    return nil
+                }
+            }
+        }
+        defer { ImageURLProtocol.handlers.withLock { $0.removeValue(forKey: host) } }
+        let configuration = URLSessionConfiguration.ephemeral
+        configuration.protocolClasses = [ImageURLProtocol.self]
+        let session = URLSession(configuration: configuration)
+        defer { session.invalidateAndCancel() }
+        let client = OpenAI(configuration: .init(token: "test", host: host), session: session)
+        let store = ImageStore(openAIClient: client)
+
+        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)
+
+        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.b64Json == png && $0.url == nil })
+
+        store.images(prompt: "A puppy", n: 2, size: ._1024)
+        try await sleepUntil { requests.withLock { $0 == 3 } }
+        #expect(store.imagesQueryInProgress)
+        store.cancelImagesQuery()
+        #expect(!store.imagesQueryInProgress)
+        #expect(store.imagesQueryError == nil)
+        #expect(store.images.isEmpty)
+    }
+}
+
+private final class ImageURLProtocol: URLProtocol {
+    static let handlers = OSAllocatedUnfairLock(initialState: [String: @Sendable (URLRequest) throws -> Data?]())
+
+    override class func canInit(with request: URLRequest) -> Bool { true }
+    override class func canonicalRequest(for request: URLRequest) -> URLRequest { request }
+
+    override func startLoading() {
+        do {
+            let host = request.url!.host!
+            let handler = try #require(Self.handlers.withLock { $0[host] })
+            guard let data = try handler(request) else { return }
+            let status = data.starts(with: Data(#"{"error"#.utf8)) ? 400 : 200
+            let response = HTTPURLResponse(url: request.url!, statusCode: status, httpVersion: nil, headerFields: ["Content-Type": "application/json"])!
+            client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
+            client?.urlProtocol(self, didLoad: data)
+            client?.urlProtocolDidFinishLoading(self)
+        } catch {
+            client?.urlProtocol(self, didFailWithError: error)
+        }
+    }
+
+    override func stopLoading() {}
+}

Demo/DemoChat/Tests/MiscStoreTests.swift

@@ -0,0 +1,90 @@
+import Foundation
+#if canImport(FoundationNetworking)
+import FoundationNetworking
+#endif
+import Testing
+import os
+@testable import DemoChat
+@testable import OpenAI
+
+@MainActor
+@Suite(.timeLimit(.minutes(1)))
+struct MiscStoreTests {
+    @Test(arguments: [200, 400])
+    func moderationUsesOmniAndPreservesConversationHandling(status: Int) async throws {
+        let host = UUID().uuidString.lowercased() + ".test"
+        let capturedBody = OSAllocatedUnfairLock(initialState: Data?.none)
+        ModerationURLProtocol.handlers.withLock { handlers in
+            handlers[host] = { request in
+                let stream = try #require(request.httpBody.map { InputStream(data: $0) } ?? request.httpBodyStream)
+                stream.open()
+                defer { stream.close() }
+                let json = try JSONSerialization.jsonObject(with: stream)
+                let body = try JSONSerialization.data(withJSONObject: json)
+                capturedBody.withLock { $0 = body }
+                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 { ModerationURLProtocol.handlers.withLock { $0.removeValue(forKey: host) } }
+        let configuration = URLSessionConfiguration.ephemeral
+        configuration.protocolClasses = [ModerationURLProtocol.self]
+        let session = URLSession(configuration: configuration)
+        defer { session.invalidateAndCancel() }
+        let client = OpenAI(configuration: .init(token: "test", host: host), session: session)
+        let store = MiscStore(openAIClient: client)
+
+        await store.sendModerationMessage(Message(id: "user", role: .user, content: "I love puppies", createdAt: Date()))
+
+        let body = try #require(capturedBody.withLock { $0 })
+        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)
+        }
+    }
+}
+
+private final class ModerationURLProtocol: URLProtocol {
+    static let handlers = OSAllocatedUnfairLock(initialState: [String: @Sendable (URLRequest) throws -> Data]())
+
+    override class func canInit(with request: URLRequest) -> Bool { true }
+    override class func canonicalRequest(for request: URLRequest) -> URLRequest { request }
+
+    override func startLoading() {
+        do {
+            let host = request.url!.host!
+            let handler = try #require(Self.handlers.withLock { $0[host] })
+            let data = try handler(request)
+            let status = data.starts(with: Data(#"{"error"#.utf8)) ? 400 : 200
+            let response = HTTPURLResponse(url: request.url!, statusCode: status, httpVersion: nil, headerFields: ["Content-Type": "application/json"])!
+            client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
+            client?.urlProtocol(self, didLoad: data)
+            client?.urlProtocolDidFinishLoading(self)
+        } catch {
+            client?.urlProtocol(self, didFailWithError: error)
+        }
+    }
+
+    override func stopLoading() {}
+}