bleam 0.0.10 → 0.0.12

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

@@ -15,90 +15,7 @@ public final class AIModule: Module {
 
   public func definition() -> ModuleDefinition {
     Name("AI")
-    Events("imagegenProgress", "textGenerationUpdate", "toolCall")
-
-    Function("isAvailable") { () -> Bool in
-#if canImport(FoundationModels)
-      if #available(iOS 26, *) {
-        return SystemLanguageModel.default.availability == .available
-      }
-#endif
-      return false
-    }
-
-    Function("availability") { () -> [String: Any] in
-#if canImport(FoundationModels)
-      if #available(iOS 26, *) {
-        switch SystemLanguageModel.default.availability {
-        case .available:
-          return ["available": true]
-        case .unavailable(.deviceNotEligible):
-          return ["available": false, "reason": "device_not_eligible"]
-        case .unavailable(.appleIntelligenceNotEnabled):
-          return ["available": false, "reason": "apple_intelligence_not_enabled"]
-        case .unavailable(.modelNotReady):
-          return ["available": false, "reason": "model_not_ready"]
-        case .unavailable:
-          return ["available": false, "reason": "unknown"]
-        }
-      }
-#endif
-      return ["available": false, "reason": "unknown"]
-    }
-
-    AsyncFunction("generate") { [weak self] (request: TextGenerationRequest, promise: Promise) in
-#if canImport(FoundationModels)
-      if #available(iOS 26, *) {
-        guard SystemLanguageModel.default.availability == .available else {
-          promise.reject("MODEL_UNAVAILABLE", "Apple Intelligence model is not available")
-          return
-        }
-
-        guard !request.prompt.isEmpty else {
-          promise.reject("INVALID_MESSAGE", "Expected prompt")
-          return
-        }
-
-        let generationID = request.id ?? UUID().uuidString
-        let task = Task { @MainActor [weak self] in
-          guard let self else { return }
-          do {
-            let tools = try Self.tools(
-              from: request.tools,
-              generationID: generationID,
-              coordinator: self.textGeneration,
-              emit: { [weak self] event in self?.sendEvent("toolCall", event) }
-            )
-            let session = Self.session(system: request.system, tools: tools)
-            let options = try Self.generationOptions(from: request)
-
-            if let schema = request.outputSchema {
-              let generationSchema = try Self.generationSchema(from: schema)
-              let response = try await session.respond(
-                to: request.prompt,
-                schema: generationSchema,
-                includeSchemaInPrompt: true,
-                options: options
-              )
-              promise.resolve(try Self.value(from: response.content, schema: schema))
-              return
-            }
-
-            let response = try await session.respond(to: request.prompt, options: options)
-            promise.resolve(response.content)
-          } catch {
-            Self.reject(error, promise: promise)
-          }
-          self.textGeneration.finish(generationID)
-        }
-        Task { @MainActor [weak self] in self?.textGeneration.start(generationID, task: task) }
-      } else {
-        promise.reject("UNSUPPORTED_OS", "Apple Foundation Models require iOS 26 or newer")
-      }
-#else
-      promise.reject("UNSUPPORTED_OS", "Apple Foundation Models are unavailable in this build")
-#endif
-    }
+    Events("imagegenProgress", "textGenerationUpdate", "localTextGenerationUpdate")
 
     AsyncFunction("streamGenerate") { [weak self] (request: TextGenerationRequest, promise: Promise) in
 #if canImport(FoundationModels)
@@ -116,13 +33,7 @@ public final class AIModule: Module {
         let task = Task { @MainActor [weak self] in
           guard let self else { return }
           do {
-            let tools = try Self.tools(
-              from: request.tools,
-              generationID: generationID,
-              coordinator: self.textGeneration,
-              emit: { [weak self] event in self?.sendEvent("toolCall", event) }
-            )
-            let session = Self.session(system: request.system, tools: tools)
+            let session = try Self.session(system: request.system, messages: request.messages)
             let options = try Self.generationOptions(from: request)
             let stream = session.streamResponse(to: request.prompt, options: options)
             var content = ""
@@ -147,11 +58,28 @@ 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("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
@@ -272,12 +200,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 +432,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,33 +503,53 @@ public final class AIModule: Module {
     promise.reject(error)
   }
 
-#if canImport(FoundationModels)
-  @available(iOS 26, *)
-  private static func session(system: String?, tools: [any Tool]) -> LanguageModelSession {
-    if let system, !system.isEmpty {
-      return LanguageModelSession(tools: tools, instructions: system)
+  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
     }
-    return LanguageModelSession(tools: tools)
+    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 tools(
-    from definitions: [ToolDefinitionRequest],
-    generationID: String,
-    coordinator: TextGenerationCoordinator,
-    emit: @escaping @MainActor @Sendable ([String: Any]) -> Void
-  ) throws -> [any Tool] {
-    try definitions.map { definition in
-      let parameters = try generationSchema(from: definition.inputSchema)
-      return DynamicTool(
-        name: definition.name,
-        description: definition.description,
-        parameters: parameters,
-        generationID: generationID,
-        coordinator: coordinator,
-        emit: emit
-      )
+  private static func session(
+    system: String?,
+    messages: [TextMessageRequest]
+  ) throws -> LanguageModelSession {
+    if !messages.isEmpty {
+      var entries: [Transcript.Entry] = []
+      if system?.isEmpty == false {
+        entries.append(.instructions(Transcript.Instructions(
+          segments: system.map { [.text(Transcript.TextSegment(content: $0))] } ?? [],
+          toolDefinitions: []
+        )))
+      }
+      for message in messages {
+        let segments: [Transcript.Segment] = [.text(Transcript.TextSegment(content: message.content))]
+        switch message.role {
+        case "user":
+          entries.append(.prompt(Transcript.Prompt(segments: segments)))
+        case "assistant":
+          entries.append(.response(Transcript.Response(assetIDs: [], segments: segments)))
+        default:
+          throw ModuleError.invalidMessage("Unsupported transcript role: \(message.role)")
+        }
+      }
+      return LanguageModelSession(transcript: Transcript(entries: entries))
+    }
+    if let system, !system.isEmpty {
+      return LanguageModelSession(instructions: system)
     }
+    return LanguageModelSession()
   }
 
   @available(iOS 26, *)
@@ -568,131 +561,6 @@ public final class AIModule: Module {
     )
   }
 
-  @available(iOS 26, *)
-  private static func generationSchema(from schema: [String: Any]) throws -> GenerationSchema {
-    do {
-      return try GenerationSchema(root: dynamicSchema(from: schema), dependencies: [])
-    } catch {
-      throw ModuleError.invalidSchema(error.localizedDescription)
-    }
-  }
-
-  @available(iOS 26, *)
-  private static func dynamicSchema(from schema: [String: Any]) throws -> DynamicGenerationSchema {
-    guard let type = schema["type"] as? String else {
-      throw ModuleError.invalidSchema("Missing schema type")
-    }
-
-    switch type {
-    case "object":
-      return try objectSchema(from: schema)
-    case "array":
-      guard let itemSchema = schema["items"] as? [String: Any] else {
-        throw ModuleError.invalidSchema("Array schema must include items")
-      }
-      return try DynamicGenerationSchema(arrayOf: dynamicSchema(from: itemSchema))
-    case "string":
-      return DynamicGenerationSchema(type: String.self)
-    case "number":
-      return DynamicGenerationSchema(type: Double.self)
-    case "boolean":
-      return DynamicGenerationSchema(type: Bool.self)
-    default:
-      throw ModuleError.invalidSchema("Unsupported schema type: \(type)")
-    }
-  }
-
-  @available(iOS 26, *)
-  private static func objectSchema(from schema: [String: Any]) throws -> DynamicGenerationSchema {
-    let propertiesDictionary = schema["properties"] as? [String: Any] ?? [:]
-    let required = Set(schema["required"] as? [String] ?? [])
-    var properties: [DynamicGenerationSchema.Property] = []
-
-    for key in propertiesDictionary.keys.sorted() {
-      guard let propertySchema = propertiesDictionary[key] as? [String: Any] else {
-        throw ModuleError.invalidSchema("Property \(key) schema must be an object")
-      }
-
-      properties.append(
-        DynamicGenerationSchema.Property(
-          name: key,
-          description: nil,
-          schema: try dynamicSchema(from: propertySchema),
-          isOptional: !required.contains(key)
-        )
-      )
-    }
-
-    return DynamicGenerationSchema(name: schema["title"] as? String ?? "Output", properties: properties)
-  }
-
-  @available(iOS 26, *)
-  private static func value(from content: GeneratedContent, schema: [String: Any]) throws -> Any {
-    guard let type = schema["type"] as? String else {
-      throw ModuleError.invalidSchema("Missing schema type")
-    }
-
-    switch type {
-    case "object":
-      let propertiesDictionary = schema["properties"] as? [String: Any] ?? [:]
-      let output = NSMutableDictionary()
-
-      for key in propertiesDictionary.keys.sorted() {
-        guard let propertySchema = propertiesDictionary[key] as? [String: Any] else {
-          throw ModuleError.invalidSchema("Property \(key) schema must be an object")
-        }
-        output[key] = try propertyValue(from: content, property: key, schema: propertySchema)
-      }
-
-      return output
-    default:
-      throw ModuleError.invalidSchema("Root schema must be an object")
-    }
-  }
-
-  @available(iOS 26, *)
-  private static func propertyValue(
-    from content: GeneratedContent,
-    property: String,
-    schema: [String: Any]
-  ) throws -> Any {
-    guard let type = schema["type"] as? String else {
-      throw ModuleError.invalidSchema("Missing schema type")
-    }
-
-    switch type {
-    case "string":
-      return try content.value(String.self, forProperty: property)
-    case "number":
-      return try content.value(Double.self, forProperty: property)
-    case "boolean":
-      return try content.value(Bool.self, forProperty: property)
-    case "object":
-      let nestedContent = try content.value(GeneratedContent.self, forProperty: property)
-      return try value(from: nestedContent, schema: schema)
-    case "array":
-      guard let itemSchema = schema["items"] as? [String: Any], let itemType = itemSchema["type"] as? String else {
-        throw ModuleError.invalidSchema("Array schema must include items")
-      }
-
-      switch itemType {
-      case "string":
-        return try content.value([String].self, forProperty: property)
-      case "number":
-        return try content.value([Double].self, forProperty: property)
-      case "boolean":
-        return try content.value([Bool].self, forProperty: property)
-      case "object":
-        let values = try content.value([GeneratedContent].self, forProperty: property)
-        return try values.map { try value(from: $0, schema: itemSchema) }
-      default:
-        throw ModuleError.invalidSchema("Unsupported array item type: \(itemType)")
-      }
-    default:
-      throw ModuleError.invalidSchema("Unsupported schema type: \(type)")
-    }
-  }
-
   @available(iOS 26, *)
   private static func reject(_ error: Error, promise: Promise) {
     if error is CancellationError {
@@ -716,55 +584,14 @@ public final class AIModule: Module {
       }
       return
     }
-    if let toolError = error as? LanguageModelSession.ToolCallError {
-      promise.reject("TOOL_ERROR", toolError.localizedDescription)
-      return
-    }
-
     promise.reject(error)
   }
 #endif
 }
 
-#if canImport(FoundationModels)
-@available(iOS 26, *)
-private struct DynamicTool: Tool {
-  typealias Arguments = GeneratedContent
-
-  let name: String
-  let description: String
-  let parameters: GenerationSchema
-  let generationID: String
-  let coordinator: TextGenerationCoordinator
-  let emit: @MainActor @Sendable ([String: Any]) -> Void
-
-  func call(arguments: GeneratedContent) async throws -> String {
-    let data = Data(arguments.jsonString.utf8)
-    let value = try JSONSerialization.jsonObject(with: data)
-    return try await coordinator.call(
-      generationID: generationID,
-      name: name,
-      arguments: value,
-      emit: emit
-    )
-  }
-}
-#endif
-
-private struct ToolExecutionError: LocalizedError {
-  let message: String
-  var errorDescription: String? { message }
-}
-
 @MainActor
 private final class TextGenerationCoordinator {
-  private struct PendingToolCall {
-    let generationID: String
-    let continuation: CheckedContinuation<String, Error>
-  }
-
   private var tasks: [String: Task<Void, Never>] = [:]
-  private var toolCalls: [String: PendingToolCall] = [:]
 
   func start(_ id: String, task: Task<Void, Never>) {
     tasks[id]?.cancel()
@@ -773,67 +600,22 @@ private final class TextGenerationCoordinator {
 
   func finish(_ id: String) {
     tasks[id] = nil
-    failToolCalls(for: id, error: CancellationError())
   }
 
   func cancel(_ id: String) -> Bool {
     guard let task = tasks.removeValue(forKey: id) else { return false }
     task.cancel()
-    failToolCalls(for: id, error: CancellationError())
     return true
   }
-
-  func call(
-    generationID: String,
-    name: String,
-    arguments: Any,
-    emit: @escaping @MainActor @Sendable ([String: Any]) -> Void
-  ) async throws -> String {
-    try Task.checkCancellation()
-    let callID = UUID().uuidString
-    return try await withCheckedThrowingContinuation { continuation in
-      toolCalls[callID] = PendingToolCall(generationID: generationID, continuation: continuation)
-      emit([
-        "generationId": generationID,
-        "callId": callID,
-        "name": name,
-        "arguments": arguments,
-      ])
-    }
-  }
-
-  func resolve(_ request: ToolResultRequest) -> Bool {
-    guard let pending = toolCalls.removeValue(forKey: request.callId),
-      pending.generationID == request.generationId
-    else { return false }
-
-    if let error = request.error {
-      pending.continuation.resume(throwing: ToolExecutionError(message: error))
-    } else {
-      pending.continuation.resume(returning: request.output ?? "")
-    }
-    return true
-  }
-
-  private func failToolCalls(for generationID: String, error: Error) {
-    let matches = toolCalls.filter { $0.value.generationID == generationID }
-    for (id, pending) in matches {
-      toolCalls[id] = nil
-      pending.continuation.resume(throwing: error)
-    }
-  }
 }
 
 private enum ModuleError: Error, LocalizedError {
-  case invalidSchema(String)
   case invalidMessage(String)
   case invalidModel(String)
   case imageGeneration(String)
 
   var errorDescription: String? {
     switch self {
-    case .invalidSchema(let message):
-      return "Invalid schema: \(message)"
     case .invalidMessage(let message), .invalidModel(let message), .imageGeneration(let message):
       return message
     }
@@ -849,6 +631,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 +647,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 +704,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 +804,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 +889,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 +910,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 +1033,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,29 +1041,43 @@ 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 {
   @Field var id: String?
   @Field var system: String?
   @Field var prompt: String = ""
-  @Field var outputSchema: [String: Any]?
+  @Field var messages: [TextMessageRequest] = []
   @Field var temperature: Double?
   @Field var maxTokens: Int?
-  @Field var tools: [ToolDefinitionRequest] = []
 }
 
-private struct ToolDefinitionRequest: Record {
-  @Field var name: String = ""
-  @Field var description: String = ""
-  @Field var inputSchema: [String: Any] = [:]
+private struct TextMessageRequest: Record {
+  @Field var role: String = ""
+  @Field var content: String = ""
 }
 
-private struct ToolResultRequest: Record {
-  @Field var generationId: String = ""
-  @Field var callId: String = ""
-  @Field var output: String?
-  @Field var error: String?
+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 ImageGenerationRequest: Record {