kangkang/康康/AI/MNNBackend.swift

import Foundation

/// MNN(CPU / SME2)推理后端,封装 `MNNLLMBridge` 的文本流式生成。
/// 与 `LLMSession`/`VLSession` 同款 actor 隔离;跨调用的串行化由上游 `AIRuntime` 闸门保证。
///
/// VL(图→文)需 MNN OMNI 构建(OpenCV 解码图像),当前文本构建不支持;`analyze` 抛错,
/// 上层在 VL 路径回退 MLX(见 `AIRuntime`)。
actor MNNBackend {
    private var bridge: MNNLLMBridge?

    var isLoaded: Bool { bridge?.isLoaded ?? false }

    /// 从 MNN 模型目录加载(目录含 MNN llm 的 config.json + llm.mnn + 权重 + tokenizer)。
    func load(folderURL: URL) throws {
        let configPath = folderURL.appendingPathComponent("config.json").path
        guard FileManager.default.fileExists(atPath: configPath) else {
            throw AIRuntimeError.notReady
        }
        guard let b = MNNLLMBridge(configPath: configPath) else {
            throw AIRuntimeError.modelLoadFailed("MNN createLLM/load 失败")
        }
        bridge = b
    }

    func unload() { bridge = nil }

    /// 文本流式生成。`bridge.generateText` 同步阻塞、逐段回调,放在 detached 线程跑,
    /// 把每段文本 yield 成 `TokenChunk`(含即时 tok/s)。流被取消时调用 `bridge.cancel()`。
    func generate(prompt: String, maxTokens: Int) -> AsyncThrowingStream<TokenChunk, Error> {
        guard let bridge else {
            return AsyncThrowingStream { $0.finish(throwing: AIRuntimeError.notReady) }
        }
        let box = MNNUncheckedBox(bridge)
        return AsyncThrowingStream { continuation in
            let meter = MNNRateMeter()
            let task = Task.detached(priority: .userInitiated) {
                _ = box.value.generateText(prompt, maxTokens: Int32(maxTokens)) { piece in
                    let rate = meter.tick()
                    continuation.yield(TokenChunk(text: piece, decodeRate: rate))
                }
                continuation.finish()
            }
            continuation.onTermination = { _ in
                box.value.cancel()
                task.cancel()
            }
        }
    }

    /// 图→文(VL)。一次性收集(JSON 抽取不需流式)。桥接里把图片路径拼成 <img> 标签,
    /// MNN Omni 内部 imread 加载(需 OMNI 框架);blocking 调用放 detached 线程。
    func analyze(imageURLs: [URL], prompt: String, maxTokens: Int) async throws -> String {
        guard let bridge else { throw AIRuntimeError.notReady }
        let paths = imageURLs.map(\.path)
        let box = MNNUncheckedBox(bridge)
        return try await withCheckedThrowingContinuation { cont in
            Task.detached(priority: .userInitiated) {
                let sink = MNNTextSink()
                do {
                    _ = try box.value.analyzeImages(paths, prompt: prompt, maxTokens: Int32(maxTokens)) { piece in
                        sink.append(piece)
                    }
                    cont.resume(returning: sink.text)
                } catch {
                    cont.resume(throwing: AIRuntimeError.inferenceFailed(error.localizedDescription))
                }
            }
        }
    }
}

/// 单线程串行回调聚合文本,无竞争。
private nonisolated final class MNNTextSink: @unchecked Sendable {
    private(set) var text = ""
    func append(_ s: String) { text += s }
}

/// 把非 Sendable 的 ObjC 桥对象安全带过 detached 边界。
/// 安全性来自 `AIRuntime` 闸门:同一时刻只有一个生成在跑,桥不会被并发访问。
private nonisolated struct MNNUncheckedBox<T>: @unchecked Sendable {
    let value: T
    init(_ value: T) { self.value = value }
}

/// 即时解码速率计:回调在单线程串行调用,内部计数无竞争。
private nonisolated final class MNNRateMeter: @unchecked Sendable {
    private let start = Date()
    private var produced = 0
    func tick() -> Double {
        produced += 1
        let elapsed = Date().timeIntervalSince(start)
        return elapsed > 0 ? Double(produced) / elapsed : 0
    }
}