spb/zyquo-local Public MIT

1//2//  ModelStore.swift3//  Zyquo Local4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import AppKit10import Foundation11import Observation1213/// The on-disk model library: scans the Models folder, validates model14/// directories, tracks sizes and last-used dates, deletes models.15@MainActor16@Observable17final class ModelStore {18    private(set) var models: [LocalModel] = []1920    /// Root: ~/Library/Application Support/ZyquoLocal/Models21    let modelsRoot: URL2223    /// Sidecar metadata (last used, per-model params) keyed by repoID.24    private var meta: [String: ModelMeta] = [:]2526    struct ModelMeta: Codable {27        var lastUsed: Date?28        var defaultParams: GenerationParams?29    }3031    var totalSizeBytes: Int64 {32        models.reduce(0) { $0 + $1.sizeBytes }33    }3435    init(root: URL = PersistenceService.appSupportDirectory.appendingPathComponent("Models")) {36        self.modelsRoot = root37        try? FileManager.default.createDirectory(at: root, withIntermediateDirectories: true)38        loadMeta()39        rescan()40    }4142    /// Directory a repo downloads into: Models/{org}/{repo}.43    func directory(for repoID: String) -> URL {44        modelsRoot.appendingPathComponent(repoID)45    }4647    func model(for repoID: String) -> LocalModel? {48        models.first { $0.repoID == repoID }49    }5051    /// Required files for a valid model directory (docs/MLX-RESEARCH.md §3).52    static func isValidModelDirectory(_ dir: URL) -> Bool {53        let fm = FileManager.default54        guard fm.fileExists(atPath: dir.appendingPathComponent("config.json").path) else { return false }55        guard fm.fileExists(atPath: dir.appendingPathComponent("tokenizer_config.json").path) else { return false }56        guard fm.fileExists(atPath: dir.appendingPathComponent("tokenizer.json").path) else { return false }57        guard let contents = try? fm.contentsOfDirectory(atPath: dir.path) else { return false }58        return contents.contains { $0.hasSuffix(".safetensors") && !$0.hasSuffix(".partial") }59    }6061    /// Scans Models/{org}/{repo} two levels deep and rebuilds the library.62    func rescan() {63        let fm = FileManager.default64        var found: [LocalModel] = []65        let orgs = (try? fm.contentsOfDirectory(66            at: modelsRoot, includingPropertiesForKeys: nil,67            options: .skipsHiddenFiles)) ?? []68        for org in orgs where org.hasDirectoryPath || (try? org.resourceValues(forKeys: [.isDirectoryKey]).isDirectory) == true {69            let repos = (try? fm.contentsOfDirectory(70                at: org, includingPropertiesForKeys: nil,71                options: .skipsHiddenFiles)) ?? []72            for repo in repos {73                guard (try? repo.resourceValues(forKeys: [.isDirectoryKey]).isDirectory) == true else { continue }74                guard Self.isValidModelDirectory(repo) else { continue }75                let repoID = "\(org.lastPathComponent)/\(repo.lastPathComponent)"76                let size = (try? fm.allocatedSizeOfDirectory(at: repo)) ?? 077                let config = Self.readConfig(directory: repo)78                found.append(LocalModel(79                    repoID: repoID,80                    directory: repo,81                    sizeBytes: size,82                    architecture: config.architecture,83                    quantization: Self.quantLabel(repoID: repoID, config: config),84                    contextWindow: InferenceEngine.readContextWindow(directory: repo),85                    lastUsed: meta[repoID]?.lastUsed,86                    defaultParams: meta[repoID]?.defaultParams87                ))88            }89        }90        models = found.sorted { ($0.lastUsed ?? .distantPast) > ($1.lastUsed ?? .distantPast) }91    }9293    /// Deletes a model directory. Returns the reclaimed bytes.94    @discardableResult95    func delete(repoID: String) -> Int64 {96        guard let model = model(for: repoID) else { return 0 }97        let reclaimed = model.sizeBytes98        try? FileManager.default.removeItem(at: model.directory)99        // Prune the org folder when empty.100        let org = model.directory.deletingLastPathComponent()101        if let rest = try? FileManager.default.contentsOfDirectory(atPath: org.path), rest.isEmpty {102            try? FileManager.default.removeItem(at: org)103        }104        meta[repoID] = nil105        saveMeta()106        rescan()107        return reclaimed108    }109110    func revealInFinder(repoID: String) {111        guard let model = model(for: repoID) else { return }112        NSWorkspace.shared.activateFileViewerSelecting([model.directory])113    }114115    func markUsed(repoID: String) {116        var m = meta[repoID] ?? ModelMeta()117        m.lastUsed = Date()118        meta[repoID] = m119        saveMeta()120        if let i = models.firstIndex(where: { $0.repoID == repoID }) {121            models[i].lastUsed = m.lastUsed122        }123    }124125    func setDefaultParams(_ params: GenerationParams?, for repoID: String) {126        var m = meta[repoID] ?? ModelMeta()127        m.defaultParams = params128        meta[repoID] = m129        saveMeta()130        if let i = models.firstIndex(where: { $0.repoID == repoID }) {131            models[i].defaultParams = params132        }133    }134135    // MARK: - Config parsing136137    struct ConfigInfo {138        var architecture: String?139        var quantBits: Int?140    }141142    static func readConfig(directory: URL) -> ConfigInfo {143        let url = directory.appendingPathComponent("config.json")144        guard let data = try? Data(contentsOf: url),145            let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]146        else { return ConfigInfo() }147        var bits: Int?148        if let quant = json["quantization"] as? [String: Any] {149            bits = quant["bits"] as? Int150        }151        return ConfigInfo(architecture: json["model_type"] as? String, quantBits: bits)152    }153154    /// "4bit", "8bit", "bf16"… parsed from repo name, falling back to config.155    static func quantLabel(repoID: String, config: ConfigInfo) -> String? {156        let name = repoID.lowercased()157        for candidate in ["4bit", "5bit", "6bit", "8bit", "bf16", "fp16", "mxfp4"] {158            if name.contains(candidate) { return candidate }159        }160        if let bits = config.quantBits { return "\(bits)bit" }161        return nil162    }163164    // MARK: - Meta persistence165166    private var metaURL: URL {167        PersistenceService.appSupportDirectory.appendingPathComponent("library-meta.json")168    }169170    private func loadMeta() {171        guard let data = try? Data(contentsOf: metaURL),172            let decoded = try? JSONDecoder().decode([String: ModelMeta].self, from: data)173        else { return }174        meta = decoded175    }176177    private func saveMeta() {178        try? JSONEncoder().encode(meta).write(to: metaURL, options: .atomic)179    }180}181