@capgo/capacitor-llm 8.0.11 → 8.0.13

package/ios/Sources/LLMPlugin/LLMPlugin.swift

@@ -1,5 +1,7 @@
 import Foundation
 import Capacitor
+import ObjectiveC
+import Darwin
 #if canImport(FoundationModels)
 import FoundationModels
 #endif
@@ -16,7 +18,7 @@ import MediaPipeTasksGenAI
 
 @objc(LLMPlugin)
 public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
-    private let pluginVersion: String = "8.0.11"
+    private let pluginVersion: String = "8.0.13"
     #if targetEnvironment(macCatalyst)
     private let catalystNotSupportedMessage = "CapgoLLM is not available on Mac Catalyst. Use an iOS/iPadOS target instead."
     #endif
@@ -34,6 +36,23 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
     private enum ModelType {
         case appleIntelligence
         case mediaPipe
+        case execuTorch
+    }
+
+    private final class ExecuTorchChatSession {
+        private var turns: [(role: String, content: String)] = []
+
+        func buildPrompt(userMessage: String) -> String {
+            var parts = turns.map { "\($0.role): \($0.content)" }
+            parts.append("User: \(userMessage)")
+            parts.append("Assistant:")
+            return parts.joined(separator: "\n")
+        }
+
+        func addTurn(userMessage: String, assistantMessage: String) {
+            turns.append((role: "User", content: userMessage))
+            turns.append((role: "Assistant", content: assistantMessage))
+        }
     }
 
     private var currentModelType: ModelType = .appleIntelligence
@@ -56,6 +75,14 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
     private var llmInference: LlmInference?
     #endif
 
+    private var execuTorchRunner: DynamicExecuTorchRunner?
+    private var execuTorchChats: [String: ExecuTorchChatSession] = [:]
+    private let execuTorchGenerationLock = NSLock()
+    private let execuTorchNotLinkedMessage = "ExecuTorch is not linked in this iOS app. " +
+        "It is not installed by CocoaPods; add the ExecuTorch Swift Package products " +
+        "to the app target on iOS 17+ or choose another engine."
+    private var execuTorchIsGenerating = false
+
     @objc func setModel(_ call: CAPPluginCall) {
         guard let path = call.getString("path") else {
             call.reject("Path is required")
@@ -74,13 +101,33 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
         let topk = call.getInt("topk") ?? 40
         let temperature = call.getFloat("temperature") ?? 0.8
         let modelType = call.getString("modelType") // Optional model type parameter
+        let engine = call.getString("engine")?.lowercased() ?? "auto"
+        let tokenizerPath = call.getString("tokenizerPath")
+        let specialTokens = call.getArray("specialTokens", String.self) ?? []
+        let sequenceLength = call.getInt("sequenceLength") ?? maxTokens
+
+        guard sequenceLength > 0 else {
+            call.reject("sequenceLength must be > 0")
+            return
+        }
+
+        guard !isExecuTorchGenerating() else {
+            call.reject("Cannot switch model while ExecuTorch is generating. Wait for aiFinished.")
+            return
+        }
 
         modelPath = path
+        execuTorchChats.removeAll()
+        #if !targetEnvironment(macCatalyst) && (COCOAPODS || canImport(MediaPipeTasksGenAI))
+        llmInference = nil
+        #endif
+        execuTorchRunner = nil
         isReady = false
 
         // Check if user wants to use Apple Intelligence
-        if path.lowercased() == "apple intelligence" {
+        if engine == "apple" || path.lowercased() == "apple intelligence" {
             print("[CapgoLLM] Switching to Apple Intelligence")
+            execuTorchRunner = nil
             currentModelType = .appleIntelligence
             isReady = true
             notifyListeners("readinessChange", data: ["readiness": "ready"])
@@ -88,7 +135,22 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
             return
         }
 
+        if shouldUseExecuTorch(engine: engine, path: path, tokenizerPath: tokenizerPath) {
+            currentModelType = .execuTorch
+            setExecuTorchModel(
+                call,
+                path: path,
+                tokenizerPath: tokenizerPath,
+                modelType: modelType,
+                specialTokens: specialTokens,
+                sequenceLength: sequenceLength,
+                temperature: temperature
+            )
+            return
+        }
+
         #if !targetEnvironment(macCatalyst) && canImport(MediaPipeTasksGenAI)
+        currentModelType = .mediaPipe
         Task {
             do {
                 let modelURL: URL
@@ -153,6 +215,8 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
                 notifyListeners("readinessChange", data: ["readiness": "ready"])
                 call.resolve()
             } catch {
+                llmInference = nil
+                isReady = false
                 print("[CapgoLLM] Error loading model: \(error.localizedDescription)")
                 notifyListeners("readinessChange", data: ["readiness": "Failed to load model: \(error.localizedDescription)"])
                 call.reject("Failed to load model: \(error.localizedDescription)")
@@ -163,6 +227,108 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
         #endif
     }
 
+    private func shouldUseExecuTorch(engine: String, path: String, tokenizerPath: String?) -> Bool {
+        if path.lowercased() == "apple intelligence" {
+            return false
+        }
+        if engine == "executorch" {
+            return true
+        }
+        if engine == "apple" || engine == "mediapipe" {
+            return false
+        }
+        if tokenizerPath?.isEmpty == false {
+            return true
+        }
+        return path.lowercased().hasSuffix(".pte")
+    }
+
+    private func setExecuTorchModel(
+        _ call: CAPPluginCall,
+        path: String,
+        tokenizerPath: String?,
+        modelType: String?,
+        specialTokens: [String],
+        sequenceLength: Int,
+        temperature: Float
+    ) {
+        execuTorchRunner = nil
+
+        guard sequenceLength > 0 else {
+            call.reject("sequenceLength must be > 0")
+            return
+        }
+
+        guard DynamicExecuTorchRunner.isAvailable else {
+            isReady = false
+            notifyListeners("readinessChange", data: ["readiness": execuTorchNotLinkedMessage])
+            call.reject(execuTorchNotLinkedMessage)
+            return
+        }
+
+        guard let tokenizerPath = tokenizerPath, !tokenizerPath.isEmpty else {
+            call.reject("tokenizerPath is required for ExecuTorch models")
+            return
+        }
+
+        Task {
+            do {
+                let modelURL = try resolveModelURL(path: path, modelType: modelType, defaultExtension: "pte")
+                let tokenizerURL = try resolveModelURL(path: tokenizerPath, modelType: nil, defaultExtension: "model")
+
+                print("[CapgoLLM] Initializing ExecuTorch with model at: \(modelURL.path)")
+                print("[CapgoLLM] Initializing ExecuTorch with tokenizer at: \(tokenizerURL.path)")
+
+                execuTorchRunner = try DynamicExecuTorchRunner(
+                    modelPath: modelURL.path,
+                    tokenizerPath: tokenizerURL.path,
+                    specialTokens: specialTokens,
+                    sequenceLength: sequenceLength,
+                    temperature: Double(temperature)
+                )
+                currentModelType = .execuTorch
+                isReady = true
+
+                notifyListeners("readinessChange", data: ["readiness": "ready"])
+                call.resolve()
+            } catch {
+                execuTorchRunner = nil
+                isReady = false
+                print("[CapgoLLM] Error loading ExecuTorch model: \(error.localizedDescription)")
+                notifyListeners("readinessChange", data: ["readiness": "Failed to load ExecuTorch model: \(error.localizedDescription)"])
+                call.reject("Failed to load ExecuTorch model: \(error.localizedDescription)")
+            }
+        }
+    }
+
+    private func resolveModelURL(path: String, modelType: String?, defaultExtension: String) throws -> URL {
+        if path.hasPrefix("/") {
+            return URL(fileURLWithPath: path)
+        }
+
+        let fileName: String
+        let fileExtension: String
+
+        if let providedType = modelType, !providedType.isEmpty {
+            fileName = (path as NSString).deletingPathExtension
+            fileExtension = providedType
+        } else {
+            fileName = (path as NSString).deletingPathExtension
+            let extractedExtension = (path as NSString).pathExtension
+            fileExtension = extractedExtension.isEmpty ? defaultExtension : extractedExtension
+        }
+
+        guard let bundlePath = Bundle.main.path(forResource: fileName, ofType: fileExtension) else {
+            throw NSError(
+                domain: "CapgoLLM",
+                code: 404,
+                userInfo: [NSLocalizedDescriptionKey: "Model file not found in bundle: \(path)"]
+            )
+        }
+
+        return URL(fileURLWithPath: bundlePath)
+    }
+
     @objc func createChat(_ call: CAPPluginCall) {
         #if targetEnvironment(macCatalyst)
         call.reject(catalystNotSupportedMessage)
@@ -205,6 +371,18 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
             #else
             call.reject("MediaPipe is not available on this platform. Please install via CocoaPods on iOS.")
             #endif
+
+        case .execuTorch:
+            guard execuTorchRunner != nil else {
+                call.reject("ExecuTorch model not loaded")
+                return
+            }
+            let id = UUID().uuidString
+            execuTorchChats[id] = ExecuTorchChatSession()
+            print("[CapgoLLM] Created ExecuTorch chat with ID: \(id)")
+            call.resolve([
+                "id": id
+            ])
         }
     }
 
@@ -247,6 +425,9 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
 
         case .mediaPipe:
             return isReady ? "ready" : "not_ready"
+
+        case .execuTorch:
+            return isReady ? "ready" : "not_ready"
         }
     }
 
@@ -384,7 +565,71 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
             #else
             call.reject("MediaPipe is not available on this platform. Please install via CocoaPods on iOS.")
             #endif
+
+        case .execuTorch:
+            guard isReady, let runner = execuTorchRunner else {
+                call.reject("ExecuTorch model not loaded")
+                return
+            }
+            guard let chat = execuTorchChats[chatId] else {
+                call.reject("chat not found")
+                return
+            }
+            guard beginExecuTorchGeneration() else {
+                call.reject("ExecuTorch is already generating. Wait for aiFinished before sending another message.")
+                return
+            }
+
+            Task {
+                defer {
+                    self.finishExecuTorchGeneration()
+                }
+
+                do {
+                    let prompt = chat.buildPrompt(userMessage: message)
+                    var fullResponse = ""
+                    print("[CapgoLLM] Attempting to generate ExecuTorch response for message: \(message)")
+                    try runner.generate(prompt) { token in
+                        self.notifyListeners("textFromAi", data: [
+                            "chatId": chatId,
+                            "text": token,
+                            "isChunk": true
+                        ])
+                        fullResponse += token
+                    }
+                    chat.addTurn(userMessage: message, assistantMessage: fullResponse)
+                    self.notifyListeners("aiFinished", data: ["chatId": chatId])
+                    call.resolve()
+                } catch {
+                    print("[CapgoLLM] ExecuTorch error details: \(error)")
+                    call.reject("Failed to generate ExecuTorch response: \(error.localizedDescription)")
+                }
+            }
+        }
+    }
+
+    private func isExecuTorchGenerating() -> Bool {
+        execuTorchGenerationLock.lock()
+        defer { execuTorchGenerationLock.unlock() }
+        return execuTorchIsGenerating
+    }
+
+    private func beginExecuTorchGeneration() -> Bool {
+        execuTorchGenerationLock.lock()
+        defer { execuTorchGenerationLock.unlock() }
+
+        if execuTorchIsGenerating {
+            return false
         }
+
+        execuTorchIsGenerating = true
+        return true
+    }
+
+    private func finishExecuTorchGeneration() {
+        execuTorchGenerationLock.lock()
+        execuTorchIsGenerating = false
+        execuTorchGenerationLock.unlock()
     }
 
     @objc func downloadModel(_ call: CAPPluginCall) {
@@ -469,6 +714,165 @@ public class LLMPlugin: CAPPlugin, CAPBridgedPlugin {
     }
 }
 
+private typealias ExecuTorchAllocFunction = @convention(c) (AnyClass, Selector) -> Unmanaged<AnyObject>
+private typealias ExecuTorchInitRunnerFunction = @convention(c) (
+    AnyObject,
+    Selector,
+    NSString,
+    NSString,
+    NSArray
+) -> Unmanaged<AnyObject>
+private typealias ExecuTorchLoadFunction = @convention(c) (
+    AnyObject,
+    Selector,
+    UnsafeMutablePointer<NSError?>
+) -> Bool
+private typealias ExecuTorchConfigBlock = @convention(block) (NSObject) -> Void
+private typealias ExecuTorchInitConfigFunction = @convention(c) (
+    AnyObject,
+    Selector,
+    ExecuTorchConfigBlock
+) -> Unmanaged<AnyObject>
+private typealias ExecuTorchTokenCallback = @convention(block) (NSString) -> Void
+private typealias ExecuTorchGenerateFunction = @convention(c) (
+    AnyObject,
+    Selector,
+    NSString,
+    AnyObject,
+    ExecuTorchTokenCallback,
+    UnsafeMutablePointer<NSError?>
+) -> Bool
+private typealias ExecuTorchVoidFunction = @convention(c) (AnyObject, Selector) -> Void
+
+private final class DynamicExecuTorchRunner {
+    private static let runnerInitSelector = NSSelectorFromString("initWithModelPath:tokenizerPath:specialTokens:")
+    private static let runnerLoadSelector = NSSelectorFromString("loadWithError:")
+    private static let runnerResetSelector = NSSelectorFromString("reset")
+    private static let runnerGenerateSelector = NSSelectorFromString("generateWithPrompt:config:tokenCallback:error:")
+    private static let configInitSelector = NSSelectorFromString("initWithBlock:")
+
+    static var isAvailable: Bool {
+        guard let runnerClass = findClass("ExecuTorchLLMTextRunner"),
+              let configClass = findClass("ExecuTorchLLMConfig") else {
+            return false
+        }
+
+        return hasInstanceMethod(runnerClass, runnerInitSelector) &&
+            hasInstanceMethod(runnerClass, runnerLoadSelector) &&
+            hasInstanceMethod(runnerClass, runnerResetSelector) &&
+            hasInstanceMethod(runnerClass, runnerGenerateSelector) &&
+            hasInstanceMethod(configClass, configInitSelector)
+    }
+
+    private let runner: AnyObject
+    private let sequenceLength: Int
+    private let temperature: Double
+
+    init(
+        modelPath: String,
+        tokenizerPath: String,
+        specialTokens: [String],
+        sequenceLength: Int,
+        temperature: Double
+    ) throws {
+        guard let runnerClass = Self.findClass("ExecuTorchLLMTextRunner") else {
+            throw Self.error("ExecuTorchLLMTextRunner was not found. Link ExecuTorch in the app target.")
+        }
+        try Self.requireInstanceMethod(Self.runnerInitSelector, on: runnerClass)
+        try Self.requireInstanceMethod(Self.runnerLoadSelector, on: runnerClass)
+
+        let alloc: ExecuTorchAllocFunction = try Self.objcMessage()
+        let initRunner: ExecuTorchInitRunnerFunction = try Self.objcMessage()
+        let load: ExecuTorchLoadFunction = try Self.objcMessage()
+
+        let allocatedRunner = alloc(runnerClass, NSSelectorFromString("alloc")).takeRetainedValue()
+        runner = initRunner(
+            allocatedRunner,
+            Self.runnerInitSelector,
+            modelPath as NSString,
+            tokenizerPath as NSString,
+            specialTokens as NSArray
+        ).takeRetainedValue()
+        self.sequenceLength = sequenceLength
+        self.temperature = temperature
+
+        var loadError: NSError?
+        guard load(runner, Self.runnerLoadSelector, &loadError) else {
+            throw loadError ?? Self.error("ExecuTorch model failed to load")
+        }
+    }
+
+    func generate(_ prompt: String, tokenCallback: @escaping (String) -> Void) throws {
+        guard let configClass = Self.findClass("ExecuTorchLLMConfig") else {
+            throw Self.error("ExecuTorchLLMConfig was not found. Link ExecuTorch in the app target.")
+        }
+        guard let runnerClass = object_getClass(runner) else {
+            throw Self.error("ExecuTorch runner class was not found")
+        }
+        try Self.requireInstanceMethod(Self.runnerResetSelector, on: runnerClass)
+        try Self.requireInstanceMethod(Self.runnerGenerateSelector, on: runnerClass)
+        try Self.requireInstanceMethod(Self.configInitSelector, on: configClass)
+
+        let sendVoid: ExecuTorchVoidFunction = try Self.objcMessage()
+        let alloc: ExecuTorchAllocFunction = try Self.objcMessage()
+        let initConfig: ExecuTorchInitConfigFunction = try Self.objcMessage()
+        let generate: ExecuTorchGenerateFunction = try Self.objcMessage()
+
+        sendVoid(runner, Self.runnerResetSelector)
+
+        let allocatedConfig = alloc(configClass, NSSelectorFromString("alloc")).takeRetainedValue()
+        let configBlock: ExecuTorchConfigBlock = { [temperature, sequenceLength] config in
+            config.setValue(temperature, forKey: "temperature")
+            config.setValue(sequenceLength, forKey: "sequenceLength")
+        }
+        let config = initConfig(
+            allocatedConfig,
+            Self.configInitSelector,
+            configBlock
+        ).takeRetainedValue()
+
+        let callback: ExecuTorchTokenCallback = { token in
+            tokenCallback(token as String)
+        }
+        var generateError: NSError?
+        guard generate(
+            runner,
+            Self.runnerGenerateSelector,
+            prompt as NSString,
+            config,
+            callback,
+            &generateError
+        ) else {
+            throw generateError ?? Self.error("ExecuTorch generation failed")
+        }
+    }
+
+    private static func objcMessage<Function>() throws -> Function {
+        guard let handle = dlopen(nil, RTLD_NOW), let symbol = dlsym(handle, "objc_msgSend") else {
+            throw error("Unable to resolve objc_msgSend")
+        }
+        return unsafeBitCast(symbol, to: Function.self)
+    }
+
+    private static func requireInstanceMethod(_ selector: Selector, on cls: AnyClass) throws {
+        guard hasInstanceMethod(cls, selector) else {
+            throw error("ExecuTorch API selector missing: \(NSStringFromSelector(selector))")
+        }
+    }
+
+    private static func hasInstanceMethod(_ cls: AnyClass, _ selector: Selector) -> Bool {
+        class_getInstanceMethod(cls, selector) != nil
+    }
+
+    private static func findClass(_ name: String) -> AnyClass? {
+        NSClassFromString(name) ?? NSClassFromString("ExecuTorchLLM.\(name)")
+    }
+
+    private static func error(_ description: String) -> NSError {
+        NSError(domain: "CapgoLLM.ExecuTorch", code: 1, userInfo: [NSLocalizedDescriptionKey: description])
+    }
+}
+
 // Download delegate class to handle progress updates
 class DownloadDelegate: NSObject, URLSessionDownloadDelegate {
     weak var plugin: CAPPlugin?

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@capgo/capacitor-llm",
-  "version": "8.0.11",
+  "version": "8.0.13",
   "description": "Adds support for LLM locally runned for Capacitor",
   "main": "dist/plugin.cjs.js",
   "module": "dist/esm/index.js",