bleam 0.0.10 → 0.0.11

package/templates/native/modules/bleam-runtime/ios/AIModule.swift

@@ -15,7 +15,7 @@ public final class AIModule: Module {
 
   public func definition() -> ModuleDefinition {
     Name("AI")
-    Events("imagegenProgress", "textGenerationUpdate", "toolCall")
+    Events("imagegenProgress", "textGenerationUpdate", "localTextGenerationUpdate", "toolCall")
 
     Function("isAvailable") { () -> Bool in
 #if canImport(FoundationModels)
@@ -147,13 +147,47 @@ public final class AIModule: Module {
     }
 
     AsyncFunction("cancelTextGeneration") { [weak self] (id: String) -> Bool in
-      await self?.textGeneration.cancel(id) ?? false
+      let foundationCanceled = await self?.textGeneration.cancel(id) ?? false
+      let localCanceled = (try? await ImageGenerationServiceClient.shared.cancelText(id: id)) ?? false
+      return foundationCanceled || localCanceled
     }
 
     AsyncFunction("resolveToolCall") { [weak self] (request: ToolResultRequest) -> Bool in
       await self?.textGeneration.resolve(request) ?? false
     }
 
+    AsyncFunction("generateLocalText") { (request: LocalTextGenerationRequest, promise: Promise) in
+      Task { @MainActor in
+        do {
+          let event = try await ImageGenerationServiceClient.shared.runText(
+            request: Self.bonsaiTextRequest(from: request)
+          )
+          promise.resolve(event.tokenIDs)
+        } catch {
+          Self.rejectLocalTextGeneration(error, promise: promise)
+        }
+      }
+    }
+
+    AsyncFunction("streamLocalText") { [weak self] (request: LocalTextGenerationRequest, promise: Promise) in
+      Task { @MainActor in
+        do {
+          let event = try await ImageGenerationServiceClient.shared.runText(
+            request: Self.bonsaiTextRequest(from: request),
+            onUpdate: { [weak self] event in
+              self?.sendEvent("localTextGenerationUpdate", [
+                "id": event.id,
+                "tokenIds": event.tokenIDs,
+              ])
+            }
+          )
+          promise.resolve(event.tokenIDs)
+        } catch {
+          Self.rejectLocalTextGeneration(error, promise: promise)
+        }
+      }
+    }
+
     AsyncFunction("generateImage") { (request: ImageGenerationRequest, promise: Promise) in
       do {
         let imageRequest = try Self.imageRequest(from: request)
@@ -272,12 +306,23 @@ public final class AIModule: Module {
     guard !request.model.isEmpty else {
       throw ModuleError.invalidModel("Missing image model")
     }
-    guard request.model == "flux2-klein-4b-8bit-abliterated" else {
+    guard request.model == "flux2-klein-4b-8bit-abliterated"
+      || request.model == "ternary-bonsai-image-4b"
+    else {
       throw ModuleError.invalidModel("Unsupported image model: \(request.model)")
     }
     guard !request.prompt.isEmpty else {
       throw ModuleError.invalidMessage("Expected prompt")
     }
+    if request.model == "ternary-bonsai-image-4b",
+      request.sourceImagePath != nil
+        || request.maskImagePath != nil
+        || request.outpaintWidth != nil
+        || !request.referenceImagePaths.isEmpty
+        || !request.loraPaths.isEmpty
+    {
+      throw ModuleError.invalidModel("Ternary Bonsai Image 4B currently supports text-to-image generation")
+    }
     if let sourceImagePath = request.sourceImagePath, sourceImagePath.isEmpty {
       throw ModuleError.invalidMessage("Source image path must not be empty")
     }
@@ -493,6 +538,40 @@ public final class AIModule: Module {
     )
   }
 
+  private static func bonsaiTextRequest(
+    from request: LocalTextGenerationRequest
+  ) throws -> BonsaiTextGenerationRequest {
+    guard !request.id.isEmpty else {
+      throw ModuleError.invalidMessage("Expected generation id")
+    }
+    guard !request.model.isEmpty else {
+      throw ModuleError.invalidModel("Expected Bonsai model id")
+    }
+    var isDirectory: ObjCBool = false
+    guard FileManager.default.fileExists(atPath: request.modelPath, isDirectory: &isDirectory),
+      isDirectory.boolValue
+    else {
+      throw ModuleError.invalidModel("Missing Bonsai model directory: \(request.modelPath)")
+    }
+    guard !request.inputIds.isEmpty else {
+      throw ModuleError.invalidMessage("Expected tokenized prompt")
+    }
+    return BonsaiTextGenerationRequest(
+      id: request.id,
+      model: request.model,
+      modelPath: request.modelPath,
+      inputIDs: request.inputIds,
+      endTokenIDs: request.endTokenIds,
+      maximumNewTokens: request.maxTokens,
+      temperature: request.temperature,
+      topP: request.topP,
+      topK: request.topK,
+      repetitionPenalty: request.repetitionPenalty,
+      repetitionContextSize: request.repetitionContextSize,
+      seed: request.seed
+    )
+  }
+
   private static func metadataDictionary(_ metadata: GenerationWorkerMetadata?) -> [String: Any]? {
     guard let metadata else { return nil }
     var value: [String: Any] = [
@@ -530,6 +609,22 @@ public final class AIModule: Module {
     promise.reject(error)
   }
 
+  private static func rejectLocalTextGeneration(_ error: Error, promise: Promise) {
+    if error is CancellationError {
+      promise.reject("GENERATION_CANCELED", "Generation canceled")
+      return
+    }
+    if error as? GenerationServiceClientError == .concurrentRequest {
+      promise.reject("CONCURRENT_REQUEST", error.localizedDescription)
+      return
+    }
+    if let failure = error as? GenerationWorkerFailure {
+      promise.reject("GENERATION_FAILED", failure.localizedDescription)
+      return
+    }
+    promise.reject(error)
+  }
+
 #if canImport(FoundationModels)
   @available(iOS 26, *)
   private static func session(system: String?, tools: [any Tool]) -> LanguageModelSession {
@@ -849,6 +944,7 @@ private final class ImageGenerationServiceClient {
   }
   private var connection: NSXPCConnection?
   private var pendingRun: PendingRun?
+  private var pendingTextRun: PendingTextRun?
   private let eventReceiver = GenerationServiceEventReceiver()
 
   private final class PendingRun {
@@ -864,13 +960,29 @@ private final class ImageGenerationServiceClient {
     }
   }
 
+  private final class PendingTextRun {
+    let id: String
+    let continuation: CheckedContinuation<BonsaiTextGenerationEvent, Error>
+    let onUpdate: (BonsaiTextGenerationEvent) -> Void
+
+    init(
+      id: String,
+      continuation: CheckedContinuation<BonsaiTextGenerationEvent, Error>,
+      onUpdate: @escaping (BonsaiTextGenerationEvent) -> Void
+    ) {
+      self.id = id
+      self.continuation = continuation
+      self.onUpdate = onUpdate
+    }
+  }
+
   func run(
     request: GenerationWorkerRequest,
     onProgress: @escaping (GenerationWorkerEvent) -> Void = { _ in }
   ) async throws -> GenerationWorkerEvent {
     let terminalEvent: GenerationWorkerEvent = try await withCheckedThrowingContinuation { continuation in
       do {
-        guard pendingRun == nil else {
+        guard pendingRun == nil, pendingTextRun == nil else {
           continuation.resume(throwing: GenerationServiceClientError.concurrentRequest)
           return
         }
@@ -905,6 +1017,71 @@ private final class ImageGenerationServiceClient {
     return terminalEvent
   }
 
+  func runText(
+    request: BonsaiTextGenerationRequest,
+    onUpdate: @escaping (BonsaiTextGenerationEvent) -> Void = { _ in }
+  ) async throws -> BonsaiTextGenerationEvent {
+    let terminalEvent: BonsaiTextGenerationEvent = try await withCheckedThrowingContinuation { continuation in
+      do {
+        guard pendingRun == nil, pendingTextRun == nil else {
+          continuation.resume(throwing: GenerationServiceClientError.concurrentRequest)
+          return
+        }
+        eventReceiver.onTextEvent = { [weak self] event in
+          self?.deliverText(event)
+        }
+        pendingTextRun = PendingTextRun(
+          id: request.id,
+          continuation: continuation,
+          onUpdate: onUpdate
+        )
+        let requestData = try JSONEncoder().encode(request)
+        let service = try remoteProxy(errorHandler: { [weak self] error in
+          Task { @MainActor in
+            self?.failPendingTextRun(Self.mapConnectionError(error))
+          }
+        })
+        service.startTextGeneration(requestData: requestData) { [weak self] accepted in
+          Task { @MainActor in
+            guard accepted else {
+              self?.failPendingTextRun(GenerationServiceClientError.concurrentRequest)
+              return
+            }
+          }
+        }
+      } catch {
+        pendingTextRun = nil
+        eventReceiver.onTextEvent = nil
+        continuation.resume(throwing: error)
+      }
+    }
+
+    switch terminalEvent.type {
+    case .completed:
+      return terminalEvent
+    case .canceled:
+      throw CancellationError()
+    case .failed:
+      throw GenerationWorkerFailure(message: terminalEvent.message ?? "Text generation failed.")
+    case .update:
+      throw GenerationServiceClientError.invalidResponse
+    }
+  }
+
+  func cancelText(id: String) async throws -> Bool {
+    guard pendingTextRun?.id == id else { return false }
+    let service = try remoteProxy(errorHandler: { [weak self] error in
+      Task { @MainActor in
+        self?.failPendingTextRun(Self.mapConnectionError(error))
+      }
+    })
+    return try await withCheckedThrowingContinuation { continuation in
+      service.cancelTextGeneration(id: id) { canceled in
+        continuation.resume(returning: canceled)
+      }
+    }
+  }
+
   func cancel() async throws -> Bool {
     guard pendingRun != nil else {
       return false
@@ -940,7 +1117,7 @@ private final class ImageGenerationServiceClient {
   }
 
   func unload() async throws -> Bool {
-    guard pendingRun == nil else { return false }
+    guard pendingRun == nil, pendingTextRun == nil else { return false }
     let service = try remoteProxy(errorHandler: { _ in })
     let unloaded = try await withCheckedThrowingContinuation { continuation in
       service.shutdown { accepted in
@@ -1025,6 +1202,17 @@ private final class ImageGenerationServiceClient {
     }
   }
 
+  private func deliverText(_ event: BonsaiTextGenerationEvent) {
+    guard let pendingTextRun, event.id == pendingTextRun.id else { return }
+    if event.type == .update {
+      pendingTextRun.onUpdate(event)
+      return
+    }
+    self.pendingTextRun = nil
+    eventReceiver.onTextEvent = nil
+    pendingTextRun.continuation.resume(returning: event)
+  }
+
   private func failPendingRun(_ error: Error) {
     guard let pendingRun else { return }
     self.pendingRun = nil
@@ -1035,9 +1223,17 @@ private final class ImageGenerationServiceClient {
     pendingRun.continuation.resume(throwing: error)
   }
 
+  private func failPendingTextRun(_ error: Error) {
+    guard let pendingTextRun else { return }
+    self.pendingTextRun = nil
+    eventReceiver.onTextEvent = nil
+    pendingTextRun.continuation.resume(throwing: error)
+  }
+
   private func handleConnectionLost() {
     connection = nil
     failPendingRun(GenerationServiceClientError.serviceInterrupted)
+    failPendingTextRun(GenerationServiceClientError.serviceInterrupted)
   }
 
   private func invalidate() {
@@ -1150,6 +1346,7 @@ private enum GenerationServiceClientError: LocalizedError, Equatable {
 
 private final class GenerationServiceEventReceiver: NSObject, GenerationServiceClientProtocol {
   @MainActor var onEvent: ((GenerationWorkerEvent) -> Void)?
+  @MainActor var onTextEvent: ((BonsaiTextGenerationEvent) -> Void)?
 
   func generationWorkerDidEmit(eventData: Data) {
     guard let event = try? JSONDecoder().decode(GenerationWorkerEvent.self, from: eventData) else { return }
@@ -1157,6 +1354,14 @@ private final class GenerationServiceEventReceiver: NSObject, GenerationServiceC
       onEvent?(event)
     }
   }
+
+
+  func textGenerationWorkerDidEmit(eventData: Data) {
+    guard let event = try? JSONDecoder().decode(BonsaiTextGenerationEvent.self, from: eventData) else { return }
+    Task { @MainActor in
+      onTextEvent?(event)
+    }
+  }
 }
 
 private struct TextGenerationRequest: Record {
@@ -1169,6 +1374,21 @@ private struct TextGenerationRequest: Record {
   @Field var tools: [ToolDefinitionRequest] = []
 }
 
+private struct LocalTextGenerationRequest: Record {
+  @Field var id: String = ""
+  @Field var model: String = ""
+  @Field var modelPath: String = ""
+  @Field var inputIds: [Int] = []
+  @Field var endTokenIds: [Int] = []
+  @Field var maxTokens: Int?
+  @Field var temperature: Double?
+  @Field var topP: Double?
+  @Field var topK: Int?
+  @Field var repetitionPenalty: Double?
+  @Field var repetitionContextSize: Int?
+  @Field var seed: Int?
+}
+
 private struct ToolDefinitionRequest: Record {
   @Field var name: String = ""
   @Field var description: String = ""

package/templates/native/modules/bleam-runtime/ios/GenerationContracts.swift

@@ -3,12 +3,15 @@ import Foundation
 @objc protocol GenerationServiceProtocol {
   func startGeneration(requestData: Data, reply: @escaping (Bool) -> Void)
   func cancelGeneration(reply: @escaping (Bool) -> Void)
+  func startTextGeneration(requestData: Data, reply: @escaping (Bool) -> Void)
+  func cancelTextGeneration(id: String, reply: @escaping (Bool) -> Void)
   func memoryTelemetry(reply: @escaping (Data?) -> Void)
   func shutdown(reply: @escaping (Bool) -> Void)
 }
 
 @objc protocol GenerationServiceClientProtocol {
   func generationWorkerDidEmit(eventData: Data)
+  func textGenerationWorkerDidEmit(eventData: Data)
 }
 
 struct GenerationWorkerRequest: Codable, Sendable {
@@ -133,3 +136,47 @@ struct GenerationWorkerLoRA: Codable, Sendable {
   let path: String
   let scale: Double
 }
+
+struct BonsaiTextGenerationRequest: Codable, Sendable {
+  let id: String
+  let model: String
+  let modelPath: String
+  let inputIDs: [Int]
+  let endTokenIDs: [Int]
+  let maximumNewTokens: Int?
+  let temperature: Double?
+  let topP: Double?
+  let topK: Int?
+  let repetitionPenalty: Double?
+  let repetitionContextSize: Int?
+  let seed: Int?
+}
+
+struct BonsaiTextGenerationEvent: Codable, Sendable {
+  enum EventType: String, Codable, Sendable {
+    case update
+    case completed
+    case canceled
+    case failed
+  }
+
+  let type: EventType
+  let id: String
+  let tokenIDs: [Int]
+  let finishReason: String?
+  let message: String?
+
+  init(
+    type: EventType,
+    id: String,
+    tokenIDs: [Int] = [],
+    finishReason: String? = nil,
+    message: String? = nil
+  ) {
+    self.type = type
+    self.id = id
+    self.tokenIDs = tokenIDs
+    self.finishReason = finishReason
+    self.message = message
+  }
+}

package/templates/native/modules/bleam-runtime/ios/PlatformModule.swift

@@ -6,70 +6,10 @@ import ExpoModulesCore
 import Foundation
 import Photos
 
-@objc private protocol PlatformHelperProtocol {
-  func health(reply: @escaping (Data?) -> Void)
-  func downloadPlatformUpdate(requestData: Data, reply: @escaping (Data?) -> Void)
-  func applyPlatformUpdate(requestData: Data, reply: @escaping (Bool) -> Void)
-}
-
 public final class PlatformModule: Module {
   public func definition() -> ModuleDefinition {
     Name("Platform")
 
-    AsyncFunction("helperStatus") { (promise: Promise) in
-      guard Self.isSupportedInstallLocation else {
-        UserDefaults.standard.set(false, forKey: "BleamPlatformUpdatesSupported")
-        promise.resolve([
-          "healthy": false,
-          "protocolVersion": 1,
-          "helperVersion": "unavailable",
-          "platformVersion": "unavailable",
-          "capabilities": [],
-          "updatesSupported": false,
-          "installPath": Bundle.main.bundleURL.path,
-          "reason": "reinstall_required",
-        ])
-        return
-      }
-      let serviceName = "\(Bundle.main.bundleIdentifier ?? "dev.bleam.app").PlatformHelper"
-      guard let connection = Self.makeConnection(serviceName: serviceName) else {
-        promise.reject("HELPER_UNAVAILABLE", "The app helper is unavailable")
-        return
-      }
-      connection.remoteObjectInterface = NSXPCInterface(with: PlatformHelperProtocol.self)
-      connection.resume()
-
-      guard let service = connection.remoteObjectProxyWithErrorHandler({ error in
-        connection.invalidate()
-        promise.reject("HELPER_UNAVAILABLE", error.localizedDescription)
-      }) as? PlatformHelperProtocol else {
-        connection.invalidate()
-        promise.reject("HELPER_UNAVAILABLE", "The app helper is unavailable")
-        return
-      }
-
-      service.health { data in
-        defer { connection.invalidate() }
-        guard
-          let data,
-          var status = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
-        else {
-          promise.reject("HELPER_INVALID_RESPONSE", "The app helper returned an invalid response")
-          return
-        }
-        let supported = Self.isSupportedInstallLocation
-        status["updatesSupported"] = supported
-        status["installPath"] = Bundle.main.bundleURL.path
-        if !supported {
-          status["reason"] = "reinstall_required"
-        }
-        UserDefaults.standard.set(status["healthy"], forKey: "BleamHelperHealthy")
-        UserDefaults.standard.set(status["platformVersion"], forKey: "BleamHelperPlatformVersion")
-        UserDefaults.standard.set(supported, forKey: "BleamPlatformUpdatesSupported")
-        promise.resolve(status)
-      }
-    }
-
     AsyncFunction("permissionState") { (capability: String, promise: Promise) in
       guard let state = Self.permissionState(capability) else {
         promise.reject("INVALID_CAPABILITY", "Unknown permission capability: \(capability)")
@@ -85,132 +25,6 @@ public final class PlatformModule: Module {
       }
       promise.resolve(states)
     }
-
-    AsyncFunction("downloadPlatformUpdate") { (promise: Promise) in
-      guard Self.isSupportedInstallLocation else {
-        promise.reject("REINSTALL_REQUIRED", "Platform updates require the app to be installed directly in ~/Applications")
-        return
-      }
-      guard
-        let serviceURL = Bundle.main.object(forInfoDictionaryKey: "BleamUpdateServiceURL") as? String,
-        let bundleIdentifier = Bundle.main.bundleIdentifier,
-        let teamIdentifier = Bundle.main.object(forInfoDictionaryKey: "BleamTeamIdentifier") as? String,
-        let platformVersion = Bundle.main.object(forInfoDictionaryKey: "BleamRuntimeVersion") as? String,
-        let platformFingerprint = Bundle.main.object(forInfoDictionaryKey: "BleamPlatformFingerprint") as? String,
-        let requestData = try? JSONSerialization.data(withJSONObject: [
-          "serviceURL": serviceURL,
-          "bundleIdentifier": bundleIdentifier,
-          "teamIdentifier": teamIdentifier,
-          "platformVersion": platformVersion,
-          "platformFingerprint": platformFingerprint,
-        ])
-      else {
-        promise.reject("UPDATES_UNAVAILABLE", "Platform updates are not configured")
-        return
-      }
-
-      guard let connection = Self.makeConnection(serviceName: "\(bundleIdentifier).PlatformHelper") else {
-        promise.reject("HELPER_UNAVAILABLE", "The app helper is unavailable")
-        return
-      }
-      connection.remoteObjectInterface = NSXPCInterface(with: PlatformHelperProtocol.self)
-      connection.resume()
-      guard let service = connection.remoteObjectProxyWithErrorHandler({ error in
-        connection.invalidate()
-        promise.reject("PLATFORM_DOWNLOAD_FAILED", error.localizedDescription)
-      }) as? PlatformHelperProtocol else {
-        connection.invalidate()
-        promise.reject("HELPER_UNAVAILABLE", "The app helper is unavailable")
-        return
-      }
-
-      service.downloadPlatformUpdate(requestData: requestData) { data in
-        defer { connection.invalidate() }
-        guard
-          let data,
-          let result = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
-        else {
-          promise.reject("PLATFORM_DOWNLOAD_FAILED", "The platform update could not be downloaded")
-          return
-        }
-        if let error = result["error"] as? String {
-          UserDefaults.standard.set(error, forKey: "BleamPlatformDownloadError")
-          promise.reject("PLATFORM_VERIFICATION_FAILED", error)
-          return
-        }
-        UserDefaults.standard.set(result["releaseId"], forKey: "BleamPlatformDownloadRelease")
-        promise.resolve(result)
-      }
-    }
-
-    AsyncFunction("updateAndRelaunch") { (releaseId: String, promise: Promise) in
-      guard Self.isSupportedInstallLocation else {
-        promise.reject("REINSTALL_REQUIRED", "Platform updates require the app to be installed directly in ~/Applications")
-        return
-      }
-      guard
-        let bundleIdentifier = Bundle.main.bundleIdentifier,
-        let teamIdentifier = Bundle.main.object(forInfoDictionaryKey: "BleamTeamIdentifier") as? String,
-        let requestData = try? JSONSerialization.data(withJSONObject: [
-          "releaseId": releaseId,
-          "bundleIdentifier": bundleIdentifier,
-          "teamIdentifier": teamIdentifier,
-          "installedAppPath": Bundle.main.bundleURL.path,
-          "appPID": ProcessInfo.processInfo.processIdentifier,
-        ]),
-        let connection = Self.makeConnection(serviceName: "\(bundleIdentifier).PlatformHelper")
-      else {
-        promise.reject("PLATFORM_UPDATE_UNAVAILABLE", "The platform update cannot be applied")
-        return
-      }
-      connection.remoteObjectInterface = NSXPCInterface(with: PlatformHelperProtocol.self)
-      connection.resume()
-      guard let service = connection.remoteObjectProxyWithErrorHandler({ error in
-        connection.invalidate()
-        promise.reject("PLATFORM_UPDATE_FAILED", error.localizedDescription)
-      }) as? PlatformHelperProtocol else {
-        connection.invalidate()
-        promise.reject("HELPER_UNAVAILABLE", "The app helper is unavailable")
-        return
-      }
-      service.applyPlatformUpdate(requestData: requestData) { accepted in
-        connection.invalidate()
-        UserDefaults.standard.set(accepted, forKey: "BleamPlatformApplyAccepted")
-        guard accepted else {
-          promise.reject("PLATFORM_UPDATE_REJECTED", "The platform update is not verified or the app is not installed in ~/Applications")
-          return
-        }
-        promise.resolve(nil)
-        DispatchQueue.main.asyncAfter(deadline: .now() + 0.2) {
-          exit(0)
-        }
-      }
-    }
-  }
-
-  private static func makeConnection(serviceName: String) -> NSXPCConnection? {
-#if targetEnvironment(macCatalyst)
-    let connectionClass: AnyObject = NSXPCConnection.self as AnyObject
-    let allocated = connectionClass.perform(NSSelectorFromString("alloc"))?.takeUnretainedValue() as AnyObject?
-    let selector = NSSelectorFromString("initWithMachServiceName:options:")
-    guard let allocated, allocated.responds(to: selector) else { return nil }
-
-    typealias Initializer = @convention(c) (
-      AnyObject,
-      Selector,
-      NSString,
-      UInt
-    ) -> Unmanaged<AnyObject>
-    let initialize = unsafeBitCast(allocated.method(for: selector), to: Initializer.self)
-    return initialize(
-      allocated,
-      selector,
-      serviceName as NSString,
-      0
-    ).takeRetainedValue() as? NSXPCConnection
-#else
-    return NSXPCConnection(machServiceName: serviceName)
-#endif
   }
 
   private static let permissionCapabilities = [
@@ -321,9 +135,4 @@ public final class PlatformModule: Module {
     }
   }
 
-  private static var isSupportedInstallLocation: Bool {
-    let applications = URL(fileURLWithPath: NSHomeDirectory(), isDirectory: true)
-      .appendingPathComponent("Applications", isDirectory: true).standardizedFileURL
-    return Bundle.main.bundleURL.standardizedFileURL.deletingLastPathComponent() == applications
-  }
 }