kangkang/康康/AI/AIRuntime.swift

import Foundation
import MLX

enum AIRuntimeError: Error, LocalizedError {
    case notReady
    case modelLoadFailed(String)
    case inferenceFailed(String)

    var errorDescription: String? {
        switch self {
        case .notReady:                  return String(appLoc: "AI 模型尚未准备好")
        case .modelLoadFailed(let m):    return String(appLoc: "模型加载失败:\(m)")
        case .inferenceFailed(let m):    return String(appLoc: "推理失败:\(m)")
        }
    }
}

actor AIRuntime {
    static let shared = AIRuntime()

    enum Status: Sendable, Equatable {
        case notReady
        case loading
        case ready
        case error(String)
    }

    private(set) var status: Status = .notReady
    private(set) var vlStatus: Status = .notReady
    private(set) var lastDecodeRate: Double = 0

    private var llmSession: LLMSession?
    private var vlSession: VLSession?

    private init() {}

    /// App 启动时调用一次:给 MLX 的 GPU 缓冲池设上限,避免 reuse cache 在大模型常驻之上
    /// 继续膨胀、把峰值推过单 App 内存上限。仅真机生效(模拟器走 CPU,且部分 Metal 路径会 abort)。
    /// 与 increased-memory-limit entitlement + LLM/VL 互斥卸载配合,三管齐下防 jetsam OOM。
    nonisolated static func configureMLXMemory() {
        #if !targetEnvironment(simulator)
        // 256MB cache 上限:够复用、不至于在 3GB 模型之上再囤几百 MB 空闲缓冲。
        MLX.GPU.set(cacheLimit: 256 * 1024 * 1024)
        #endif
    }

    /// 加载模型。首次调用会真正加载,后续幂等。
    func prepare() async throws {
        switch status {
        case .ready:
            return
        case .loading:
            // 已有其他调用方在加载;本次 prepare 直接返回,
            // 调用方需稍后 await prepare() 再判 status,或自行轮询 / 显示加载 UI。
            // W3 引入 prepare 队列时优化。
            return
        case .error, .notReady:
            break
        }

        guard ModelStore.shared.isReady(.llm) else {
            status = .error("LLM 模型未就绪")
            throw AIRuntimeError.notReady
        }

        // OOM 闸门(§3.1):LLM(~1GB)与 VL(~3GB)不可同时常驻,叠加会冲过单 App 内存上限被 jetsam 杀。
        // 加载 LLM 前先卸 VL,释放其 ModelContainer + MLX 显存缓存。
        unloadVL()

        status = .loading
        do {
            let session = try await LLMSession.load(
                folderURL: ModelStore.shared.localURL(for: .llm)
            )
            self.llmSession = session
            status = .ready
        } catch {
            status = .error("\(error)")
            throw AIRuntimeError.modelLoadFailed("\(error)")
        }
    }

    /// 流式生成。调用前应先 await prepare()。
    /// 注意:返回流是同步创建的,但跨 actor 调用 LLMSession 需要 await。
    func generate(prompt: String, maxTokens: Int = 256) -> AsyncThrowingStream<TokenChunk, Error> {
        // 在 actor 隔离上下文中捕获快照,Task 内不再访问 self.status / self.llmSession
        let snapshotStatus = status
        let snapshotSession = llmSession

        return AsyncThrowingStream { continuation in
            let task = Task {
                guard snapshotStatus == .ready, let session = snapshotSession else {
                    continuation.finish(throwing: AIRuntimeError.notReady)
                    return
                }
                do {
                    // session.generate 跨 actor 边界,需要 await
                    let stream = await session.generate(prompt: prompt, maxTokens: maxTokens)
                    for try await chunk in stream {
                        // 消费者(UI)提前关闭/取消时,下面的 checkCancellation 让本 Task 尽快退出,
                        // 连带丢弃 session 流并触发其 onTermination,停止底层 MLX 解码,不空耗 GPU。
                        try Task.checkCancellation()
                        // Task 闭包在 generate() 内启动,继承 AIRuntime 的 actor 隔离;
                        // 调用同 actor 的 recordRate 不需要 await
                        self.recordRate(chunk.decodeRate)
                        continuation.yield(chunk)
                    }
                    continuation.finish()
                } catch {
                    continuation.finish(throwing: AIRuntimeError.inferenceFailed("\(error)"))
                }
            }
            // 消费者取消/流终止时取消内部 Task(与 LLMSession / HealthExportService 一致)。
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    private func recordRate(_ rate: Double) {
        if rate > 0 { lastDecodeRate = rate }
    }

    // MARK: - VL

    /// 加载 VL 模型。幂等,首调真正 load。
    func prepareVL() async throws {
        switch vlStatus {
        case .ready, .loading:
            return
        case .error, .notReady:
            break
        }

        guard ModelStore.shared.isReady(.vl) else {
            vlStatus = .error("VL 模型未就绪")
            throw AIRuntimeError.notReady
        }

        // OOM 闸门(§3.1):加载 VL(~3GB)前先卸 LLM(~1GB),否则两者常驻叠加冲过内存上限被 jetsam 杀
        // —— 这正是「异常项快拍识别时 App 自动退出」的主因。
        unloadLLM()

        vlStatus = .loading
        do {
            let session = try await VLSession.load(
                folderURL: ModelStore.shared.localURL(for: .vl)
            )
            self.vlSession = session
            vlStatus = .ready
        } catch {
            vlStatus = .error("\(error)")
            throw AIRuntimeError.modelLoadFailed("\(error)")
        }
    }

    // MARK: - 卸载(OOM 闸门)

    /// 卸载 LLM,释放 ModelContainer 引用并清 MLX 显存缓存。幂等。
    /// 注:若此刻有 generate() 的流仍在跑,它持有 session 快照,真正释放要等流结束;
    /// 但快拍/归档场景下没有并发文本流,卸载即时生效。
    private func unloadLLM() {
        guard llmSession != nil else { return }
        llmSession = nil
        status = .notReady
        MLX.GPU.clearCache()
    }

    /// 卸载 VL,释放 ModelContainer 引用并清 MLX 显存缓存。幂等。
    private func unloadVL() {
        guard vlSession != nil else { return }
        vlSession = nil
        vlStatus = .notReady
        MLX.GPU.clearCache()
    }

    /// 图像 → JSON 字符串(由 VLPrompts.reportExtraction 引导)。
    /// 调用方负责解析 + 失败回退(§3.2)。
    /// AIRuntime 是 actor,本调用与 LLM.generate() 自然串行,不会 OOM。
    func analyzeReport(imageURLs: [URL],
                       prompt: String,
                       maxTokens: Int = 512) async throws -> String {
        guard vlStatus == .ready, let session = vlSession else {
            throw AIRuntimeError.notReady
        }
        do {
            return try await session.analyze(
                imageURLs: imageURLs,
                prompt: prompt,
                maxTokens: maxTokens
            )
        } catch {
            throw AIRuntimeError.inferenceFailed("\(error)")
        }
    }
}