diff --git a/Fluid.xcodeproj/project.pbxproj b/Fluid.xcodeproj/project.pbxproj index 673a31ed..8bf6dda1 100644 --- a/Fluid.xcodeproj/project.pbxproj +++ b/Fluid.xcodeproj/project.pbxproj @@ -13,6 +13,7 @@ 7C5AF14C2F15041600DE21B0 /* MediaRemoteAdapter in Embed Frameworks */ = {isa = PBXBuildFile; productRef = 7C5AF14A2F15041600DE21B0 /* MediaRemoteAdapter */; settings = {ATTRIBUTES = (CodeSignOnCopy, RemoveHeadersOnCopy, ); }; }; 7C9A71022F58B00000FB7CAF /* TranscribeCpp in Frameworks */ = {isa = PBXBuildFile; productRef = 7C9A71012F58B00000FB7CAF /* TranscribeCpp */; }; 7C91B0012F42AA0100C0DEF0 /* HotkeyShortcutTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7C91B0022F42AA0100C0DEF0 /* HotkeyShortcutTests.swift */; }; + A60100000000000000000002 /* AIModelRefreshTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = A60100000000000000000001 /* AIModelRefreshTests.swift */; }; 7CDB0A2D2F3C4D5600FB7CAD /* DictationE2ETests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7CDB0A292F3C4D5600FB7CAD /* DictationE2ETests.swift */; }; 7CDB0A2E2F3C4D5600FB7CAD /* AudioFixtureLoader.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7CDB0A2A2F3C4D5600FB7CAD /* AudioFixtureLoader.swift */; }; 86CAA2D4EF18433096185602 /* LLMClientRequestBodyTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 343B29013F4441D6A797D12D /* LLMClientRequestBodyTests.swift */; }; @@ -49,6 +50,7 @@ /* End PBXCopyFilesBuildPhase section */ /* Begin PBXFileReference section */ + A60100000000000000000001 /* AIModelRefreshTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AIModelRefreshTests.swift; sourceTree = ""; }; 343B29013F4441D6A797D12D /* LLMClientRequestBodyTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LLMClientRequestBodyTests.swift; sourceTree = ""; }; 980330F3CE464336ADCE3E23 /* TemperatureSupportTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TemperatureSupportTests.swift; sourceTree = ""; }; A62300000000000000000001 /* AudioBufferConverterTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AudioBufferConverterTests.swift; sourceTree = ""; }; @@ -124,6 +126,7 @@ children = ( 7CDB0A262F3C4D5600FB7CAD /* Helpers */, 7CDB0A272F3C4D5600FB7CAD /* Resources */, + A60100000000000000000001 /* AIModelRefreshTests.swift */, 7CDB0A292F3C4D5600FB7CAD /* DictationE2ETests.swift */, 7C91B0022F42AA0100C0DEF0 /* HotkeyShortcutTests.swift */, 343B29013F4441D6A797D12D /* LLMClientRequestBodyTests.swift */, @@ -281,6 +284,7 @@ isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( + A60100000000000000000002 /* AIModelRefreshTests.swift in Sources */, 7CDB0A2E2F3C4D5600FB7CAD /* AudioFixtureLoader.swift in Sources */, 7CDB0A2D2F3C4D5600FB7CAD /* DictationE2ETests.swift in Sources */, 7C91B0012F42AA0100C0DEF0 /* HotkeyShortcutTests.swift in Sources */, diff --git a/Sources/Fluid/Persistence/BackupService.swift b/Sources/Fluid/Persistence/BackupService.swift index 03528d67..9b4e5020 100644 --- a/Sources/Fluid/Persistence/BackupService.swift +++ b/Sources/Fluid/Persistence/BackupService.swift @@ -10,6 +10,9 @@ struct BackupFileVersion: Codable, Equatable { struct SettingsBackupPayload: Codable, Equatable { let selectedProviderID: String let selectedModelByProvider: [String: String] + // Optional so backups created before custom model persistence still decode. + // swiftlint:disable:next discouraged_optional_collection + let customModelsByProvider: [String: [String]]? let savedProviders: [SettingsStore.SavedProvider] let modelReasoningConfigs: [String: SettingsStore.ModelReasoningConfig] let privateAIPrefixKVCacheEnabled: Bool? diff --git a/Sources/Fluid/Persistence/SettingsStore.swift b/Sources/Fluid/Persistence/SettingsStore.swift index 8a823273..944aa6d5 100644 --- a/Sources/Fluid/Persistence/SettingsStore.swift +++ b/Sources/Fluid/Persistence/SettingsStore.swift @@ -22,6 +22,8 @@ final class SettingsStore: ObservableObject { static let privateAIDictationMinimumOutputTokens = 256 static let privateAIDictationRoundTripTokenCost = 2.75 static let privateAIBackendPreferenceDefaultsKey = "FluidIntelligenceBackendPreference" + static let customModelsByProviderDefaultsKey = "CustomModelsByProvider" + static let legacyModelCandidatesDefaultsKey = "LegacyModelCandidatesByProvider" private static let forcedOnboardingResetIntroducedAt = Date(timeIntervalSince1970: 1_782_091_732) private let defaults = UserDefaults.standard private let keychain = KeychainService.shared @@ -1412,6 +1414,127 @@ final class SettingsStore: ObservableObject { } } + var customModelsByProvider: [String: [String]] { + get { + (self.defaults.dictionary(forKey: Self.customModelsByProviderDefaultsKey) as? [String: [String]]) ?? [:] + } + set { + objectWillChange.send() + self.defaults.set(newValue, forKey: Self.customModelsByProviderDefaultsKey) + } + } + + var hasStoredCustomModelsByProvider: Bool { + self.defaults.object(forKey: Self.customModelsByProviderDefaultsKey) != nil + } + + func clearStoredCustomModelsByProvider() { + objectWillChange.send() + self.defaults.removeObject(forKey: Self.customModelsByProviderDefaultsKey) + } + + var legacyModelCandidatesByProvider: [String: [String]] { + get { + (self.defaults.dictionary( + forKey: Self.legacyModelCandidatesDefaultsKey + ) as? [String: [String]]) ?? [:] + } + set { + objectWillChange.send() + self.defaults.set(newValue, forKey: Self.legacyModelCandidatesDefaultsKey) + } + } + + func clearStoredLegacyModelCandidatesByProvider() { + objectWillChange.send() + self.defaults.removeObject(forKey: Self.legacyModelCandidatesDefaultsKey) + } + + static func availableModelsAfterRestore( + customModelsByProvider: [String: [String]], + savedProviders: [SavedProvider] + ) -> [String: [String]] { + func normalized(_ models: [String]) -> [String] { + var seen: Set = [] + return models.compactMap { model in + let trimmed = model.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty, seen.insert(trimmed).inserted else { return nil } + return trimmed + } + } + + func providerKey(_ providerID: String) -> String { + let trimmed = providerID.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return "" } + let lower = trimmed.lowercased() + if ModelRepository.shared.isBuiltIn(lower) { + return lower + } + return trimmed.hasPrefix("custom:") ? trimmed : "custom:\(trimmed)" + } + + var restored: [String: [String]] = [:] + for provider in savedProviders { + let key = providerKey(provider.id) + let models = normalized(provider.models) + if !key.isEmpty, !models.isEmpty { + restored[key] = models + } + } + + for (providerID, customModels) in customModelsByProvider { + let key = providerKey(providerID) + guard !key.isEmpty else { continue } + let defaults = ModelRepository.shared.isBuiltIn(key) + ? ModelRepository.shared.defaultModels(for: key) + : restored[key] ?? [] + let models = normalized(defaults + customModels) + if models.isEmpty { + restored.removeValue(forKey: key) + } else { + restored[key] = models + } + } + return restored + } + + func restoreModelCatalogState( + customModelsByProvider: [String: [String]], + savedProviders: [SavedProvider] + ) { + self.availableModels = [] + self.clearStoredLegacyModelCandidatesByProvider() + self.customModelsByProvider = customModelsByProvider + self.availableModelsByProvider = Self.availableModelsAfterRestore( + customModelsByProvider: customModelsByProvider, + savedProviders: savedProviders + ) + } + + func prepareLegacyModelCatalogRestore(savedProviders: [SavedProvider]) { + self.availableModels = [] + self.availableModelsByProvider = Self.availableModelsAfterRestore( + customModelsByProvider: [:], + savedProviders: savedProviders + ) + self.clearStoredCustomModelsByProvider() + self.clearStoredLegacyModelCandidatesByProvider() + } + + func customModelsByProviderForBackup() -> [String: [String]] { + guard !self.hasStoredCustomModelsByProvider else { + return self.customModelsByProvider + } + let savedModelsByProvider = Dictionary( + self.savedProviders.map { ($0.id, $0.models) }, + uniquingKeysWith: { _, newer in newer } + ) + return AIModelCatalog.migratedLegacyCustomModels( + cachedModelsByProvider: self.availableModelsByProvider, + savedModelsByProvider: savedModelsByProvider + ) + } + var enableDebugLogs: Bool { get { let value = self.defaults.object(forKey: Keys.enableDebugLogs) @@ -2994,6 +3117,7 @@ final class SettingsStore: ObservableObject { SettingsBackupPayload( selectedProviderID: self.selectedProviderID, selectedModelByProvider: self.selectedModelByProvider, + customModelsByProvider: self.customModelsByProviderForBackup(), savedProviders: self.savedProviders, modelReasoningConfigs: self.modelReasoningConfigs, privateAIPrefixKVCacheEnabled: self.privateAIPrefixKVCacheEnabled, @@ -3095,6 +3219,14 @@ final class SettingsStore: ObservableObject { self.savedProviders = payload.savedProviders self.selectedProviderID = payload.selectedProviderID self.selectedModelByProvider = payload.selectedModelByProvider + if let customModelsByProvider = payload.customModelsByProvider { + self.restoreModelCatalogState( + customModelsByProvider: customModelsByProvider, + savedProviders: self.savedProviders + ) + } else { + self.prepareLegacyModelCatalogRestore(savedProviders: self.savedProviders) + } self.modelReasoningConfigs = payload.modelReasoningConfigs if let privateAIPrefixKVCacheEnabled = payload.privateAIPrefixKVCacheEnabled { self.privateAIPrefixKVCacheEnabled = privateAIPrefixKVCacheEnabled diff --git a/Sources/Fluid/UI/AISettings/AIEnhancementSettingsViewModel.swift b/Sources/Fluid/UI/AISettings/AIEnhancementSettingsViewModel.swift index 08a869b1..ce8627f4 100644 --- a/Sources/Fluid/UI/AISettings/AIEnhancementSettingsViewModel.swift +++ b/Sources/Fluid/UI/AISettings/AIEnhancementSettingsViewModel.swift @@ -18,8 +18,10 @@ final class AIEnhancementSettingsViewModel: ObservableObject { // Model Management @Published var availableModelsByProvider: [String: [String]] = [:] + @Published var customModelsByProvider: [String: [String]] = [:] @Published var selectedModelByProvider: [String: String] = [:] @Published var availableModels: [String] = [] + private var legacyModelCandidatesByProvider: [String: [String]] = [:] @Published var selectedModel: String = "" { didSet { guard self.selectedModel != "__ADD_MODEL__" else { return } @@ -101,6 +103,18 @@ final class AIEnhancementSettingsViewModel: ObservableObject { } } + struct ManualModelAddition: Equatable { + let modelID: String + let visibleModels: [String] + let customModels: [String] + } + + struct ManualModelDeletion: Equatable { + let visibleModels: [String] + let customModels: [String] + let selectedModel: String + } + @Published var cachedProviderItems: [ProviderItemData] = [] @Published var cachedVerifiedProviderItems: [ProviderItemData] = [] @Published var cachedUnverifiedProviderItems: [ProviderItemData] = [] @@ -157,7 +171,10 @@ final class AIEnhancementSettingsViewModel: ObservableObject { self.settings.reconcilePromptStateAfterProfileChanges() self.selectedProviderID = self.settings.selectedProviderID + let hasStoredCustomModels = self.settings.hasStoredCustomModelsByProvider self.availableModelsByProvider = self.settings.availableModelsByProvider + self.customModelsByProvider = self.settings.customModelsByProvider + self.legacyModelCandidatesByProvider = self.settings.legacyModelCandidatesByProvider self.selectedModelByProvider = self.settings.selectedModelByProvider self.providerAPIKeys = self.settings.providerAPIKeys self.savedProviders = self.settings.savedProviders @@ -169,6 +186,12 @@ final class AIEnhancementSettingsViewModel: ObservableObject { self.isDictationPromptOff = self.settings.isDictationPromptOff self.isEditPromptOff = self.settings.isEditPromptOff + let legacyCachedModels = self.availableModelsByProvider + let legacySavedModels = Dictionary( + self.savedProviders.map { ($0.id, $0.models) }, + uniquingKeysWith: { _, newer in newer } + ) + if !self.selectedProviderID.isEmpty, !ModelRepository.shared.isBuiltIn(self.selectedProviderID), self.savedProviders.contains(where: { $0.id == self.selectedProviderID }) == false @@ -190,8 +213,40 @@ final class AIEnhancementSettingsViewModel: ObservableObject { let clean = Array(Set(models.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) })).sorted() if !clean.isEmpty { normalized[newKey] = clean } } + var normalizedCustom: [String: [String]] = [:] + for (key, models) in self.customModelsByProvider { + let providerKey = self.providerKey(for: key) + let clean = AIModelCatalog.normalized(models) + if !providerKey.isEmpty, !clean.isEmpty { + normalizedCustom[providerKey] = clean + } + } + if !hasStoredCustomModels { + if self.legacyModelCandidatesByProvider.isEmpty { + self.legacyModelCandidatesByProvider = Self.legacyModelCandidates( + cachedModelsByProvider: legacyCachedModels, + savedModelsByProvider: legacySavedModels + ) + self.settings.legacyModelCandidatesByProvider = self.legacyModelCandidatesByProvider + } + normalizedCustom = Self.migratedLegacyCustomModels( + cachedModelsByProvider: legacyCachedModels, + savedModelsByProvider: legacySavedModels + ) + } + for (key, customModels) in normalizedCustom { + let hasDiscoveredModels = normalized[key] != nil + normalized[key] = Self.modelsByMergingCustomModels( + normalized[key] ?? [], + customModels: customModels, + providerKey: key, + useDefaultModels: !hasDiscoveredModels + ) + } self.availableModelsByProvider = normalized self.settings.availableModelsByProvider = normalized + self.customModelsByProvider = normalizedCustom + self.settings.customModelsByProvider = normalizedCustom // Normalize selected model by provider var normalizedSel: [String: String] = [:] @@ -535,6 +590,7 @@ final class AIEnhancementSettingsViewModel: ObservableObject { func saveSavedProviders() { self.settings.savedProviders = self.savedProviders self.settings.availableModelsByProvider = self.availableModelsByProvider + self.settings.customModelsByProvider = self.customModelsByProvider self.settings.selectedModelByProvider = self.selectedModelByProvider self.settings.selectedProviderID = self.selectedProviderID self.refreshProviderItems() @@ -554,34 +610,137 @@ final class AIEnhancementSettingsViewModel: ObservableObject { return self.settings.isReasoningModel(self.selectedModel) } + static func manualModelAddition( + _ enteredModel: String, + visibleModels: [String], + customModels: [String] + ) -> ManualModelAddition? { + guard let visibleAddition = AIModelCatalog.adding(enteredModel, to: visibleModels) else { return nil } + let wasAlreadyVisible = AIModelCatalog.normalized(visibleModels).contains(visibleAddition.modelID) + let updatedCustomModels: [String] + if wasAlreadyVisible { + updatedCustomModels = AIModelCatalog.normalized(customModels) + } else { + updatedCustomModels = AIModelCatalog.normalized(customModels + [visibleAddition.modelID]) + } + return ManualModelAddition( + modelID: visibleAddition.modelID, + visibleModels: visibleAddition.models, + customModels: updatedCustomModels + ) + } + + static func migratedLegacyCustomModels( + cachedModelsByProvider: [String: [String]], + savedModelsByProvider: [String: [String]] + ) -> [String: [String]] { + AIModelCatalog.migratedLegacyCustomModels( + cachedModelsByProvider: cachedModelsByProvider, + savedModelsByProvider: savedModelsByProvider + ) + } + + static func legacyModelCandidates( + cachedModelsByProvider: [String: [String]], + savedModelsByProvider: [String: [String]] + ) -> [String: [String]] { + self.migratedLegacyCustomModels( + cachedModelsByProvider: cachedModelsByProvider, + savedModelsByProvider: savedModelsByProvider + ) + } + + static func modelsByMergingCustomModels( + _ discoveredModels: [String], + customModels: [String], + providerKey: String, + useDefaultModels: Bool + ) -> [String] { + let discoveredFallback = useDefaultModels && ModelRepository.shared.isBuiltIn(providerKey) + ? ModelRepository.shared.defaultModels(for: providerKey) + : discoveredModels + return AIModelCatalog.merged( + discoveredModels: discoveredFallback, + customModels: customModels + ) + } + + static func reconciledLegacyCustomModels( + legacyModels: [String], + discoveredModels: [String] + ) -> [String] { + let discovered = Set(AIModelCatalog.normalized(discoveredModels)) + return AIModelCatalog.normalized(legacyModels).filter { !discovered.contains($0) } + } + + static func customModelsAfterReconcilingLegacyCandidates( + customModels: [String], + legacyCandidates: [String], + discoveredModels: [String] + ) -> [String] { + let candidates = Set(AIModelCatalog.normalized(legacyCandidates)) + let retainedCustomModels = AIModelCatalog.normalized(customModels).filter { + !candidates.contains($0) + } + return AIModelCatalog.normalized( + retainedCustomModels + self.reconciledLegacyCustomModels( + legacyModels: legacyCandidates, + discoveredModels: discoveredModels + ) + ) + } + + static func visibleModelsForManualAddition( + _ cachedModels: [String], + providerKey: String + ) -> [String] { + self.modelsByMergingCustomModels( + cachedModels, + customModels: [], + providerKey: providerKey, + useDefaultModels: cachedModels.isEmpty + ) + } + func addNewModel() { - guard !self.newModelName.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return } - let modelName = self.newModelName.trimmingCharacters(in: .whitespacesAndNewlines) let key = self.providerKey(for: self.selectedProviderID) + let cachedModels = self.availableModelsByProvider[key] ?? self.availableModels + let visibleModels = Self.visibleModelsForManualAddition(cachedModels, providerKey: key) + let customModels = self.customModelsByProvider[key] ?? [] + guard let addition = Self.manualModelAddition( + self.newModelName, + visibleModels: visibleModels, + customModels: customModels + ) else { return } + let modelName = addition.modelID + let list = addition.visibleModels + + if addition.customModels.isEmpty { + self.customModelsByProvider.removeValue(forKey: key) + } else { + self.customModelsByProvider[key] = addition.customModels + } + self.settings.customModelsByProvider = self.customModelsByProvider - var list = self.availableModelsByProvider[key] ?? self.availableModels - if !list.contains(modelName) { - list.append(modelName) - self.availableModelsByProvider[key] = list - self.settings.availableModelsByProvider = self.availableModelsByProvider - - if let providerIndex = savedProviders.firstIndex(where: { $0.id == selectedProviderID }) { - let updatedProvider = SettingsStore.SavedProvider( - id: self.savedProviders[providerIndex].id, - name: self.savedProviders[providerIndex].name, - baseURL: self.savedProviders[providerIndex].baseURL, - models: list - ) - self.savedProviders[providerIndex] = updatedProvider - self.saveSavedProviders() - } + self.availableModelsByProvider[key] = list + self.settings.availableModelsByProvider = self.availableModelsByProvider - self.availableModels = list - self.selectedModel = modelName - self.selectedModelByProvider[key] = modelName - self.settings.selectedModelByProvider = self.selectedModelByProvider + if let providerIndex = savedProviders.firstIndex(where: { $0.id == selectedProviderID }) { + let updatedProvider = SettingsStore.SavedProvider( + id: self.savedProviders[providerIndex].id, + name: self.savedProviders[providerIndex].name, + baseURL: self.savedProviders[providerIndex].baseURL, + models: list + ) + self.savedProviders[providerIndex] = updatedProvider + self.saveSavedProviders() } + self.availableModels = list + self.selectedModel = modelName + self.selectedModelByProvider[key] = modelName + self.settings.selectedModelByProvider = self.selectedModelByProvider + self.showingAddModel = false self.newModelName = "" } @@ -1083,11 +1242,15 @@ final class AIEnhancementSettingsViewModel: ObservableObject { self.saveSavedProviders() let key = self.providerKey(for: self.selectedProviderID) self.availableModelsByProvider.removeValue(forKey: key) + self.customModelsByProvider.removeValue(forKey: key) + self.legacyModelCandidatesByProvider.removeValue(forKey: key) self.selectedModelByProvider.removeValue(forKey: key) self.providerAPIKeys.removeValue(forKey: key) self.saveProviderAPIKeys() self.settings.verifiedProviderFingerprints.removeValue(forKey: key) self.settings.availableModelsByProvider = self.availableModelsByProvider + self.settings.customModelsByProvider = self.customModelsByProvider + self.settings.legacyModelCandidatesByProvider = self.legacyModelCandidatesByProvider self.settings.selectedModelByProvider = self.selectedModelByProvider self.selectedProviderID = "" self.openAIBaseURL = "" @@ -1098,6 +1261,30 @@ final class AIEnhancementSettingsViewModel: ObservableObject { self.selectSoleVerifiedProviderIfNeeded() } + private func modelsAfterRefresh( + discoveredModels: [String], + providerKey: String + ) -> [String] { + if let legacyModels = self.legacyModelCandidatesByProvider.removeValue(forKey: providerKey) { + let reconciledCustomModels = Self.customModelsAfterReconcilingLegacyCandidates( + customModels: self.customModelsByProvider[providerKey] ?? [], + legacyCandidates: legacyModels, + discoveredModels: discoveredModels + ) + if reconciledCustomModels.isEmpty { + self.customModelsByProvider.removeValue(forKey: providerKey) + } else { + self.customModelsByProvider[providerKey] = reconciledCustomModels + } + self.settings.legacyModelCandidatesByProvider = self.legacyModelCandidatesByProvider + self.settings.customModelsByProvider = self.customModelsByProvider + } + return AIModelCatalog.merged( + discoveredModels: discoveredModels, + customModels: self.customModelsByProvider[providerKey] ?? [] + ) + } + func saveEditedProvider() { let name = self.editProviderName.trimmingCharacters(in: .whitespacesAndNewlines) let base = self.editProviderBaseURL.trimmingCharacters(in: .whitespacesAndNewlines) @@ -1125,29 +1312,86 @@ final class AIEnhancementSettingsViewModel: ObservableObject { self.invalidateVerification(for: self.selectedProviderID) } - func deleteSelectedModel() { - let key = self.providerKey(for: self.selectedProviderID) - var list = self.availableModelsByProvider[key] ?? self.availableModels - list.removeAll { $0 == self.selectedModel } - if list.isEmpty { list = ModelRepository.shared.defaultModels(for: key) } - self.availableModelsByProvider[key] = list + func deleteSelectedModel(for providerID: String) { + let key = self.providerKey(for: providerID) + let selectedModel = self.selectedModelByProvider[key] ?? "" + let visibleModels = self.availableModelsByProvider[key] ?? [] + guard let deletion = Self.manualModelDeletion( + selectedModel, + visibleModels: visibleModels, + customModels: self.customModelsByProvider[key] ?? [], + fallbackModels: ModelRepository.shared.defaultModels(for: key) + ) else { return } + + if deletion.customModels.isEmpty { + self.customModelsByProvider.removeValue(forKey: key) + } else { + self.customModelsByProvider[key] = deletion.customModels + } + self.legacyModelCandidatesByProvider[key]?.removeAll { $0 == selectedModel } + if self.legacyModelCandidatesByProvider[key]?.isEmpty == true { + self.legacyModelCandidatesByProvider.removeValue(forKey: key) + } + self.settings.customModelsByProvider = self.customModelsByProvider + self.settings.legacyModelCandidatesByProvider = self.legacyModelCandidatesByProvider + + self.availableModelsByProvider[key] = deletion.visibleModels self.settings.availableModelsByProvider = self.availableModelsByProvider - if let providerIndex = savedProviders.firstIndex(where: { $0.id == selectedProviderID }) { + if let providerIndex = savedProviders.firstIndex(where: { $0.id == providerID }) { let updatedProvider = SettingsStore.SavedProvider( id: self.savedProviders[providerIndex].id, name: self.savedProviders[providerIndex].name, baseURL: self.savedProviders[providerIndex].baseURL, - models: list + models: deletion.visibleModels ) self.savedProviders[providerIndex] = updatedProvider self.saveSavedProviders() } - self.availableModels = list - self.selectedModel = list.first ?? "" - self.selectedModelByProvider[key] = self.selectedModel + self.selectedModelByProvider[key] = deletion.selectedModel self.settings.selectedModelByProvider = self.selectedModelByProvider + if self.selectedProviderID == providerID { + self.availableModels = deletion.visibleModels + self.selectedModel = deletion.selectedModel + } + } + + func deleteSelectedModel() { + self.deleteSelectedModel(for: self.selectedProviderID) + } + + func canDeleteSelectedModel(for providerID: String) -> Bool { + let key = self.providerKey(for: providerID) + guard let selectedModel = self.selectedModelByProvider[key] else { return false } + return self.customModelsByProvider[key]?.contains(selectedModel) == true + } + + func canDeleteSelectedModel() -> Bool { + self.canDeleteSelectedModel(for: self.selectedProviderID) + } + + static func manualModelDeletion( + _ selectedModel: String, + visibleModels: [String], + customModels: [String], + fallbackModels: [String] + ) -> ManualModelDeletion? { + let customModels = AIModelCatalog.normalized(customModels) + guard customModels.contains(selectedModel) else { return nil } + + let remainingCustomModels = customModels.filter { $0 != selectedModel } + var remainingVisibleModels = AIModelCatalog.normalized( + visibleModels.filter { $0 != selectedModel } + ) + if remainingVisibleModels.isEmpty { + remainingVisibleModels = AIModelCatalog.normalized(fallbackModels) + } + return ManualModelDeletion( + visibleModels: remainingVisibleModels, + customModels: remainingCustomModels, + selectedModel: remainingVisibleModels.first ?? "" + ) } func fetchModelsForCurrentProvider() async { @@ -1181,8 +1425,12 @@ final class AIEnhancementSettingsViewModel: ObservableObject { // Keep existing models if fetch returned empty self.fetchModelsError = "No models returned from API" } else { - self.availableModels = models - self.availableModelsByProvider[key] = models + let mergedModels = self.modelsAfterRefresh( + discoveredModels: models, + providerKey: key + ) + self.availableModels = mergedModels + self.availableModelsByProvider[key] = mergedModels self.settings.availableModelsByProvider = self.availableModelsByProvider self.fetchedModelsProviders.insert(key) @@ -1191,15 +1439,15 @@ final class AIEnhancementSettingsViewModel: ObservableObject { id: self.savedProviders[providerIndex].id, name: self.savedProviders[providerIndex].name, baseURL: self.savedProviders[providerIndex].baseURL, - models: models + models: mergedModels ) self.savedProviders[providerIndex] = updatedProvider self.saveSavedProviders() } // Select first model if current selection not in list - if !models.contains(self.selectedModel) { - self.selectedModel = models.first ?? "" + if !mergedModels.contains(self.selectedModel) { + self.selectedModel = mergedModels.first ?? "" self.selectedModelByProvider[key] = self.selectedModel self.settings.selectedModelByProvider = self.selectedModelByProvider } @@ -1285,19 +1533,23 @@ final class AIEnhancementSettingsViewModel: ObservableObject { return } - self.availableModelsByProvider[key] = models + let selectedForProvider = self.selectedModelByProvider[key] ?? "" + let mergedModels = self.modelsAfterRefresh( + discoveredModels: models, + providerKey: key + ) + self.availableModelsByProvider[key] = mergedModels self.settings.availableModelsByProvider = self.availableModelsByProvider self.fetchedModelsProviders.insert(key) if providerID == self.selectedProviderID { - self.availableModels = models - if !models.contains(self.selectedModel) { - self.selectedModel = models.first ?? "" + self.availableModels = mergedModels + if !mergedModels.contains(self.selectedModel) { + self.selectedModel = mergedModels.first ?? "" } } - let selectedForProvider = self.selectedModelByProvider[key] ?? "" - if !models.contains(selectedForProvider), let first = models.first { + if !mergedModels.contains(selectedForProvider), let first = mergedModels.first { self.selectedModelByProvider[key] = first self.settings.selectedModelByProvider = self.selectedModelByProvider } @@ -1307,7 +1559,7 @@ final class AIEnhancementSettingsViewModel: ObservableObject { id: self.savedProviders[providerIndex].id, name: self.savedProviders[providerIndex].name, baseURL: self.savedProviders[providerIndex].baseURL, - models: models + models: mergedModels ) self.savedProviders[providerIndex] = updatedProvider self.saveSavedProviders() diff --git a/Sources/Fluid/UI/AISettings/AIModelCatalog.swift b/Sources/Fluid/UI/AISettings/AIModelCatalog.swift new file mode 100644 index 00000000..831a2da5 --- /dev/null +++ b/Sources/Fluid/UI/AISettings/AIModelCatalog.swift @@ -0,0 +1,78 @@ +import Foundation + +enum AIModelCatalog { + struct Addition: Equatable { + let modelID: String + let models: [String] + } + + static func merged( + discoveredModels: [String], + customModels: [String] + ) -> [String] { + self.normalized(discoveredModels + customModels) + } + + static func adding(_ enteredModel: String, to models: [String]) -> Addition? { + guard let modelID = self.normalized([enteredModel]).first else { return nil } + return Addition( + modelID: modelID, + models: self.normalized(models + [modelID]) + ) + } + + static func normalized(_ models: [String]) -> [String] { + var seen: Set = [] + return models.compactMap { model in + let trimmed = model.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty, seen.insert(trimmed).inserted else { return nil } + return trimmed + } + } + + static func migratedLegacyCustomModels( + cachedModelsByProvider: [String: [String]], + savedModelsByProvider: [String: [String]] + ) -> [String: [String]] { + var migrated: [String: [String]] = [:] + + func providerKey(_ providerID: String) -> String { + let trimmed = providerID.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return "" } + let lower = trimmed.lowercased() + if ModelRepository.shared.isBuiltIn(lower) { + return lower + } + return trimmed.hasPrefix("custom:") ? trimmed : "custom:\(trimmed)" + } + + func appendedManualModels(_ models: [String]) -> [String] { + let normalizedInOrder = models + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + guard normalizedInOrder.count > 1, + let firstAppendedIndex = (1.. some View { + if self.viewModel.canDeleteSelectedModel(for: providerID) { + self.companionIconButton(systemName: "trash", help: "Delete custom model") { + self.activateProvider(providerID) + self.viewModel.deleteSelectedModel(for: providerID) + } + } + } + private func reasoningButton(for providerID: String) -> some View { let hasEnabledConfig = self.viewModel.isReasoningEnabled(for: providerID) @@ -2395,7 +2435,7 @@ extension AIEnhancementSettingsView { controlHeight: AISettingsLayout.controlHeight ) - if !ModelRepository.shared.isBuiltIn(self.viewModel.selectedProviderID) { + if self.viewModel.canDeleteSelectedModel() { Button(action: { self.viewModel.deleteSelectedModel() }) { HStack(spacing: 4) { Image(systemName: "trash"); Text("Delete") }.font(.caption) } diff --git a/Tests/FluidDictationIntegrationTests/AIModelRefreshTests.swift b/Tests/FluidDictationIntegrationTests/AIModelRefreshTests.swift new file mode 100644 index 00000000..24a27e77 --- /dev/null +++ b/Tests/FluidDictationIntegrationTests/AIModelRefreshTests.swift @@ -0,0 +1,386 @@ +@testable import FluidVoice_Debug +import XCTest + +@MainActor +final class AIModelRefreshTests: XCTestCase { + func testCustomModelsPersistByProvider() { + let defaults = UserDefaults.standard + let previousValue = defaults.object(forKey: SettingsStore.customModelsByProviderDefaultsKey) + defer { + if let previousValue { + defaults.set(previousValue, forKey: SettingsStore.customModelsByProviderDefaultsKey) + } else { + defaults.removeObject(forKey: SettingsStore.customModelsByProviderDefaultsKey) + } + } + + let providerKey = "custom:issue-601-test" + SettingsStore.shared.customModelsByProvider = [providerKey: ["model/custom:nitro"]] + + XCTAssertEqual( + SettingsStore.shared.customModelsByProvider[providerKey], + ["model/custom:nitro"] + ) + } + + func testRestoringLegacyBackupClearsStoredCustomModels() { + let defaults = UserDefaults.standard + let previousValue = defaults.object(forKey: SettingsStore.customModelsByProviderDefaultsKey) + let previousAvailableModels = SettingsStore.shared.availableModels + let previousAvailableModelsByProvider = SettingsStore.shared.availableModelsByProvider + let previousLegacyCandidates = SettingsStore.shared.legacyModelCandidatesByProvider + defer { + if let previousValue { + defaults.set(previousValue, forKey: SettingsStore.customModelsByProviderDefaultsKey) + } else { + defaults.removeObject(forKey: SettingsStore.customModelsByProviderDefaultsKey) + } + SettingsStore.shared.availableModels = previousAvailableModels + SettingsStore.shared.availableModelsByProvider = previousAvailableModelsByProvider + SettingsStore.shared.legacyModelCandidatesByProvider = previousLegacyCandidates + } + + SettingsStore.shared.customModelsByProvider = ["openai": ["stale-local-model"]] + SettingsStore.shared.availableModels = ["stale-local-model"] + SettingsStore.shared.availableModelsByProvider = ["openai": ["gpt-default", "stale-local-model"]] + SettingsStore.shared.legacyModelCandidatesByProvider = ["openai": ["stale-local-model"]] + SettingsStore.shared.prepareLegacyModelCatalogRestore(savedProviders: []) + + XCTAssertFalse(SettingsStore.shared.hasStoredCustomModelsByProvider) + XCTAssertTrue(SettingsStore.shared.customModelsByProvider.isEmpty) + XCTAssertTrue(SettingsStore.shared.availableModels.isEmpty) + XCTAssertTrue(SettingsStore.shared.availableModelsByProvider.isEmpty) + XCTAssertTrue(SettingsStore.shared.legacyModelCandidatesByProvider.isEmpty) + } + + func testRestoringLegacyBackupRehydratesSavedProviderCatalog() { + let defaults = UserDefaults.standard + let previousCustomModels = defaults.object( + forKey: SettingsStore.customModelsByProviderDefaultsKey + ) + let previousLegacyCandidates = defaults.object( + forKey: SettingsStore.legacyModelCandidatesDefaultsKey + ) + let previousAvailableModels = SettingsStore.shared.availableModels + let previousAvailableModelsByProvider = SettingsStore.shared.availableModelsByProvider + defer { + if let previousCustomModels { + defaults.set( + previousCustomModels, + forKey: SettingsStore.customModelsByProviderDefaultsKey + ) + } else { + defaults.removeObject(forKey: SettingsStore.customModelsByProviderDefaultsKey) + } + if let previousLegacyCandidates { + defaults.set( + previousLegacyCandidates, + forKey: SettingsStore.legacyModelCandidatesDefaultsKey + ) + } else { + defaults.removeObject(forKey: SettingsStore.legacyModelCandidatesDefaultsKey) + } + SettingsStore.shared.availableModels = previousAvailableModels + SettingsStore.shared.availableModelsByProvider = previousAvailableModelsByProvider + } + + SettingsStore.shared.availableModelsByProvider = ["custom:provider-id": ["stale-model"]] + SettingsStore.shared.prepareLegacyModelCatalogRestore( + savedProviders: [ + SettingsStore.SavedProvider( + id: "provider-id", + name: "Provider", + baseURL: "https://example.com", + models: ["provider-model"] + ), + ] + ) + + XCTAssertEqual( + SettingsStore.shared.availableModelsByProvider["custom:provider-id"], + ["provider-model"] + ) + } + + func testBackupMigratesLegacyManualModelsBeforeSettingsLoad() { + let defaults = UserDefaults.standard + let previousCustomModels = defaults.object( + forKey: SettingsStore.customModelsByProviderDefaultsKey + ) + let previousAvailableModelsByProvider = SettingsStore.shared.availableModelsByProvider + let previousSavedProviders = SettingsStore.shared.savedProviders + defer { + if let previousCustomModels { + defaults.set( + previousCustomModels, + forKey: SettingsStore.customModelsByProviderDefaultsKey + ) + } else { + defaults.removeObject(forKey: SettingsStore.customModelsByProviderDefaultsKey) + } + SettingsStore.shared.availableModelsByProvider = previousAvailableModelsByProvider + SettingsStore.shared.savedProviders = previousSavedProviders + } + + SettingsStore.shared.clearStoredCustomModelsByProvider() + SettingsStore.shared.availableModelsByProvider = [ + "openai": ["gpt-a", "gpt-z", "custom-mid"], + ] + SettingsStore.shared.savedProviders = [ + SettingsStore.SavedProvider( + id: "provider-id", + name: "Provider", + baseURL: "https://example.com", + models: ["model-a", "model-z", "manual-provider"] + ), + ] + + let payload = SettingsStore.shared.makeBackupPayload() + + XCTAssertEqual(payload.customModelsByProvider?["openai"], ["custom-mid"]) + XCTAssertEqual( + payload.customModelsByProvider?["custom:provider-id"], + ["manual-provider"] + ) + } + + func testRestoringCustomModelsRehydratesAvailableCatalog() { + let defaults = UserDefaults.standard + let previousValue = defaults.object(forKey: SettingsStore.customModelsByProviderDefaultsKey) + let previousAvailableModels = SettingsStore.shared.availableModels + let previousAvailableModelsByProvider = SettingsStore.shared.availableModelsByProvider + let previousLegacyCandidates = SettingsStore.shared.legacyModelCandidatesByProvider + defer { + if let previousValue { + defaults.set(previousValue, forKey: SettingsStore.customModelsByProviderDefaultsKey) + } else { + defaults.removeObject(forKey: SettingsStore.customModelsByProviderDefaultsKey) + } + SettingsStore.shared.availableModels = previousAvailableModels + SettingsStore.shared.availableModelsByProvider = previousAvailableModelsByProvider + SettingsStore.shared.legacyModelCandidatesByProvider = previousLegacyCandidates + } + + SettingsStore.shared.availableModelsByProvider = [:] + SettingsStore.shared.restoreModelCatalogState( + customModelsByProvider: [ + "openai": ["gpt-custom"], + "custom:provider-id": ["provider-custom"], + ], + savedProviders: [ + SettingsStore.SavedProvider( + id: "provider-id", + name: "Provider", + baseURL: "https://example.com", + models: ["provider-discovered"] + ), + ] + ) + + XCTAssertEqual(SettingsStore.shared.customModelsByProvider["openai"], ["gpt-custom"]) + XCTAssertEqual( + SettingsStore.shared.availableModelsByProvider["openai"], + AIModelCatalog.normalized( + ModelRepository.shared.defaultModels(for: "openai") + ["gpt-custom"] + ) + ) + XCTAssertEqual( + SettingsStore.shared.availableModelsByProvider["custom:provider-id"], + ["provider-discovered", "provider-custom"] + ) + } + + func testEnteringDiscoveredModelSelectsWithoutPersistingAsCustom() { + XCTAssertEqual( + AIEnhancementSettingsViewModel.manualModelAddition( + " model/discovered ", + visibleModels: ["model/discovered", "model/other"], + customModels: [] + ), + AIEnhancementSettingsViewModel.ManualModelAddition( + modelID: "model/discovered", + visibleModels: ["model/discovered", "model/other"], + customModels: [] + ) + ) + } + + func testEnteringNewManualModelPersistsAsCustom() { + XCTAssertEqual( + AIEnhancementSettingsViewModel.manualModelAddition( + " model/custom:nitro ", + visibleModels: ["model/discovered"], + customModels: [] + ), + AIEnhancementSettingsViewModel.ManualModelAddition( + modelID: "model/custom:nitro", + visibleModels: ["model/discovered", "model/custom:nitro"], + customModels: ["model/custom:nitro"] + ) + ) + } + + func testLegacyAppendedModelsMigrateWithoutPromotingSortedCatalogs() { + XCTAssertEqual( + AIEnhancementSettingsViewModel.migratedLegacyCustomModels( + cachedModelsByProvider: [ + "OpenAI": ["gpt-a", "gpt-z", " custom-mid "], + "sorted-provider": ["model-a", "model-b", "model-c"], + ], + savedModelsByProvider: [ + "legacy-provider": ["model-a", "model-z", "manual-a", "manual-b"], + ] + ), + [ + "openai": ["custom-mid"], + "custom:legacy-provider": ["manual-a", "manual-b"], + ] + ) + } + + func testLegacyModelsReconcileAgainstFirstFreshCatalog() { + XCTAssertEqual( + AIEnhancementSettingsViewModel.reconciledLegacyCustomModels( + legacyModels: ["gpt-4.1", "gpt-5-custom"], + discoveredModels: ["gpt-4.1"] + ), + ["gpt-5-custom"] + ) + XCTAssertEqual( + AIEnhancementSettingsViewModel.reconciledLegacyCustomModels( + legacyModels: ["model-a", "model-b", "model-c"], + discoveredModels: ["model-a", "model-b", "model-c"] + ), + [] + ) + } + + func testLegacyCandidateReconciliationReplacesCandidateSubset() { + XCTAssertEqual( + AIEnhancementSettingsViewModel.customModelsAfterReconcilingLegacyCandidates( + customModels: ["official-model", "new-manual-model", "legacy-only-model"], + legacyCandidates: ["official-model", "legacy-only-model"], + discoveredModels: ["official-model"] + ), + ["new-manual-model", "legacy-only-model"] + ) + XCTAssertEqual( + AIEnhancementSettingsViewModel.customModelsAfterReconcilingLegacyCandidates( + customModels: ["official-model"], + legacyCandidates: ["official-model"], + discoveredModels: ["official-model"] + ), + [] + ) + } + + func testLegacyCandidatesExcludeSortedDiscoveryOnlyCatalogs() { + XCTAssertEqual( + AIEnhancementSettingsViewModel.legacyModelCandidates( + cachedModelsByProvider: [ + "openai": ["gpt-a", "gpt-b", "gpt-retired"], + "groq": ["model-a", "model-z", "manual-model"], + ], + savedModelsByProvider: [:] + ), + ["groq": ["manual-model"]] + ) + } + + func testCustomModelsMergeWithBuiltInDefaultsWhenCacheIsMissing() { + let defaultModels = ModelRepository.shared.defaultModels(for: "openai") + XCTAssertFalse(defaultModels.isEmpty) + + XCTAssertEqual( + AIEnhancementSettingsViewModel.modelsByMergingCustomModels( + [], + customModels: ["gpt-custom"], + providerKey: "openai", + useDefaultModels: true + ), + AIModelCatalog.normalized(defaultModels + ["gpt-custom"]) + ) + } + + func testManualAdditionStartsWithBuiltInDefaultsWhenCacheIsMissing() { + let defaultModels = ModelRepository.shared.defaultModels(for: "openai") + let visibleModels = AIEnhancementSettingsViewModel.visibleModelsForManualAddition( + [], + providerKey: "openai" + ) + + XCTAssertEqual(visibleModels, AIModelCatalog.normalized(defaultModels)) + XCTAssertEqual( + AIEnhancementSettingsViewModel.manualModelAddition( + "gpt-custom", + visibleModels: visibleModels, + customModels: [] + )?.visibleModels, + AIModelCatalog.normalized(defaultModels + ["gpt-custom"]) + ) + } + + func testRefreshDropsSelectedNonCustomModelMissingFromCatalog() { + let selectedModel = "model/retired" + let merged = AIModelCatalog.merged( + discoveredModels: ["model/a"], + customModels: [] + ) + + XCTAssertFalse(merged.contains(selectedModel)) + } + + func testRefreshKeepsSelectedCustomModel() { + let selectedModel = "model/custom:nitro" + let merged = AIModelCatalog.merged( + discoveredModels: ["model/a"], + customModels: [selectedModel] + ) + + XCTAssertTrue(merged.contains(selectedModel)) + } + + func testAddingExistingModelReturnsNormalizedSelectionWithoutDuplicate() { + XCTAssertEqual( + AIModelCatalog.adding(" model/a ", to: ["model/a", "model/b"]), + AIModelCatalog.Addition( + modelID: "model/a", + models: ["model/a", "model/b"] + ) + ) + } + + func testDeletingCustomModelRemovesOnlyManualCatalogEntries() { + XCTAssertEqual( + AIEnhancementSettingsViewModel.manualModelDeletion( + "model/custom", + visibleModels: ["model/discovered", "model/custom"], + customModels: ["model/custom"], + fallbackModels: ["model/fallback"] + ), + AIEnhancementSettingsViewModel.ManualModelDeletion( + visibleModels: ["model/discovered"], + customModels: [], + selectedModel: "model/discovered" + ) + ) + XCTAssertNil( + AIEnhancementSettingsViewModel.manualModelDeletion( + "model/discovered", + visibleModels: ["model/discovered", "model/custom"], + customModels: ["model/custom"], + fallbackModels: ["model/fallback"] + ) + ) + } + + func testRefreshDropsBlankModelIDs() { + XCTAssertEqual( + AIModelCatalog.merged( + discoveredModels: ["", " \n "], + customModels: [" model/custom "] + ), + ["model/custom"] + ) + } +}