spb/zyquo-local Public MIT

1//2//  InferenceEngine.swift3//  Zyquo Local4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation10import MLX11import MLXHuggingFace12import MLXLLM13import MLXLMCommon14import Tokenizers1516/// Events streamed by the engine during generation.17enum GenerationEvent: Sendable {18    case token(String)19    case stats(GenerationStats)20    case finished(GenerationFinishReason)21}2223enum GenerationFinishReason: String, Sendable {24    case stop25    case length26    case cancelled27}2829/// Final per-response statistics.30struct GenerationStats: Sendable {31    var timeToFirstToken: TimeInterval32    var tokensPerSecond: Double33    var promptTokenCount: Int34    var generationTokenCount: Int35    var peakMemoryBytes: Int36}3738/// Engine lifecycle state.39enum EngineState: Sendable, Equatable {40    case unloaded41    case loading(repoID: String)42    case ready(repoID: String)43    case generating(repoID: String)44}4546/// Errors surfaced by the engine, with human-readable messages.47enum EngineError: LocalizedError {48    case noModelLoaded49    case modelDirectoryInvalid(URL)50    case loadFailed(String, underlying: String)51    case unsupportedArchitecture(String)52    case outOfMemory(String)5354    var errorDescription: String? {55        switch self {56        case .noModelLoaded:57            "No model is loaded. Load a model from the Library first."58        case .modelDirectoryInvalid(let url):59            "The model folder is missing required files (config.json, weights, tokenizer): \(url.path)"60        case .loadFailed(let repo, let underlying):61            "Could not load \(repo): \(underlying)"62        case .unsupportedArchitecture(let arch):63            "This model's architecture “\(arch)” is not supported by the MLX engine yet."64        case .outOfMemory(let repo):65            "Not enough memory to load \(repo). Try a smaller quantization (e.g. 4-bit) or a smaller model."66        }67    }68}6970/// The single owner of all inference. One model loaded at a time (v1).71/// States: unloaded → loading → ready ⇄ generating.72actor InferenceEngine {73    private(set) var state: EngineState = .unloaded7475    private var container: ModelContainer?76    private var loadedModel: LocalModel?77    private var session: ChatSession?78    private var generationTask: Task<Void, Never>?7980    /// Context window of the loaded model (from config.json), defaulting81    /// conservatively when absent.82    private(set) var contextWindow: Int = 40968384    // MARK: - Load / unload8586    /// Loads a model from its local directory, replacing any loaded model.87    func load(model: LocalModel) async throws {88        await stopGeneration()89        unloadInternal()90        state = .loading(repoID: model.repoID)91        do {92            guard FileManager.default.fileExists(93                atPath: model.directory.appendingPathComponent("config.json").path)94            else {95                throw EngineError.modelDirectoryInvalid(model.directory)96            }97            contextWindow = Self.readContextWindow(directory: model.directory) ?? 409698            let container = try await loadModelContainer(99                from: model.directory,100                using: #huggingFaceTokenizerLoader()101            )102            self.container = container103            self.loadedModel = model104            state = .ready(repoID: model.repoID)105        } catch let error as EngineError {106            state = .unloaded107            throw error108        } catch {109            state = .unloaded110            let text = String(describing: error)111            if text.localizedCaseInsensitiveContains("Unsupported model type") {112                throw EngineError.unsupportedArchitecture(model.architecture ?? "unknown")113            }114            if text.localizedCaseInsensitiveContains("memory") {115                throw EngineError.outOfMemory(model.repoID)116            }117            throw EngineError.loadFailed(model.repoID, underlying: error.localizedDescription)118        }119    }120121    /// Unloads the current model and verifiably frees memory.122    func unload() async {123        await stopGeneration()124        unloadInternal()125    }126127    private func unloadInternal() {128        session = nil129        container = nil130        loadedModel = nil131        MemoryAdvisor.reclaimMemory()132        state = .unloaded133    }134135    var currentModel: LocalModel? { loadedModel }136137    // MARK: - Sessions138139    /// Binds (or rebinds) the engine to a conversation. Rebuilding the session140    /// drops the KV cache, so this is only done when switching conversations,141    /// or when the system prompt changed.142    func startSession(conversation: Conversation) throws {143        guard let container else { throw EngineError.noModelLoaded }144        session = ChatSession(145            container: container,146            conversation: conversation,147            contextWindow: contextWindow,148            params: conversation.params149        )150    }151152    /// Ensures the active session matches the conversation, preserving the153    /// KV cache when it does.154    func ensureSession(conversation: Conversation) throws {155        if session?.conversationID != conversation.id {156            try startSession(conversation: conversation)157        }158    }159160    var contextUsage: (used: Int, window: Int)? {161        session?.estimatedContextUsage162    }163164    // MARK: - Generation165166    /// Streams a response to `prompt` inside the bound session.167    /// Cancellation: cancel the consuming task or call `stopGeneration()` —168    /// the underlying generation loop actually stops.169    func generate(prompt: String, params: GenerationParams) throws170        -> AsyncThrowingStream<GenerationEvent, Error>171    {172        guard let session else { throw EngineError.noModelLoaded }173        guard let repoID = loadedModel?.repoID else { throw EngineError.noModelLoaded }174175        let (stream, continuation) = AsyncThrowingStream.makeStream(of: GenerationEvent.self)176        state = .generating(repoID: repoID)177        MemoryAdvisor.resetPeakMemory()178179        let task = Task {180            let start = ContinuousClock.now181            var firstTokenTime: TimeInterval?182            var response = ""183            var finished: GenerationFinishReason = .cancelled184            do {185                for try await generation in session.stream(prompt: prompt, params: params) {186                    if Task.isCancelled { break }187                    switch generation {188                    case .chunk(let text):189                        if firstTokenTime == nil {190                            firstTokenTime = Self.seconds(since: start)191                        }192                        response += text193                        continuation.yield(.token(text))194                    case .info(let info):195                        let stats = GenerationStats(196                            timeToFirstToken: firstTokenTime ?? info.promptTime,197                            tokensPerSecond: info.tokensPerSecond,198                            promptTokenCount: info.promptTokenCount,199                            generationTokenCount: info.generationTokenCount,200                            peakMemoryBytes: MemoryAdvisor.peakMemoryBytes201                        )202                        finished = switch info.stopReason {203                        case .stop: .stop204                        case .length: .length205                        case .cancelled: .cancelled206                        }207                        continuation.yield(.stats(stats))208                    case .toolCall:209                        break210                    }211                }212                session.noteResponse(response)213                continuation.yield(.finished(Task.isCancelled ? .cancelled : finished))214                continuation.finish()215            } catch {216                continuation.finish(throwing: error)217            }218            self.finishGeneration()219        }220        generationTask = task221        continuation.onTermination = { termination in222            if case .cancelled = termination { task.cancel() }223        }224        return stream225    }226227    /// Stops any in-flight generation and waits for it to wind down.228    func stopGeneration() async {229        generationTask?.cancel()230        _ = await generationTask?.value231        generationTask = nil232    }233234    private func finishGeneration() {235        if let repoID = loadedModel?.repoID {236            state = .ready(repoID: repoID)237        }238        generationTask = nil239    }240241    // MARK: - Helpers242243    private static func seconds(since start: ContinuousClock.Instant) -> TimeInterval {244        let duration = start.duration(to: .now)245        return Double(duration.components.seconds)246            + Double(duration.components.attoseconds) * 1e-18247    }248249    /// Reads max_position_embeddings from config.json.250    static func readContextWindow(directory: URL) -> Int? {251        let url = directory.appendingPathComponent("config.json")252        guard let data = try? Data(contentsOf: url),253            let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]254        else { return nil }255        if let value = json["max_position_embeddings"] as? Int { return value }256        if let text = json["text_config"] as? [String: Any],257            let value = text["max_position_embeddings"] as? Int258        {259            return value260        }261        return nil262    }263}264