526 lines
18 KiB
Swift
Executable File
526 lines
18 KiB
Swift
Executable File
// MARK: – ModelDownloadFeature.swift
|
||
|
||
// A full‐featured TCA reducer for managing on‑device ML models.
|
||
// Dependencies: ComposableArchitecture, IdentifiedCollections, Dependencies
|
||
|
||
import AppKit
|
||
import ComposableArchitecture
|
||
import Dependencies
|
||
import HexCore
|
||
import IdentifiedCollections
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────
|
||
|
||
// MARK: – Data Models
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────
|
||
|
||
public struct ModelInfo: Equatable, Identifiable {
|
||
public let name: String
|
||
public var isDownloaded: Bool
|
||
|
||
public var id: String { name }
|
||
public init(name: String, isDownloaded: Bool) {
|
||
self.name = name
|
||
self.isDownloaded = isDownloaded
|
||
}
|
||
}
|
||
|
||
public struct CuratedModelInfo: Equatable, Identifiable, Codable {
|
||
public let displayName: String
|
||
public let internalName: String
|
||
public let size: String
|
||
public let accuracyStars: Int
|
||
public let speedStars: Int
|
||
public let storageSize: String
|
||
public var isDownloaded: Bool
|
||
public var id: String { internalName }
|
||
|
||
public var badge: String? {
|
||
switch parakeetModel {
|
||
case .englishV2:
|
||
return "BEST FOR ENGLISH"
|
||
case .multilingualV3:
|
||
return "BEST FOR MULTILINGUAL"
|
||
case nil:
|
||
return nil
|
||
}
|
||
}
|
||
|
||
var parakeetModel: ParakeetModel? {
|
||
ParakeetModel(rawValue: internalName)
|
||
}
|
||
|
||
var isParakeet: Bool {
|
||
parakeetModel != nil
|
||
}
|
||
|
||
public init(
|
||
displayName: String,
|
||
internalName: String,
|
||
size: String,
|
||
accuracyStars: Int,
|
||
speedStars: Int,
|
||
storageSize: String,
|
||
isDownloaded: Bool
|
||
) {
|
||
self.displayName = displayName
|
||
self.internalName = internalName
|
||
self.size = size
|
||
self.accuracyStars = accuracyStars
|
||
self.speedStars = speedStars
|
||
self.storageSize = storageSize
|
||
self.isDownloaded = isDownloaded
|
||
}
|
||
|
||
// Codable (isDownloaded is set at runtime)
|
||
private enum CodingKeys: String, CodingKey { case displayName, internalName, size, accuracyStars, speedStars, storageSize }
|
||
public init(from decoder: Decoder) throws {
|
||
let c = try decoder.container(keyedBy: CodingKeys.self)
|
||
displayName = try c.decode(String.self, forKey: .displayName)
|
||
internalName = try c.decode(String.self, forKey: .internalName)
|
||
size = try c.decode(String.self, forKey: .size)
|
||
accuracyStars = try c.decode(Int.self, forKey: .accuracyStars)
|
||
speedStars = try c.decode(Int.self, forKey: .speedStars)
|
||
storageSize = try c.decode(String.self, forKey: .storageSize)
|
||
isDownloaded = false
|
||
}
|
||
}
|
||
|
||
// Convenience helper for loading the bundled models.json once.
|
||
private enum CuratedModelLoader {
|
||
private static let bundledModels: [CuratedModelInfo] = {
|
||
guard let url = Bundle.main.url(forResource: "models", withExtension: "json") ??
|
||
Bundle.main.url(forResource: "models", withExtension: "json", subdirectory: "Data")
|
||
else {
|
||
assertionFailure("models.json not found in bundle")
|
||
return []
|
||
}
|
||
do { return try JSONDecoder().decode([CuratedModelInfo].self, from: Data(contentsOf: url)) }
|
||
catch { assertionFailure("Failed to decode models.json - \(error)"); return [] }
|
||
}()
|
||
|
||
static func load() -> [CuratedModelInfo] {
|
||
bundledModels
|
||
}
|
||
}
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────
|
||
|
||
// MARK: – Domain
|
||
|
||
// ──────────────────────────────────────────────────────────────────────────
|
||
|
||
@Reducer
|
||
public struct ModelDownloadFeature {
|
||
@ObservableState
|
||
public struct State: Equatable {
|
||
// Shared user settings
|
||
@Shared(.hexSettings) var hexSettings: HexSettings
|
||
@Shared(.modelBootstrapState) var modelBootstrapState: ModelBootstrapState
|
||
|
||
// Remote data
|
||
public var availableModels: IdentifiedArrayOf<ModelInfo> = []
|
||
public var curatedModels = IdentifiedArrayOf(uniqueElements: CuratedModelLoader.load())
|
||
public var recommendedModel: String = ""
|
||
|
||
// UI state
|
||
public var isLoadingModels = false
|
||
public var isDownloading = false
|
||
public var downloadProgress: Double = 0
|
||
public var downloadError: String?
|
||
public var downloadingModelName: String?
|
||
|
||
// Track which model generated a progress update to handle switching models
|
||
public var activeDownloadID: UUID?
|
||
|
||
// Convenience computed vars
|
||
var selectedModel: String { hexSettings.selectedModel }
|
||
|
||
/// The downloaded model matching the current selection, pattern-aware so
|
||
/// legacy or glob-style selections (e.g. "distil*large-v3") still resolve.
|
||
private var downloadedModelMatchingSelection: ModelInfo? {
|
||
guard !selectedModel.isEmpty else { return nil }
|
||
return availableModels.first { model in
|
||
model.isDownloaded && ModelPatternMatcher.namesMatch(model.name, selectedModel)
|
||
}
|
||
}
|
||
|
||
var selectedModelNameForDisplay: String? {
|
||
guard !selectedModel.isEmpty else { return nil }
|
||
if let downloaded = downloadedModelMatchingSelection {
|
||
return downloaded.name
|
||
}
|
||
if modelBootstrapState.isModelReady,
|
||
let identifier = modelBootstrapState.modelIdentifier,
|
||
ModelPatternMatcher.namesMatch(identifier, selectedModel)
|
||
{
|
||
return selectedModel
|
||
}
|
||
return hexSettings.hasCompletedModelBootstrap ? selectedModel : nil
|
||
}
|
||
|
||
var selectedModelIsDownloaded: Bool {
|
||
downloadedModelMatchingSelection != nil
|
||
}
|
||
|
||
var anyModelDownloaded: Bool {
|
||
availableModels.contains(where: { $0.isDownloaded })
|
||
}
|
||
}
|
||
|
||
// MARK: Actions
|
||
|
||
public enum Action: BindableAction {
|
||
case binding(BindingAction<State>)
|
||
// Requests
|
||
case fetchModels
|
||
case selectModel(String)
|
||
case downloadModel(String)
|
||
// Effects
|
||
case modelsLoaded(recommended: String, available: [ModelInfo])
|
||
case modelsLoadFailed
|
||
case downloadProgress(id: UUID, progress: Double)
|
||
case downloadCompleted(id: UUID, result: Result<String, Error>)
|
||
case cancelDownload
|
||
|
||
case deleteModel(String)
|
||
case modelDeleted(String)
|
||
case modelDeletionFailed(Error)
|
||
case openModelLocation(String)
|
||
}
|
||
|
||
// MARK: Dependencies
|
||
|
||
@Dependency(\.transcription) var transcription
|
||
|
||
public init() {}
|
||
|
||
// MARK: Reducer
|
||
|
||
public var body: some ReducerOf<Self> {
|
||
BindingReducer()
|
||
Reduce(reduce)
|
||
}
|
||
|
||
// MARK: - Helpers (pattern matching)
|
||
|
||
private func resolvePattern(_ pattern: String, from available: [ModelInfo]) -> String? {
|
||
ModelPatternMatcher.resolvePattern(pattern, from: available.map { ($0.name, $0.isDownloaded) })
|
||
}
|
||
|
||
private func curatedDisplayName(for model: String, curated: IdentifiedArrayOf<CuratedModelInfo>) -> String {
|
||
if let match = curated.first(where: { ModelPatternMatcher.matches($0.internalName, model) }) {
|
||
return match.displayName
|
||
}
|
||
return model
|
||
.replacingOccurrences(of: "-", with: " ")
|
||
.replacingOccurrences(of: "_", with: " ")
|
||
.capitalized
|
||
}
|
||
|
||
private func updateBootstrapState(_ state: inout State) {
|
||
let model = state.hexSettings.selectedModel
|
||
guard !model.isEmpty else {
|
||
state.$modelBootstrapState.withLock { bootstrap in
|
||
bootstrap.modelIdentifier = ""
|
||
bootstrap.modelDisplayName = ""
|
||
bootstrap.isModelReady = false
|
||
bootstrap.progress = 0
|
||
bootstrap.lastError = nil
|
||
}
|
||
return
|
||
}
|
||
let displayName = curatedDisplayName(for: model, curated: state.curatedModels)
|
||
let isDownloaded = state.selectedModelIsDownloaded
|
||
state.$modelBootstrapState.withLock { bootstrap in
|
||
bootstrap.modelIdentifier = model
|
||
bootstrap.modelDisplayName = displayName
|
||
bootstrap.isModelReady = isDownloaded
|
||
if isDownloaded {
|
||
bootstrap.lastError = nil
|
||
bootstrap.progress = 1
|
||
}
|
||
}
|
||
}
|
||
|
||
private func reduce(state: inout State, action: Action) -> Effect<Action> {
|
||
switch action {
|
||
// MARK: – UI bindings
|
||
|
||
case .binding:
|
||
return .none
|
||
|
||
case let .selectModel(model):
|
||
// If the curated item is a glob (e.g., "distil*large-v3"),
|
||
// resolve it to a concrete available model so both tabs stay in sync
|
||
let resolved = resolvePattern(model, from: Array(state.availableModels)) ?? model
|
||
state.$hexSettings.withLock { $0.selectedModel = resolved }
|
||
updateBootstrapState(&state)
|
||
return .none
|
||
|
||
// MARK: – Fetch Models
|
||
|
||
case .fetchModels:
|
||
guard !state.isLoadingModels else { return .none }
|
||
state.isLoadingModels = true
|
||
return .run { send in
|
||
do {
|
||
async let recommendedSupportTask = transcription.getRecommendedModels()
|
||
async let availableNamesTask = transcription.getAvailableModels()
|
||
let recommendedSupport = try await recommendedSupportTask
|
||
let names = try await availableNamesTask
|
||
let recommended = recommendedSupport.default
|
||
let infos = try await withThrowingTaskGroup(of: ModelInfo.self) { group -> [ModelInfo] in
|
||
for name in names {
|
||
group.addTask {
|
||
ModelInfo(
|
||
name: name,
|
||
isDownloaded: await transcription.isModelDownloaded(name)
|
||
)
|
||
}
|
||
}
|
||
return try await group.reduce(into: []) { $0.append($1) }
|
||
}
|
||
await send(.modelsLoaded(recommended: recommended, available: infos))
|
||
} catch {
|
||
await send(.modelsLoadFailed)
|
||
}
|
||
}
|
||
|
||
case .modelsLoadFailed:
|
||
state.isLoadingModels = false
|
||
return .none
|
||
|
||
case let .modelsLoaded(recommended, available):
|
||
state.isLoadingModels = false
|
||
// Ensure our curated Parakeet options are visible even if WhisperKit doesn't list them
|
||
var availablePlus = available
|
||
for model in ParakeetModel.allCases.reversed() {
|
||
if !availablePlus.contains(where: { $0.name == model.identifier }) {
|
||
availablePlus.insert(ModelInfo(name: model.identifier, isDownloaded: false), at: 0)
|
||
}
|
||
}
|
||
|
||
if availablePlus.contains(where: { $0.name == state.preferredParakeetIdentifier }) {
|
||
state.recommendedModel = state.preferredParakeetIdentifier
|
||
} else {
|
||
state.recommendedModel = recommended
|
||
}
|
||
state.availableModels = IdentifiedArrayOf(uniqueElements: availablePlus)
|
||
|
||
// If the selected model is a pattern, resolve it now to the first available match
|
||
if state.hexSettings.selectedModel.contains("*") || state.hexSettings.selectedModel.contains("?") {
|
||
if let resolved = resolvePattern(state.hexSettings.selectedModel, from: available) {
|
||
state.$hexSettings.withLock { $0.selectedModel = resolved }
|
||
}
|
||
}
|
||
|
||
// Merge curated + download status with pattern support
|
||
var curated = CuratedModelLoader.load()
|
||
for idx in curated.indices {
|
||
let internalName = curated[idx].internalName
|
||
if let match = available.first(where: { ModelPatternMatcher.matches(internalName, $0.name) }) {
|
||
curated[idx].isDownloaded = match.isDownloaded
|
||
} else {
|
||
curated[idx].isDownloaded = false
|
||
}
|
||
}
|
||
state.curatedModels = IdentifiedArrayOf(uniqueElements: curated)
|
||
// If the selection isn't installed but another model is, switch to the
|
||
// installed one so transcription keeps working. Never clear the user's
|
||
// selection outright: availability scans can produce false negatives
|
||
// (e.g. after a dependency changes its cache layout), and wiping the
|
||
// setting turns a transient glitch into a permanent silent failure.
|
||
if !state.selectedModelIsDownloaded,
|
||
let installedModel = state.curatedModels.first(where: \.isDownloaded)
|
||
{
|
||
let fallback = resolvePattern(installedModel.internalName, from: Array(state.availableModels)) ?? installedModel.internalName
|
||
if fallback != state.selectedModel {
|
||
state.$hexSettings.withLock { $0.selectedModel = fallback }
|
||
}
|
||
}
|
||
updateBootstrapState(&state)
|
||
if !state.anyModelDownloaded && !state.hexSettings.hasCompletedModelBootstrap {
|
||
let preferred = state.recommendedModel.isEmpty ? state.hexSettings.selectedModel : state.recommendedModel
|
||
if !preferred.isEmpty {
|
||
state.$hexSettings.withLock { $0.selectedModel = preferred }
|
||
updateBootstrapState(&state)
|
||
}
|
||
}
|
||
return .none
|
||
|
||
// MARK: – Download
|
||
|
||
case let .downloadModel(requestedModel):
|
||
guard !requestedModel.isEmpty, !state.isDownloading else { return .none }
|
||
// Resolve glob/legacy selections to a concrete model name up front so
|
||
// the completion handler updates the matching rows and writes a
|
||
// concrete name back into settings.
|
||
let model = resolvePattern(requestedModel, from: Array(state.availableModels)) ?? requestedModel
|
||
state.downloadError = nil
|
||
state.isDownloading = true
|
||
state.downloadProgress = 0
|
||
state.downloadingModelName = model
|
||
state.activeDownloadID = UUID()
|
||
let downloadID = state.activeDownloadID!
|
||
if !state.anyModelDownloaded {
|
||
let displayName = curatedDisplayName(for: model, curated: state.curatedModels)
|
||
state.$modelBootstrapState.withLock {
|
||
$0.modelIdentifier = model
|
||
$0.modelDisplayName = displayName
|
||
$0.isModelReady = false
|
||
$0.progress = 0
|
||
$0.lastError = nil
|
||
}
|
||
}
|
||
return .run { send in
|
||
do {
|
||
try await transcription.downloadModel(model) { progress in
|
||
let fractionCompleted = progress.fractionCompleted
|
||
Task {
|
||
await send(.downloadProgress(id: downloadID, progress: fractionCompleted))
|
||
}
|
||
}
|
||
await send(.downloadCompleted(id: downloadID, result: .success(model)))
|
||
} catch is CancellationError {
|
||
} catch {
|
||
await send(.downloadCompleted(id: downloadID, result: .failure(error)))
|
||
}
|
||
}
|
||
.cancellable(id: downloadID)
|
||
|
||
case let .downloadProgress(id, progress):
|
||
guard state.activeDownloadID == id else { return .none }
|
||
guard state.downloadProgress != progress else { return .none }
|
||
state.downloadProgress = progress
|
||
if !state.modelBootstrapState.isModelReady {
|
||
state.$modelBootstrapState.withLock { $0.progress = progress }
|
||
}
|
||
return .none
|
||
|
||
case let .downloadCompleted(id, result):
|
||
guard state.activeDownloadID == id else { return .none }
|
||
state.isDownloading = false
|
||
state.downloadingModelName = nil
|
||
state.activeDownloadID = nil
|
||
state.downloadProgress = 0
|
||
var failureMessage: String?
|
||
switch result {
|
||
case let .success(name):
|
||
state.availableModels[id: name]?.isDownloaded = true
|
||
if let idx = state.curatedModels.firstIndex(where: { $0.internalName == name }) {
|
||
state.curatedModels[idx].isDownloaded = true
|
||
}
|
||
state.$hexSettings.withLock { settings in
|
||
settings.selectedModel = name
|
||
settings.hasCompletedModelBootstrap = true
|
||
}
|
||
state.downloadError = nil
|
||
case let .failure(err):
|
||
let ns = err as NSError
|
||
var message = ns.localizedDescription
|
||
if let url = ns.userInfo[NSURLErrorFailingURLErrorKey] as? URL,
|
||
let host = url.host
|
||
{
|
||
message += " (\(host))"
|
||
} else if let str = ns.userInfo[NSURLErrorFailingURLStringErrorKey] as? String,
|
||
let u = URL(string: str), let host = u.host
|
||
{
|
||
message += " (\(host))"
|
||
}
|
||
state.downloadError = message
|
||
failureMessage = message
|
||
}
|
||
state.$modelBootstrapState.withLock { bootstrap in
|
||
if let failureMessage {
|
||
bootstrap.isModelReady = false
|
||
bootstrap.lastError = failureMessage
|
||
bootstrap.progress = 0
|
||
} else {
|
||
bootstrap.isModelReady = true
|
||
bootstrap.lastError = nil
|
||
bootstrap.progress = 1
|
||
}
|
||
}
|
||
updateBootstrapState(&state)
|
||
return .none
|
||
|
||
case .cancelDownload:
|
||
guard let id = state.activeDownloadID else { return .none }
|
||
state.isDownloading = false
|
||
state.downloadingModelName = nil
|
||
state.activeDownloadID = nil
|
||
state.downloadProgress = 0
|
||
state.$modelBootstrapState.withLock { $0.progress = 0 }
|
||
updateBootstrapState(&state)
|
||
return .cancel(id: id)
|
||
|
||
case let .deleteModel(model):
|
||
guard !model.isEmpty else { return .none }
|
||
let resolved = resolvePattern(model, from: Array(state.availableModels)) ?? model
|
||
if ModelPatternMatcher.namesMatch(model, state.selectedModel) {
|
||
state.$modelBootstrapState.withLock { $0.isModelReady = false }
|
||
}
|
||
return .run { send in
|
||
do {
|
||
try await transcription.deleteModel(resolved)
|
||
await send(.modelDeleted(resolved))
|
||
} catch {
|
||
await send(.modelDeletionFailed(error))
|
||
}
|
||
}
|
||
|
||
case let .modelDeleted(model):
|
||
state.availableModels[id: model]?.isDownloaded = false
|
||
for index in state.curatedModels.indices
|
||
where ModelPatternMatcher.matches(state.curatedModels[index].internalName, model) {
|
||
state.curatedModels[index].isDownloaded = false
|
||
}
|
||
if ModelPatternMatcher.namesMatch(state.selectedModel, model) {
|
||
let fallback = state.availableModels.first { $0.isDownloaded }?.name ?? ""
|
||
state.$hexSettings.withLock { $0.selectedModel = fallback }
|
||
updateBootstrapState(&state)
|
||
}
|
||
return .send(.fetchModels)
|
||
|
||
case let .modelDeletionFailed(error):
|
||
state.downloadError = error.localizedDescription
|
||
updateBootstrapState(&state)
|
||
return .none
|
||
|
||
case let .openModelLocation(model):
|
||
return openModelLocationEffect(for: model)
|
||
}
|
||
}
|
||
|
||
// MARK: Helpers
|
||
|
||
private func openModelLocationEffect(for model: String) -> Effect<Action> {
|
||
// Parakeet caches live under FluidAudio's directory, not the WhisperKit
|
||
// models folder. Route "Show in Finder" to the matching root so users
|
||
// don't end up staring at an empty WhisperKit folder thinking the
|
||
// Parakeet download silently failed.
|
||
let usesParakeetRoot = ParakeetModel(rawValue: model) != nil
|
||
return .run { _ in
|
||
let base = try usesParakeetRoot
|
||
? URL.hexParakeetModelsDirectory
|
||
: URL.hexModelsDirectory
|
||
NSWorkspace.shared.selectFile(nil, inFileViewerRootedAtPath: base.path)
|
||
}
|
||
}
|
||
}
|
||
|
||
extension ModelDownloadFeature.State {
|
||
var preferredParakeetIdentifier: String {
|
||
(prefersEnglishParakeet ? ParakeetModel.englishV2 : ParakeetModel.multilingualV3).identifier
|
||
}
|
||
|
||
private var prefersEnglishParakeet: Bool {
|
||
guard let language = hexSettings.outputLanguage?.lowercased(), !language.isEmpty else {
|
||
return true
|
||
}
|
||
return language.hasPrefix("en")
|
||
}
|
||
}
|