@swiftbrowser/native 0.18.2 → 0.19.0

package/swift/Packages/SwiftBrowserLocalModel/Sources/SwiftBrowserLocalModel/LocalLanguageModel.swift

@@ -0,0 +1,409 @@
+import Foundation
+import SwiftBrowserNative
+
+/// Open language models run on the device with llama.cpp, for the app's
+/// `LanguageModelSession`s where Apple Intelligence cannot answer (or where
+/// the viewer chooses one). Install one, or several to choose from, at launch:
+///
+/// ```swift
+/// LocalLanguageModel.install([.llama3_2_1B, .qwen2_5_1_5B, .llama3_2_3B])
+/// ```
+///
+/// The first is the default; the viewer's choice (`selectModel` in
+/// @swiftbrowser/react) is remembered across launches.
+///
+/// The model is a GGUF file the app fetches when the viewer asks for it
+/// (`SystemLanguageModel.default.download()` in @swiftbrowser/react, after
+/// showing its size; until it is here availability is `modelNotReady`, with
+/// the download's state), kept in Application Support and excluded from
+/// backups, loaded on the first prompt or `prewarm()`, and removed with
+/// `deleteDownload()`.
+public enum LocalLanguageModel {
+    /// A GGUF model to fetch and run.
+    public struct Model: Sendable, Equatable {
+        /// How the app names it to the viewer.
+        public var name: String
+        /// A line about it: what it is good at.
+        public var summary: String?
+        /// The file's name on the device, which is also the model's id.
+        public var fileName: String
+        public var url: URL
+        /// The file's size, checked once it has downloaded; also the download's
+        /// progress where the server does not say how big the file is.
+        public var bytes: Int64?
+        /// How many tokens of the conversation the model keeps in mind.
+        public var contextLength: Int
+        /// Layers run on the GPU with Metal (all of them by default); 0 runs on the CPU.
+        public var gpuLayers: Int32
+
+        public init(name: String, summary: String? = nil, fileName: String, url: URL, bytes: Int64? = nil, contextLength: Int = 4096, gpuLayers: Int32 = 999) {
+            self.name = name
+            self.summary = summary
+            self.fileName = fileName
+            self.url = url
+            self.bytes = bytes
+            self.contextLength = contextLength
+            self.gpuLayers = gpuLayers
+        }
+
+        /// Meta's Llama 3.2 1B Instruct, 4-bit (Q4_K_M): 808 MB, at home on a
+        /// phone with 6 GB of memory. Built with Llama; Meta's Llama 3.2
+        /// Community License applies.
+        public static let llama3_2_1B = Model(
+            name: "Llama 3.2 1B",
+            summary: "Fastest",
+            fileName: "Llama-3.2-1B-Instruct-Q4_K_M.gguf",
+            url: URL(string: "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q4_K_M.gguf")!,
+            bytes: 807_694_464
+        )
+
+        /// Alibaba's Qwen 2.5 1.5B Instruct, 4-bit (Q4_K_M): 1.1 GB, the
+        /// steadiest of these with tools. Apache 2.0.
+        public static let qwen2_5_1_5B = Model(
+            name: "Qwen 2.5 1.5B",
+            summary: "Best with tools",
+            fileName: "qwen2.5-1.5b-instruct-q4_k_m.gguf",
+            url: URL(string: "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF/resolve/main/qwen2.5-1.5b-instruct-q4_k_m.gguf")!,
+            bytes: 1_117_320_736
+        )
+
+        /// Meta's Llama 3.2 3B Instruct, 4-bit (Q4_K_M): 2.0 GB, the most
+        /// capable here, and the most memory (a phone with 6 GB runs it close
+        /// to its limit). Built with Llama; Meta's Llama 3.2 Community License applies.
+        public static let llama3_2_3B = Model(
+            name: "Llama 3.2 3B",
+            summary: "Most capable",
+            fileName: "Llama-3.2-3B-Instruct-Q4_K_M.gguf",
+            url: URL(string: "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q4_K_M.gguf")!,
+            bytes: 2_019_377_696
+        )
+    }
+
+    /// Makes `models` the app's `LanguageModels.fallback`, the first the
+    /// default. `directory` is where the files are kept (Application
+    /// Support's SwiftBrowserLocalModel folder by default);
+    /// `allowsCellularAccess` false waits for Wi-Fi to fetch them.
+    @MainActor
+    public static func install(_ models: [Model], directory: URL? = nil, allowsCellularAccess: Bool = true) {
+        LanguageModels.fallback = provider(models, directory: directory, allowsCellularAccess: allowsCellularAccess)
+    }
+
+    @MainActor
+    public static func install(_ model: Model, directory: URL? = nil, allowsCellularAccess: Bool = true) {
+        install([model], directory: directory, allowsCellularAccess: allowsCellularAccess)
+    }
+
+    /// The provider `install` sets, to use some other way. `remembers` keeps
+    /// the viewer's choice across launches (in the standard user defaults).
+    @MainActor
+    public static func provider(
+        _ models: [Model],
+        directory: URL? = nil,
+        allowsCellularAccess: Bool = true,
+        remembers: Bool = true
+    ) -> any LanguageModelProvider {
+        precondition(!models.isEmpty, "LocalLanguageModel: no models")
+        let directory = directory ?? URL.applicationSupportDirectory.appendingPathComponent("SwiftBrowserLocalModel", isDirectory: true)
+        return LlamaProvider(
+            models: models,
+            stores: models.map { ModelStore(model: $0, directory: directory, allowsCellularAccess: allowsCellularAccess) },
+            remembers: remembers
+        )
+    }
+
+    @MainActor
+    public static func provider(_ model: Model, directory: URL? = nil, allowsCellularAccess: Bool = true) -> any LanguageModelProvider {
+        provider([model], directory: directory, allowsCellularAccess: allowsCellularAccess, remembers: false)
+    }
+}
+
+/// The model's file: on the device, or on its way.
+@MainActor
+final class ModelStore {
+    let model: LocalLanguageModel.Model
+    let file: URL
+    private let directory: URL
+    private let session: URLSession
+    private var download: Task<Void, Never>?
+    private var watcher: DownloadWatcher?
+    /// Why the last download failed; the next request for availability tries again.
+    private(set) var failure: Error?
+
+    init(model: LocalLanguageModel.Model, directory: URL, allowsCellularAccess: Bool) {
+        self.model = model
+        self.directory = directory
+        file = directory.appendingPathComponent(model.fileName)
+        let configuration = URLSessionConfiguration.default
+        configuration.allowsCellularAccess = allowsCellularAccess
+        configuration.timeoutIntervalForResource = 6 * 60 * 60
+        session = URLSession(configuration: configuration)
+    }
+
+    var isReady: Bool {
+        FileManager.default.fileExists(atPath: file.path)
+    }
+
+    var isFetching: Bool {
+        download != nil
+    }
+
+    /// How much of the file has come, from 0 to 1.
+    var progress: Double {
+        guard let task = watcher?.task else { return 0 }
+        let expected = task.countOfBytesExpectedToReceive > 0 ? task.countOfBytesExpectedToReceive : model.bytes ?? 0
+        guard expected > 0 else { return 0 }
+        return min(1, Double(task.countOfBytesReceived) / Double(expected))
+    }
+
+    /// Starts fetching the file, unless it is here or on its way.
+    func fetch() {
+        guard !isReady, download == nil else { return }
+        failure = nil
+        let watcher = DownloadWatcher()
+        self.watcher = watcher
+        download = Task { @MainActor in
+            do {
+                try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
+                let (temporary, response) = try await session.download(from: model.url, delegate: watcher)
+                if let http = response as? HTTPURLResponse, !(200..<300).contains(http.statusCode) {
+                    throw LanguageModelError(.unavailable, "the model's download answered \(http.statusCode)")
+                }
+                let size = (try FileManager.default.attributesOfItem(atPath: temporary.path)[.size] as? NSNumber)?.int64Value
+                if let bytes = model.bytes, size != bytes {
+                    try? FileManager.default.removeItem(at: temporary)
+                    throw LanguageModelError(.unavailable, "the model's download was \(size ?? 0) bytes, not \(bytes)")
+                }
+                try? FileManager.default.removeItem(at: file)
+                try FileManager.default.moveItem(at: temporary, to: file)
+                var values = URLResourceValues()
+                values.isExcludedFromBackup = true
+                var kept = file
+                try? kept.setResourceValues(values)
+                failure = nil
+            } catch {
+                if !Task.isCancelled { failure = error }
+            }
+            download = nil
+            self.watcher = nil
+        }
+    }
+
+    /// Stops a download under way and removes the file.
+    func remove() async {
+        if let download {
+            download.cancel()
+            await download.value
+        }
+        failure = nil
+        try? FileManager.default.removeItem(at: file)
+    }
+}
+
+/// Hears which task a download made, so its progress can be read.
+private final class DownloadWatcher: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
+    private let lock = NSLock()
+    private var current: URLSessionTask?
+
+    var task: URLSessionTask? {
+        lock.withLock { current }
+    }
+
+    func urlSession(_ session: URLSession, didCreateTask task: URLSessionTask) {
+        lock.withLock { current = task }
+    }
+}
+
+/// The models as `LanguageModels.fallback`: the chosen one available once
+/// its file is here, loaded once for every session.
+@MainActor
+final class LlamaProvider: LanguageModelProvider {
+    let models: [LocalLanguageModel.Model]
+    private let stores: [String: ModelStore]
+    private let remembers: Bool
+    private static let selectionKey = "SwiftBrowserLocalModel.selected"
+    private(set) var selected: String
+    /// The loaded model and which it is.
+    private var engine: (id: String, task: Task<LlamaEngine, Error>)?
+
+    init(models: [LocalLanguageModel.Model], stores: [ModelStore], remembers: Bool) {
+        self.models = models
+        self.stores = Dictionary(uniqueKeysWithValues: stores.map { ($0.model.fileName, $0) })
+        self.remembers = remembers
+        let remembered = remembers ? UserDefaults.standard.string(forKey: Self.selectionKey) : nil
+        selected = models.first { $0.fileName == remembered }?.fileName ?? models[0].fileName
+    }
+
+    /// The chosen model's file.
+    var store: ModelStore { stores[selected]! }
+
+    func availability() async -> LanguageModelAvailability {
+        store.isReady ? .available : .unavailable(.modelNotReady)
+    }
+
+    var downloadState: LanguageModelDownload? {
+        Self.state(of: store)
+    }
+
+    var choices: [LanguageModelChoice] {
+        models.map { LanguageModelChoice(id: $0.fileName, name: $0.name, summary: $0.summary, download: Self.state(of: stores[$0.fileName]!)) }
+    }
+
+    var selectedChoice: String? { selected }
+
+    func select(_ id: String) {
+        guard stores[id] != nil, id != selected else { return }
+        selected = id
+        if remembers { UserDefaults.standard.set(id, forKey: Self.selectionKey) }
+        // The model chosen before goes once a response under way lets it go.
+        engine = nil
+    }
+
+    private static func state(of store: ModelStore) -> LanguageModelDownload {
+        LanguageModelDownload(
+            size: store.model.bytes,
+            isDownloaded: store.isReady,
+            progress: store.isFetching ? store.progress : nil,
+            failure: store.failure.map { ($0 as? LanguageModelError)?.message ?? $0.localizedDescription }
+        )
+    }
+
+    func startDownload() {
+        store.fetch()
+    }
+
+    func deleteDownload() async throws {
+        // The loaded model goes once a response under way has let it go.
+        if engine?.id == selected { engine = nil }
+        await store.remove()
+    }
+
+    func makeSession(instructions: String?, tools: [LanguageModelTool]) throws -> any LanguageModelProviderSession {
+        LlamaSession(provider: self, instructions: instructions, tools: tools)
+    }
+
+    /// The chosen model, loading it the first time (off the main thread).
+    func loaded() async throws -> LlamaEngine {
+        if let engine, engine.id == selected { return try await engine.task.value }
+        let id = selected
+        let path = store.file.path
+        let model = store.model
+        let loading = Task.detached(priority: .userInitiated) {
+            try LlamaEngine(path: path, contextLength: model.contextLength, gpuLayers: model.gpuLayers)
+        }
+        engine = (id, loading)
+        do {
+            return try await loading.value
+        } catch {
+            if engine?.id == id { engine = nil }
+            throw error
+        }
+    }
+}
+
+/// One conversation: its messages, each prompt answered after the ones
+/// before. Tools are offered in the instructions: the model calls one by
+/// answering with just `{"tool": name, "arguments": {...}}`, the session runs
+/// it and hands the model the result, then lets it answer (a few rounds at
+/// most); a reply that starts as such an object is held back until it is
+/// clear whether it is a call.
+@MainActor
+final class LlamaSession: LanguageModelProviderSession {
+    private let provider: LlamaProvider
+    private let tools: [LanguageModelTool]
+    private(set) var messages: [LlamaEngine.Message]
+    /// How many tool calls one prompt may make before the model must answer.
+    nonisolated static let toolRounds = 4
+
+    init(provider: LlamaProvider, instructions: String?, tools: [LanguageModelTool]) {
+        self.provider = provider
+        self.tools = tools
+        let system = Self.systemPrompt(instructions: instructions, tools: tools)
+        messages = system.map { [LlamaEngine.Message(role: "system", content: $0)] } ?? []
+    }
+
+    func prewarm() {
+        let provider = provider
+        Task { _ = try? await provider.loaded() }
+    }
+
+    func streamResponse(to prompt: String, options: LanguageModelOptions) -> AsyncThrowingStream<String, Error> {
+        let provider = provider
+        let tools = tools
+        let asked = messages + [LlamaEngine.Message(role: "user", content: prompt)]
+        return AsyncThrowingStream { continuation in
+            let cancelled = CancelFlag()
+            let task = Task { @MainActor [weak self] in
+                do {
+                    let engine = try await provider.loaded()
+                    var conversation = asked
+                    var reply = ""
+                    for round in 0...Self.toolRounds {
+                        let offersTools = !tools.isEmpty && round < Self.toolRounds
+                        reply = try await engine.generate(conversation, options: options, cancelled: cancelled) { text in
+                            // What may be a tool call is not shown.
+                            if !offersTools || !Self.mayBeCall(text) { continuation.yield(text) }
+                        }
+                        guard offersTools, let call = Self.toolCall(in: reply), let tool = tools.first(where: { $0.name == call.name }) else { break }
+                        let output: String
+                        do {
+                            output = try await tool.call(call.arguments)
+                        } catch {
+                            output = "The tool failed: \(error)"
+                        }
+                        conversation += [
+                            LlamaEngine.Message(role: "assistant", content: reply),
+                            LlamaEngine.Message(role: engine.toolRole, content: output),
+                        ]
+                        reply = ""
+                    }
+                    if !tools.isEmpty, Self.mayBeCall(reply) { continuation.yield(reply) }
+                    // What the viewer saw is what the conversation goes on from, a stopped reply included.
+                    if !reply.isEmpty { self?.messages = conversation + [LlamaEngine.Message(role: "assistant", content: reply)] }
+                    continuation.finish()
+                } catch {
+                    continuation.finish(throwing: error)
+                }
+            }
+            continuation.onTermination = { _ in
+                cancelled.cancel()
+                task.cancel()
+            }
+        }
+    }
+
+    func close() {}
+
+    /// The instructions, with the tools described after them.
+    nonisolated static func systemPrompt(instructions: String?, tools: [LanguageModelTool]) -> String? {
+        guard !tools.isEmpty else { return instructions }
+        let list = tools.map { tool in
+            "- \(tool.name): \(tool.description) Arguments (JSON schema): \(tool.parameters.json)"
+        }.joined(separator: "\n")
+        let guide = """
+        You can use these tools:
+        \(list)
+        To use a tool, reply with only a JSON object and nothing else: {"tool": "<name>", "arguments": {...}}. \
+        You will then get the tool's result and can answer. Use a tool only when it helps; otherwise answer directly.
+        """
+        return [instructions, guide].compactMap { $0 }.joined(separator: "\n\n")
+    }
+
+    /// Whether a reply so far could still turn out to be a tool call.
+    nonisolated static func mayBeCall(_ text: String) -> Bool {
+        let trimmed = text.drop(while: \.isWhitespace)
+        return trimmed.hasPrefix("{") || trimmed.hasPrefix("```")
+    }
+
+    /// The tool call a reply makes, if it is one: `{"tool"|"name": …, "arguments"|"parameters": {…}}`,
+    /// alone or in a code fence.
+    nonisolated static func toolCall(in reply: String) -> (name: String, arguments: JSONValue)? {
+        var text = reply.trimmingCharacters(in: .whitespacesAndNewlines)
+        if text.hasPrefix("```") {
+            text = text.drop(while: { $0 != "\n" }).trimmingCharacters(in: .whitespacesAndNewlines)
+            if text.hasSuffix("```") { text = String(text.dropLast(3)).trimmingCharacters(in: .whitespacesAndNewlines) }
+        }
+        guard text.hasPrefix("{"), let object = (try? JSONValue.parse(text))?.object else { return nil }
+        guard let name = (object["tool"] ?? object["name"])?.string else { return nil }
+        return (name, object["arguments"] ?? object["parameters"] ?? .object([:]))
+    }
+}