refactor: clean AI architecture — deduplicate providers, cache instances, unify resolution, fix view layer

datlechin · datlechin · commit 3e763461d118 · 2026-04-08T08:35:37.000+07:00
diff --git a/TablePro/Core/AI/AIProvider.swift b/TablePro/Core/AI/AIProvider.swift
@@ -55,6 +55,21 @@ enum AIProviderError: Error, LocalizedError {
         }
     }
 
+    /// Base HTTP error mapping — providers can override for custom status codes
+    static func mapHTTPError(statusCode: Int, body: String) -> AIProviderError {
+        let message = parseErrorMessage(from: body) ?? body
+        switch statusCode {
+        case 401:
+            return .authenticationFailed(message)
+        case 429:
+            return .rateLimited
+        case 404:
+            return .modelNotFound(message)
+        default:
+            return .serverError(statusCode, message)
+        }
+    }
+
     /// Extract human-readable message from provider JSON error responses.
     /// Supports Anthropic (`{"error":{"message":"..."}}`), OpenAI, and Gemini formats.
     static func parseErrorMessage(from body: String) -> String? {
@@ -68,3 +83,16 @@ enum AIProviderError: Error, LocalizedError {
         return message
     }
 }
+
+// MARK: - Shared Helpers
+
+extension AIProvider {
+    func collectErrorBody(from bytes: URLSession.AsyncBytes) async throws -> String {
+        var body = ""
+        for try await line in bytes.lines {
+            body += line
+            if (body as NSString).length > 2_000 { break }
+        }
+        return body
+    }
+}
diff --git a/TablePro/Core/AI/AIProviderFactory.swift b/TablePro/Core/AI/AIProviderFactory.swift
@@ -9,29 +9,53 @@ import Foundation
 
 /// Factory for creating AI provider instances
 enum AIProviderFactory {
-    /// Create an AI provider for the given configuration
+    /// Resolved provider ready for use
+    struct ResolvedProvider {
+        let provider: AIProvider
+        let model: String
+        let config: AIProviderConfig
+    }
+
+    private static var cachedProviders: [UUID: (apiKey: String?, provider: AIProvider)] = [:]
+
+    /// Create or return a cached AI provider for the given configuration
     static func createProvider(
         for config: AIProviderConfig,
         apiKey: String?
     ) -> AIProvider {
+        if let cached = cachedProviders[config.id], cached.apiKey == apiKey {
+            return cached.provider
+        }
+
+        let provider: AIProvider
         switch config.type {
         case .claude:
-            return AnthropicProvider(
+            provider = AnthropicProvider(
                 endpoint: config.endpoint,
                 apiKey: apiKey ?? ""
             )
         case .gemini:
-            return GeminiProvider(
+            provider = GeminiProvider(
                 endpoint: config.endpoint,
                 apiKey: apiKey ?? ""
             )
         case .openAI, .openRouter, .ollama, .custom:
-            return OpenAICompatibleProvider(
+            provider = OpenAICompatibleProvider(
                 endpoint: config.endpoint,
                 apiKey: apiKey,
                 providerType: config.type
             )
         }
+        cachedProviders[config.id] = (apiKey, provider)
+        return provider
+    }
+
+    static func invalidateCache() {
+        cachedProviders.removeAll()
+    }
+
+    static func invalidateCache(for configID: UUID) {
+        cachedProviders.removeValue(forKey: configID)
     }
 
     static func resolveProvider(
@@ -62,4 +86,14 @@ enum AIProviderFactory {
         }
         return config.model
     }
+
+    /// Resolve provider, model, and config in one step
+    static func resolve(for feature: AIFeature, settings: AISettings) -> ResolvedProvider? {
+        guard let (config, apiKey) = resolveProvider(for: feature, settings: settings) else {
+            return nil
+        }
+        let model = resolveModel(for: feature, config: config, settings: settings)
+        let provider = createProvider(for: config, apiKey: apiKey)
+        return ResolvedProvider(provider: provider, model: model, config: config)
+    }
 }
diff --git a/TablePro/Core/AI/AnthropicProvider.swift b/TablePro/Core/AI/AnthropicProvider.swift
@@ -46,7 +46,7 @@ final class AnthropicProvider: AIProvider {
 
                     guard httpResponse.statusCode == 200 else {
                         let errorBody = try await collectErrorBody(from: bytes)
-                        throw mapHTTPError(
+                        throw AIProviderError.mapHTTPError(
                             statusCode: httpResponse.statusCode,
                             body: errorBody
                         )
@@ -153,7 +153,7 @@ final class AnthropicProvider: AIProvider {
         }
 
         let body = String(data: data, encoding: .utf8) ?? ""
-        throw mapHTTPError(statusCode: statusCode, body: body)
+        throw AIProviderError.mapHTTPError(statusCode: statusCode, body: body)
     }
 
     // MARK: - Private
@@ -237,31 +237,4 @@ final class AnthropicProvider: AIProvider {
         return outputTokens
     }
 
-    private func collectErrorBody(
-        from bytes: URLSession.AsyncBytes
-    ) async throws -> String {
-        var body = ""
-        for try await line in bytes.lines {
-            body += line
-            if (body as NSString).length > 2_000 { break }
-        }
-        return body
-    }
-
-    private func mapHTTPError(statusCode: Int, body: String) -> AIProviderError {
-        let message = AIProviderError.parseErrorMessage(from: body) ?? body
-
-        switch statusCode {
-        case 400:
-            return .serverError(statusCode, message)
-        case 401:
-            return .authenticationFailed(message)
-        case 429:
-            return .rateLimited
-        case 404:
-            return .modelNotFound(message)
-        default:
-            return .serverError(statusCode, message)
-        }
-    }
 }
diff --git a/TablePro/Core/AI/GeminiProvider.swift b/TablePro/Core/AI/GeminiProvider.swift
@@ -219,29 +219,11 @@ final class GeminiProvider: AIProvider {
         return request
     }
 
-    private func collectErrorBody(
-        from bytes: URLSession.AsyncBytes
-    ) async throws -> String {
-        var body = ""
-        for try await line in bytes.lines {
-            body += line
-            if (body as NSString).length > 2_000 { break }
-        }
-        return body
-    }
-
-    private func mapHTTPError(statusCode: Int, body: String) -> AIProviderError {
-        let message = AIProviderError.parseErrorMessage(from: body) ?? body
-
-        switch statusCode {
-        case 401, 403:
+    func mapHTTPError(statusCode: Int, body: String) -> AIProviderError {
+        if statusCode == 403 {
+            let message = AIProviderError.parseErrorMessage(from: body) ?? body
             return .authenticationFailed(message)
-        case 429:
-            return .rateLimited
-        case 404:
-            return .modelNotFound(message)
-        default:
-            return .serverError(statusCode, message)
         }
+        return AIProviderError.mapHTTPError(statusCode: statusCode, body: body)
     }
 }
diff --git a/TablePro/Core/AI/InlineSuggestionManager.swift b/TablePro/Core/AI/InlineSuggestionManager.swift
@@ -216,13 +216,10 @@ final class InlineSuggestionManager {
     private func fetchSuggestion(textBefore: String, fullQuery: String) async throws -> String {
         let settings = AppSettingsManager.shared.ai
 
-        guard let (config, apiKey) = AIProviderFactory.resolveProvider(for: .inlineSuggest, settings: settings) else {
+        guard let resolved = AIProviderFactory.resolve(for: .inlineSuggest, settings: settings) else {
             throw AIProviderError.networkError("No AI provider configured")
         }
 
-        let model = AIProviderFactory.resolveModel(for: .inlineSuggest, config: config, settings: settings)
-        let provider = AIProviderFactory.createProvider(for: config, apiKey: apiKey)
-
         let userMessage = AIPromptTemplates.inlineSuggest(textBefore: textBefore, fullQuery: fullQuery)
         let messages = [
             AIChatMessage(role: .user, content: userMessage)
@@ -231,9 +228,9 @@ final class InlineSuggestionManager {
         let systemPrompt = await buildSystemPrompt()
 
         var accumulated = ""
-        let stream = provider.streamChat(
+        let stream = resolved.provider.streamChat(
             messages: messages,
-            model: model,
+            model: resolved.model,
             systemPrompt: systemPrompt
         )
 
diff --git a/TablePro/Core/AI/OpenAICompatibleProvider.swift b/TablePro/Core/AI/OpenAICompatibleProvider.swift
@@ -51,7 +51,7 @@ final class OpenAICompatibleProvider: AIProvider {
 
                     guard httpResponse.statusCode == 200 else {
                         let errorBody = try await collectErrorBody(from: bytes)
-                        throw mapHTTPError(
+                        throw AIProviderError.mapHTTPError(
                             statusCode: httpResponse.statusCode,
                             body: errorBody
                         )
@@ -319,31 +319,4 @@ final class OpenAICompatibleProvider: AIProvider {
         return models.compactMap { $0["name"] as? String }.sorted()
     }
 
-    // MARK: - Helpers
-
-    private func collectErrorBody(
-        from bytes: URLSession.AsyncBytes
-    ) async throws -> String {
-        var body = ""
-        for try await line in bytes.lines {
-            body += line
-            if (body as NSString).length > 2_000 { break }
-        }
-        return body
-    }
-
-    private func mapHTTPError(statusCode: Int, body: String) -> AIProviderError {
-        let message = AIProviderError.parseErrorMessage(from: body) ?? body
-
-        switch statusCode {
-        case 401:
-            return .authenticationFailed(message)
-        case 429:
-            return .rateLimited
-        case 404:
-            return .modelNotFound(message)
-        default:
-            return .serverError(statusCode, message)
-        }
-    }
 }
diff --git a/TablePro/ViewModels/AIChatViewModel.swift b/TablePro/ViewModels/AIChatViewModel.swift
@@ -51,33 +51,26 @@ final class AIChatViewModel {
 
     // MARK: - AI Action Dispatch
 
-    private var queryLanguage: String {
-        guard let type = connection?.type else { return "sql" }
-        return PluginManager.shared.editorLanguage(for: type).codeBlockTag
-    }
-
-    private var queryTypeName: String {
-        guard let type = connection?.type else { return "SQL query" }
-        return "\(PluginManager.shared.queryLanguageName(for: type)) query"
-    }
-
     func handleFixError(query: String, error: String) {
         startNewConversation()
-        let prompt = "Fix this \(queryTypeName) error:\n\nQuery:\n```\(queryLanguage)\n\(query)\n```\n\nError: \(error)"
+        let databaseType = connection?.type ?? .mysql
+        let prompt = AIPromptTemplates.fixError(query: query, error: error, databaseType: databaseType)
         sendWithContext(prompt: prompt, feature: .fixError)
     }
 
     func handleExplainSelection(_ selectedText: String) {
         guard !selectedText.isEmpty else { return }
         startNewConversation()
-        let prompt = "Explain this \(queryTypeName):\n```\(queryLanguage)\n\(selectedText)\n```"
+        let databaseType = connection?.type ?? .mysql
+        let prompt = AIPromptTemplates.explainQuery(selectedText, databaseType: databaseType)
         sendWithContext(prompt: prompt, feature: .explainQuery)
     }
 
     func handleOptimizeSelection(_ selectedText: String) {
         guard !selectedText.isEmpty else { return }
         startNewConversation()
-        let prompt = "Optimize this \(queryTypeName):\n```\(queryLanguage)\n\(selectedText)\n```"
+        let databaseType = connection?.type ?? .mysql
+        let prompt = AIPromptTemplates.optimizeQuery(selectedText, databaseType: databaseType)
         sendWithContext(prompt: prompt, feature: .optimizeQuery)
     }
 
@@ -318,8 +311,7 @@ final class AIChatViewModel {
 
         let settings = AppSettingsManager.shared.ai
 
-        // Resolve provider from feature routing or use first enabled provider
-        guard let (config, apiKey) = AIProviderFactory.resolveProvider(for: feature, settings: settings) else {
+        guard let resolved = AIProviderFactory.resolve(for: feature, settings: settings) else {
             errorMessage = String(localized: "No AI provider configured. Go to Settings > AI to add one.")
             return
         }
@@ -342,8 +334,6 @@ final class AIChatViewModel {
             }
         }
 
-        let provider = AIProviderFactory.createProvider(for: config, apiKey: apiKey)
-        let model = AIProviderFactory.resolveModel(for: feature, config: config, settings: settings)
         let systemPrompt = buildSystemPrompt(settings: settings)
 
         // Create assistant message placeholder
@@ -361,9 +351,9 @@ final class AIChatViewModel {
             do {
                 // Exclude the empty assistant placeholder from sent messages
                 let chatMessages = Array(self.messages.dropLast())
-                let stream = provider.streamChat(
+                let stream = resolved.provider.streamChat(
                     messages: chatMessages,
-                    model: model,
+                    model: resolved.model,
                     systemPrompt: systemPrompt
                 )
 
@@ -432,4 +422,50 @@ final class AIChatViewModel {
             identifierQuote: idQuote
         )
     }
+
+    // MARK: - Schema Context
+
+    func fetchSchemaContext() async {
+        let settings = AppSettingsManager.shared.ai
+        guard settings.includeSchema,
+              let connection,
+              let driver = DatabaseManager.shared.driver(for: connection.id)
+        else { return }
+
+        let tablesToFetch = Array(tables.prefix(settings.maxSchemaTables))
+        var columns: [String: [ColumnInfo]] = [:]
+        var foreignKeys: [String: [ForeignKeyInfo]] = [:]
+
+        for table in tablesToFetch {
+            if let schemaProvider {
+                let cached = await schemaProvider.getColumns(for: table.name)
+                if !cached.isEmpty {
+                    columns[table.name] = cached
+                }
+            }
+
+            if columns[table.name] == nil {
+                do {
+                    let cols = try await driver.fetchColumns(table: table.name)
+                    columns[table.name] = cols
+                } catch {
+                    Self.logger.warning(
+                        "Failed to fetch columns for table '\(table.name)': \(error.localizedDescription)"
+                    )
+                }
+            }
+        }
+
+        do {
+            let fkResult = try await driver.fetchForeignKeys(forTables: tablesToFetch.map(\.name))
+            for (table, fks) in fkResult {
+                foreignKeys[table] = fks
+            }
+        } catch {
+            Self.logger.warning("Failed to bulk fetch foreign keys: \(error.localizedDescription)")
+        }
+
+        columnsByTable = columns
+        foreignKeysByTable = foreignKeys
+    }
 }
diff --git a/TablePro/Views/AIChat/AIChatPanelView.swift b/TablePro/Views/AIChat/AIChatPanelView.swift
diff --git a/TablePro/Views/Settings/AISettingsView.swift b/TablePro/Views/Settings/AISettingsView.swift
