diff --git a/apple/TablePro/Core/AI/AIPromptTemplates.swift b/apple/TablePro/Core/AI/AIPromptTemplates.swift index f0022f6391..c2dbd8d095 100644 --- a/apple/TablePro/Core/AI/AIPromptTemplates.swift +++ b/apple/TablePro/Core/AI/AIPromptTemplates.swift @@ -29,6 +29,20 @@ enum AIPromptTemplates { return prompt } + static func onDeviceQueryActionInstructions(engineName: String?, connectionRules: String?) -> String { + var instructions = "You are the database assistant in TablePro, a database client. " + if let engineName, !engineName.isEmpty { + instructions += "The user's database is \(engineName); answer for its dialect. " + } + instructions += "Answer the request about the user's query briefly and in plain language. " + + "Base every claim on the query and the table structure given, and say so when something you need is missing. " + + "Put any SQL in a fenced code block." + if let rules = connectionRules?.trimmingCharacters(in: .whitespacesAndNewlines), !rules.isEmpty { + instructions += "\n\nRules for this connection:\n" + rules + } + return instructions + } + @MainActor static func queryInfo(for databaseType: DatabaseType) -> (typeName: String, language: String) { let snapshot = PluginMetadataRegistry.shared.snapshot(for: databaseType) let editorLanguage = snapshot?.editorLanguage ?? .sql diff --git a/apple/TablePro/Core/AI/InlineSuggestion/AIChatInlineSource.swift b/apple/TablePro/Core/AI/InlineSuggestion/AIChatInlineSource.swift index d87b3aa23e..0edde1a65f 100644 --- a/apple/TablePro/Core/AI/InlineSuggestion/AIChatInlineSource.swift +++ b/apple/TablePro/Core/AI/InlineSuggestion/AIChatInlineSource.swift @@ -82,7 +82,7 @@ final class AIChatInlineSource: InlineSuggestionSource { } } - let cleaned = cleanSuggestion(accumulated) + let cleaned = InlineSuggestionText.clean(accumulated) guard !cleaned.isEmpty else { return nil } return InlineSuggestion( @@ -107,38 +107,4 @@ final class AIChatInlineSource: InlineSuggestionSource { } return AIPromptTemplates.inlineSuggestSystemPrompt() } - - /// Clean the AI suggestion: strip thinking blocks, leading newlines, - /// and trailing whitespace, but preserve leading spaces. - private func cleanSuggestion(_ raw: String) -> String { - var result = raw - - result = stripThinkingBlocks(result) - - // Strip leading newlines only (preserve leading spaces) - while result.first?.isNewline == true { - result.removeFirst() - } - while result.last?.isWhitespace == true { - result.removeLast() - } - return result - } - - private static let thinkingRegex: NSRegularExpression? = try? NSRegularExpression( - pattern: ".*?|.*$", - options: [.caseInsensitive, .dotMatchesLineSeparators] - ) - - /// Remove `...` blocks (case-insensitive) from AI output. - /// Handles partial/unclosed tags too. - private func stripThinkingBlocks(_ text: String) -> String { - guard let regex = Self.thinkingRegex else { return text } - - return regex.stringByReplacingMatches( - in: text, - range: NSRange(location: 0, length: (text as NSString).length), - withTemplate: "" - ) - } } diff --git a/apple/TablePro/Core/AI/InlineSuggestion/AppleIntelligenceInlineSource.swift b/apple/TablePro/Core/AI/InlineSuggestion/AppleIntelligenceInlineSource.swift new file mode 100644 index 0000000000..945c5d98f3 --- /dev/null +++ b/apple/TablePro/Core/AI/InlineSuggestion/AppleIntelligenceInlineSource.swift @@ -0,0 +1,73 @@ +// +// AppleIntelligenceInlineSource.swift +// TablePro +// + +import Foundation + +@MainActor +final class AppleIntelligenceInlineSource: InlineSuggestionSource { + typealias SettingsProvider = @MainActor () -> AISettings + + static let replyTokenLimit = 256 + + internal weak var schemaProvider: SQLSchemaProvider? + internal var connectionId: UUID? + + private let accessGate: AIConnectionAccessGate + private let currentSettings: SettingsProvider + private let model: any OnDeviceLanguageModel + + init( + schemaProvider: SQLSchemaProvider?, + connectionId: UUID?, + accessGate: AIConnectionAccessGate, + currentSettings: @escaping SettingsProvider = { AppSettingsManager.shared.ai }, + model: any OnDeviceLanguageModel = SystemOnDeviceLanguageModel() + ) { + self.schemaProvider = schemaProvider + self.connectionId = connectionId + self.accessGate = accessGate + self.currentSettings = currentSettings + self.model = model + } + + var isAvailable: Bool { + let settings = currentSettings() + guard settings.enabled, settings.useAppleIntelligence, model.availability.isAvailable else { return false } + return accessGate.allowsUnpromptedAccess(to: connectionId) + } + + func requestSuggestion(context: SuggestionContext) async throws -> InlineSuggestion? { + guard isAvailable else { return nil } + let settings = currentSettings() + let prompt = AIPromptTemplates.inlineSuggest(textBefore: context.textBefore, fullQuery: context.fullText) + + var candidates: [OnDeviceRequest] = [] + if settings.includeSchema, let provider = schemaProvider, + let schema = await provider.buildSchemaContextForAI(settings: settings), !schema.isEmpty { + candidates.append(OnDeviceRequest( + instructions: AIPromptTemplates.inlineSuggestSystemPrompt(schemaContext: schema), + prompt: prompt + )) + } + candidates.append(OnDeviceRequest(instructions: AIPromptTemplates.inlineSuggestSystemPrompt(), prompt: prompt)) + + guard case .fits(let index, let allowance) = await model.fit(candidates, reservingForReply: Self.replyTokenLimit) + else { return nil } + guard !Task.isCancelled, isAvailable else { return nil } + + var accumulated = "" + let stream = model.streamResponse( + to: candidates[index], + maximumResponseTokens: min(allowance, Self.replyTokenLimit) + ) + for try await delta in stream { + accumulated += delta + } + + let cleaned = InlineSuggestionText.clean(accumulated) + guard !cleaned.isEmpty else { return nil } + return InlineSuggestion(text: cleaned, replacementText: cleaned) + } +} diff --git a/apple/TablePro/Core/AI/InlineSuggestion/InlineSuggestionSource.swift b/apple/TablePro/Core/AI/InlineSuggestion/InlineSuggestionSource.swift index 02aa9b5e86..650437e298 100644 --- a/apple/TablePro/Core/AI/InlineSuggestion/InlineSuggestionSource.swift +++ b/apple/TablePro/Core/AI/InlineSuggestion/InlineSuggestionSource.swift @@ -48,3 +48,29 @@ extension InlineSuggestionSource { func didAcceptSuggestion(_ suggestion: InlineSuggestion) {} func didDismissSuggestion(_ suggestion: InlineSuggestion) {} } + +enum InlineSuggestionText { + private static let thinkingRegex: NSRegularExpression? = try? NSRegularExpression( + pattern: ".*?|.*$", + options: [.caseInsensitive, .dotMatchesLineSeparators] + ) + + // Leading spaces stay: a completion that starts a new token needs the one before it. + static func clean(_ raw: String) -> String { + var result = raw + if let regex = thinkingRegex { + result = regex.stringByReplacingMatches( + in: result, + range: NSRange(location: 0, length: (result as NSString).length), + withTemplate: "" + ) + } + while result.first?.isNewline == true { + result.removeFirst() + } + while result.last?.isWhitespace == true { + result.removeLast() + } + return result + } +} diff --git a/apple/TablePro/Core/AI/InlineSuggestion/InlineSuggestionSourceKind.swift b/apple/TablePro/Core/AI/InlineSuggestion/InlineSuggestionSourceKind.swift index ddeebb4621..175d08b554 100644 --- a/apple/TablePro/Core/AI/InlineSuggestion/InlineSuggestionSourceKind.swift +++ b/apple/TablePro/Core/AI/InlineSuggestion/InlineSuggestionSourceKind.swift @@ -9,12 +9,14 @@ enum InlineSuggestionSourceKind: Equatable { case off case copilot case chatCompletion + case appleIntelligence static func resolve(settings: AISettings, accessAllowed: @autoclosure () -> Bool) -> InlineSuggestionSourceKind { - guard settings.enabled, settings.inlineSuggestionsEnabled, let active = settings.activeProvider else { - return .off + guard settings.enabled, settings.inlineSuggestionsEnabled else { return .off } + if settings.useAppleIntelligence { + return accessAllowed() ? .appleIntelligence : .off } - guard accessAllowed() else { return .off } + guard let active = settings.activeProvider, accessAllowed() else { return .off } return active.type == .copilot ? .copilot : .chatCompletion } } diff --git a/apple/TablePro/Core/AI/OnDevice/AppleIntelligenceAvailability.swift b/apple/TablePro/Core/AI/OnDevice/AppleIntelligenceAvailability.swift new file mode 100644 index 0000000000..d0a13fac98 --- /dev/null +++ b/apple/TablePro/Core/AI/OnDevice/AppleIntelligenceAvailability.swift @@ -0,0 +1,42 @@ +// +// AppleIntelligenceAvailability.swift +// TablePro +// + +import Foundation + +enum AppleIntelligenceAvailability: Equatable, Sendable { + enum Reason: Equatable, Sendable { + case requiresNewerMacOS + case deviceNotEligible + case appleIntelligenceNotEnabled + case modelNotReady + case unknown + } + + case available + case unavailable(Reason) + + static let displayName = "Apple Intelligence" + static let providerID = "apple-intelligence" + + var isAvailable: Bool { + self == .available + } + + var unavailableExplanation: String? { + guard case .unavailable(let reason) = self else { return nil } + switch reason { + case .requiresNewerMacOS: + return String(localized: "Apple Intelligence needs macOS 26 or later.") + case .deviceNotEligible: + return String(localized: "This Mac does not support Apple Intelligence.") + case .appleIntelligenceNotEnabled: + return String(localized: "Turn on Apple Intelligence in System Settings > Apple Intelligence & Siri.") + case .modelNotReady: + return String(localized: "The Apple Intelligence model is still downloading.") + case .unknown: + return String(localized: "Apple Intelligence is not available on this Mac right now.") + } + } +} diff --git a/apple/TablePro/Core/AI/OnDevice/OnDeviceLanguageModel.swift b/apple/TablePro/Core/AI/OnDevice/OnDeviceLanguageModel.swift new file mode 100644 index 0000000000..e4222dc9a3 --- /dev/null +++ b/apple/TablePro/Core/AI/OnDevice/OnDeviceLanguageModel.swift @@ -0,0 +1,74 @@ +// +// OnDeviceLanguageModel.swift +// TablePro +// + +import Foundation + +struct OnDeviceRequest: Equatable, Sendable { + let instructions: String + let prompt: String +} + +struct OnDeviceUnavailableError: LocalizedError { + let availability: AppleIntelligenceAvailability + + var errorDescription: String? { + availability.unavailableExplanation + } +} + +enum OnDeviceFit: Equatable, Sendable { + case fits(index: Int, maximumResponseTokens: Int) + case tooLarge +} + +protocol OnDeviceLanguageModel: Sendable { + var availability: AppleIntelligenceAvailability { get } + var contextSize: Int { get } + func tokenCount(for request: OnDeviceRequest) async -> Int + func streamResponse(to request: OnDeviceRequest, maximumResponseTokens: Int) -> AsyncThrowingStream +} + +extension OnDeviceLanguageModel { + func fit(_ candidates: [OnDeviceRequest], reservingForReply reserve: Int) async -> OnDeviceFit { + let limit = contextSize - OnDeviceTokenBudget.transcriptMargin + for (index, candidate) in candidates.enumerated() { + guard !Task.isCancelled else { return .tooLarge } + let tokens = await tokenCount(for: candidate) + if tokens + reserve <= limit { + return .fits(index: index, maximumResponseTokens: limit - tokens) + } + } + return .tooLarge + } +} + +enum OnDeviceSnapshot { + // Scalars, not Characters: a snapshot can extend the last grapheme (a ZWJ emoji, a combining mark), + // and a Character prefix check would then reject every later snapshot. + static func delta(from emitted: String, to content: String) -> String? { + guard content.unicodeScalars.starts(with: emitted.unicodeScalars) else { return nil } + return String(content.unicodeScalars.dropFirst(emitted.unicodeScalars.count)) + } +} + +enum OnDeviceTokenBudget { + // The session wraps instructions and prompt in role markers that two separate counts miss. + static let transcriptMargin = 64 + + // Measured on the macOS 27 model (scripts/probes/check-on-device-model.sh): SQL and Markdown tables + // run 2.7 to 2.9 ASCII characters a token, CJK about 2.2. ASCII at two a token and anything else at one + // keeps the estimate high where tokenCount(for:) is missing (macOS 26.0 to 26.3). + static let asciiCharactersPerToken = 2 + static let estimateOverhead = 64 + + static func estimatedTokens(in request: OnDeviceRequest) -> Int { + var ascii = 0 + var other = 0 + for scalar in request.instructions.unicodeScalars + request.prompt.unicodeScalars { + if scalar.isASCII { ascii += 1 } else { other += 1 } + } + return (ascii + asciiCharactersPerToken - 1) / asciiCharactersPerToken + other + estimateOverhead + } +} diff --git a/apple/TablePro/Core/AI/OnDevice/SystemOnDeviceLanguageModel.swift b/apple/TablePro/Core/AI/OnDevice/SystemOnDeviceLanguageModel.swift new file mode 100644 index 0000000000..157aeab07a --- /dev/null +++ b/apple/TablePro/Core/AI/OnDevice/SystemOnDeviceLanguageModel.swift @@ -0,0 +1,78 @@ +// +// SystemOnDeviceLanguageModel.swift +// TablePro +// + +import Foundation +import FoundationModels +import os + +struct SystemOnDeviceLanguageModel: OnDeviceLanguageModel { + private static let logger = Logger(subsystem: "com.TablePro", category: "AppleIntelligence") + + var availability: AppleIntelligenceAvailability { + guard #available(macOS 26, *) else { return .unavailable(.requiresNewerMacOS) } + switch SystemLanguageModel.default.availability { + case .available: + return .available + case .unavailable(.deviceNotEligible): + return .unavailable(.deviceNotEligible) + case .unavailable(.appleIntelligenceNotEnabled): + return .unavailable(.appleIntelligenceNotEnabled) + case .unavailable(.modelNotReady): + return .unavailable(.modelNotReady) + case .unavailable: + return .unavailable(.unknown) + } + } + + var contextSize: Int { + guard #available(macOS 26, *) else { return 0 } + return SystemLanguageModel.default.contextSize + } + + func tokenCount(for request: OnDeviceRequest) async -> Int { + if #available(macOS 26.4, *) { + let model = SystemLanguageModel.default + do { + let instructions = try await model.tokenCount(for: Instructions(request.instructions)) + let prompt = try await model.tokenCount(for: request.prompt) + return instructions + prompt + } catch { + Self.logger.warning("Token count failed, estimating: \(error.publicLogShape, privacy: .public)") + } + } + return OnDeviceTokenBudget.estimatedTokens(in: request) + } + + func streamResponse(to request: OnDeviceRequest, maximumResponseTokens: Int) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + guard #available(macOS 26, *) else { + continuation.finish(throwing: OnDeviceUnavailableError(availability: .unavailable(.requiresNewerMacOS))) + return + } + let task = Task { + do { + let session = LanguageModelSession(model: .default, instructions: request.instructions) + let options = GenerationOptions(maximumResponseTokens: maximumResponseTokens) + var emitted = "" + for try await snapshot in session.streamResponse(to: request.prompt, options: options) { + let content = snapshot.content + guard let delta = OnDeviceSnapshot.delta(from: emitted, to: content) else { + Self.logger.fault("On-device snapshot did not extend the text already shown") + continue + } + emitted = content + if !delta.isEmpty { + continuation.yield(delta) + } + } + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in task.cancel() } + } + } +} diff --git a/apple/TablePro/Core/AI/QueryActions/AIQueryActionAvailability.swift b/apple/TablePro/Core/AI/QueryActions/AIQueryActionAvailability.swift index f8e3e07716..5b2a3b9133 100644 --- a/apple/TablePro/Core/AI/QueryActions/AIQueryActionAvailability.swift +++ b/apple/TablePro/Core/AI/QueryActions/AIQueryActionAvailability.swift @@ -13,6 +13,7 @@ struct AIQueryActionAvailability: Equatable { init( aiEnabled: Bool, hasActiveProvider: Bool, + appleIntelligence: AppleIntelligenceAvailability? = nil, connectionPolicy: AIConnectionPolicy?, isQueryTab: Bool, isConnected: Bool, @@ -21,6 +22,7 @@ struct AIQueryActionAvailability: Equatable { isVisible = aiEnabled && connectionPolicy != .never && isQueryTab let reason = Self.reason( hasActiveProvider: hasActiveProvider, + appleIntelligence: appleIntelligence, isConnected: isConnected, hasStatement: hasStatement ) @@ -42,8 +44,17 @@ struct AIQueryActionAvailability: Equatable { return "\(base)\n\(blockedReason)" } - private static func reason(hasActiveProvider: Bool, isConnected: Bool, hasStatement: Bool) -> String? { - if !hasActiveProvider { return String(localized: "Add an AI provider in Settings > AI first.") } + private static func reason( + hasActiveProvider: Bool, + appleIntelligence: AppleIntelligenceAvailability?, + isConnected: Bool, + hasStatement: Bool + ) -> String? { + if let appleIntelligence { + if let unavailable = appleIntelligence.unavailableExplanation { return unavailable } + } else if !hasActiveProvider { + return String(localized: "Add an AI provider in Settings > AI first.") + } if !hasStatement { return String(localized: "There is no query to send yet.") } if !isConnected { return String(localized: "This connection is not available.") } return nil diff --git a/apple/TablePro/Core/AI/QueryActions/QueryContextRenderer.swift b/apple/TablePro/Core/AI/QueryActions/QueryContextRenderer.swift index 141ef223a6..05606fb5ce 100644 --- a/apple/TablePro/Core/AI/QueryActions/QueryContextRenderer.swift +++ b/apple/TablePro/Core/AI/QueryActions/QueryContextRenderer.swift @@ -83,7 +83,10 @@ enum QueryContextRenderer { case .unavailable(let reason): return "\(heading)\nStructure unavailable: \(singleLine(reason))" case .described(let structure): - var parts = [heading, columnTable(structure.columns, limit: columnLimit)] + var parts = [ + heading, + columnTable(structure.columns, limit: columnLimit, omitted: structure.omittedColumnCount) + ] if let indexes = indexSection(structure) { parts.append(indexes) } @@ -94,8 +97,11 @@ enum QueryContextRenderer { } } - static func columnTable(_ columns: [QueryContextColumn], limit: Int) -> String { - guard !columns.isEmpty else { return "No columns reported." } + static func columnTable(_ columns: [QueryContextColumn], limit: Int, omitted: Int = 0) -> String { + guard !columns.isEmpty else { + guard omitted > 0 else { return "No columns reported." } + return "None of its \(omitted) columns are named in the statement, so none are listed." + } var lines = [ "| Column | Type | Nullable | Key | Default | Comment |", "| --- | --- | --- | --- | --- | --- |" @@ -108,9 +114,10 @@ enum QueryContextRenderer { + "| \(MarkdownFence.tableCell(defaultValue)) | \(MarkdownFence.tableCell(column.comment)) |" ) } - if columns.count > limit { + let unlisted = max(0, columns.count - limit) + omitted + if unlisted > 0 { lines.append("") - lines.append("… and \(columns.count - limit) more columns not listed.") + lines.append("… and \(unlisted) more columns not listed.") } return lines.joined(separator: "\n") } diff --git a/apple/TablePro/Core/AI/QueryActions/QueryContextSnapshot.swift b/apple/TablePro/Core/AI/QueryActions/QueryContextSnapshot.swift index eee0b98163..e65df56111 100644 --- a/apple/TablePro/Core/AI/QueryActions/QueryContextSnapshot.swift +++ b/apple/TablePro/Core/AI/QueryActions/QueryContextSnapshot.swift @@ -4,6 +4,7 @@ // import Foundation +import TableProSQLGrammar struct QueryContextColumn: Equatable, Sendable { let name: String @@ -146,6 +147,7 @@ struct QueryContextTableStructure: Equatable, Sendable { let foreignKeys: [QueryContextForeignKey] let foreignKeysUnavailableReason: String? let approximateRowCount: Int? + let omittedColumnCount: Int init( columns: [QueryContextColumn], @@ -153,7 +155,8 @@ struct QueryContextTableStructure: Equatable, Sendable { indexesUnavailableReason: String? = nil, foreignKeys: [QueryContextForeignKey] = [], foreignKeysUnavailableReason: String? = nil, - approximateRowCount: Int? = nil + approximateRowCount: Int? = nil, + omittedColumnCount: Int = 0 ) { self.columns = columns self.indexes = indexes @@ -161,6 +164,20 @@ struct QueryContextTableStructure: Equatable, Sendable { self.foreignKeys = foreignKeys self.foreignKeysUnavailableReason = foreignKeysUnavailableReason self.approximateRowCount = approximateRowCount + self.omittedColumnCount = omittedColumnCount + } + + func keepingColumns(named used: Set) -> QueryContextTableStructure { + let kept = columns.filter { used.contains($0.name.lowercased()) } + return QueryContextTableStructure( + columns: kept, + indexes: indexes, + indexesUnavailableReason: indexesUnavailableReason, + foreignKeys: foreignKeys, + foreignKeysUnavailableReason: foreignKeysUnavailableReason, + approximateRowCount: approximateRowCount, + omittedColumnCount: omittedColumnCount + columns.count - kept.count + ) } } @@ -230,3 +247,30 @@ struct QueryContextSnapshot: Equatable, Sendable { tables.filter(\.isUnavailable).map(\.name) } } + +extension QueryContextSnapshot { + func keepingColumns(usedBy statement: String, grammar: SQLLexicalGrammar) -> QueryContextSnapshot { + let used = Set(QueryTableReferenceResolver.identifierTokens(in: statement, grammar: grammar).map { $0.lowercased() }) + let trimmed = tables.map { table -> QueryContextTable in + guard case .described(let structure) = table.content else { return table } + return QueryContextTable( + name: table.name, + schema: table.schema, + kind: table.kind, + content: .described(structure.keepingColumns(named: used)) + ) + } + return QueryContextSnapshot( + engineName: engineName, + languageTag: languageTag, + serverVersion: serverVersion, + databaseName: databaseName, + schemaName: schemaName, + tables: trimmed, + notFound: notFound, + outsideScope: outsideScope, + notDescribed: notDescribed, + explainPlan: explainPlan + ) + } +} diff --git a/apple/TablePro/Core/AI/QueryActions/QueryTableReferenceResolver.swift b/apple/TablePro/Core/AI/QueryActions/QueryTableReferenceResolver.swift index 3434bcbfe2..7d7b0855a0 100644 --- a/apple/TablePro/Core/AI/QueryActions/QueryTableReferenceResolver.swift +++ b/apple/TablePro/Core/AI/QueryActions/QueryTableReferenceResolver.swift @@ -29,8 +29,8 @@ enum QueryTableReferenceResolver { return parser.references() } - static func identifierTokens(in text: String) -> [String] { - let tokens = QueryTokenScanner.tokens(in: scannedText(text), grammar: .ansi, literalsAsIdentifiers: true) + static func identifierTokens(in text: String, grammar: SQLLexicalGrammar = .ansi) -> [String] { + let tokens = QueryTokenScanner.tokens(in: scannedText(text), grammar: grammar, literalsAsIdentifiers: true) var seen = Set() var result: [String] = [] for token in tokens { diff --git a/apple/TablePro/Models/AI/AIModels.swift b/apple/TablePro/Models/AI/AIModels.swift index 424c6c46dd..1c7d6a8441 100644 --- a/apple/TablePro/Models/AI/AIModels.swift +++ b/apple/TablePro/Models/AI/AIModels.swift @@ -283,6 +283,7 @@ struct AISettings: Codable, Equatable, Sendable { var activeProviderID: UUID? var inlineSuggestionsEnabled: Bool var inlineSuggestionDebounceMs: Int + var useAppleIntelligence: Bool var includeSchema: Bool var includeCurrentQuery: Bool var includeQueryResults: Bool @@ -306,6 +307,7 @@ struct AISettings: Codable, Equatable, Sendable { activeProviderID: nil, inlineSuggestionsEnabled: false, inlineSuggestionDebounceMs: AISettings.defaultInlineSuggestionDebounceMs, + useAppleIntelligence: false, includeSchema: true, includeCurrentQuery: true, includeQueryResults: false, @@ -323,6 +325,7 @@ struct AISettings: Codable, Equatable, Sendable { activeProviderID: UUID? = nil, inlineSuggestionsEnabled: Bool = false, inlineSuggestionDebounceMs: Int = AISettings.defaultInlineSuggestionDebounceMs, + useAppleIntelligence: Bool = false, includeSchema: Bool = true, includeCurrentQuery: Bool = true, includeQueryResults: Bool = false, @@ -338,6 +341,7 @@ struct AISettings: Codable, Equatable, Sendable { self.activeProviderID = activeProviderID self.inlineSuggestionsEnabled = inlineSuggestionsEnabled self.inlineSuggestionDebounceMs = inlineSuggestionDebounceMs + self.useAppleIntelligence = useAppleIntelligence self.includeSchema = includeSchema self.includeCurrentQuery = includeCurrentQuery self.includeQueryResults = includeQueryResults @@ -359,6 +363,7 @@ struct AISettings: Codable, Equatable, Sendable { inlineSuggestionDebounceMs = container.decodeSetting( .inlineSuggestionDebounceMs, default: fallback.inlineSuggestionDebounceMs ) + useAppleIntelligence = container.decodeSetting(.useAppleIntelligence, default: fallback.useAppleIntelligence) includeSchema = container.decodeSetting(.includeSchema, default: fallback.includeSchema) includeCurrentQuery = container.decodeSetting(.includeCurrentQuery, default: fallback.includeCurrentQuery) includeQueryResults = container.decodeSetting(.includeQueryResults, default: fallback.includeQueryResults) diff --git a/apple/TablePro/Resources/Localizable.xcstrings b/apple/TablePro/Resources/Localizable.xcstrings index d1d10ab57d..710ac8e804 100644 --- a/apple/TablePro/Resources/Localizable.xcstrings +++ b/apple/TablePro/Resources/Localizable.xcstrings @@ -25264,6 +25264,15 @@ } } } + }, + "Apple Intelligence could not answer: %@" : { + + }, + "Apple Intelligence is not available on this Mac right now." : { + + }, + "Apple Intelligence needs macOS 26 or later." : { + }, "AppleScript" : { "localizations" : { @@ -45315,6 +45324,9 @@ } } } + }, + "Configure an active provider or turn on Apple Intelligence to enable inline suggestions." : { + }, "Configure an active provider to enable inline suggestions." : { "localizations" : { @@ -92652,6 +92664,9 @@ } } } + }, + "Inline suggestions and Review, Explain, Optimize and Fix run on this Mac and send nothing over the network." : { + }, "Insert" : { "extractionState" : "manual", @@ -158850,6 +158865,9 @@ } } } + }, + "Send to %@" : { + }, "Sends the code to %@" : { "localizations" : { @@ -180840,6 +180858,9 @@ } } } + }, + "The Apple Intelligence model is still downloading." : { + }, "The assistant answers once the connection is up" : { "localizations" : { @@ -181421,6 +181442,9 @@ }, "The changes were edited while the save was checking other tables. Save again to include them." : { + }, + "The chat and Agent mode keep the active provider." : { + }, "The ChatGPT account could not be identified." : { "localizations" : { @@ -194825,6 +194849,9 @@ } } } + }, + "This Mac does not support Apple Intelligence." : { + }, "this Mac needs %@ free to open that database" : { "localizations" : { @@ -195576,6 +195603,12 @@ } } } + }, + "This request is too large for Apple Intelligence, even with only the columns the statement uses. Select a smaller part of the statement." : { + + }, + "This request is too large for Apple Intelligence. Select a smaller part of the statement." : { + }, "This result has no column named %@" : { "localizations" : { @@ -200781,6 +200814,9 @@ } } } + }, + "Turn on Apple Intelligence in System Settings > Apple Intelligence & Siri." : { + }, "Turn on Set time automatically in System Settings > General > Date & Time, then reconnect." : { "extractionState" : "manual", @@ -204995,6 +205031,9 @@ }, "Use for Dark Appearance" : { + }, + "Use for inline suggestions and query actions" : { + }, "Use for Light Appearance" : { diff --git a/apple/TablePro/ViewModels/AIChatViewModel+OnDevice.swift b/apple/TablePro/ViewModels/AIChatViewModel+OnDevice.swift new file mode 100644 index 0000000000..2e90e86bb0 --- /dev/null +++ b/apple/TablePro/ViewModels/AIChatViewModel+OnDevice.swift @@ -0,0 +1,185 @@ +// +// AIChatViewModel+OnDevice.swift +// TablePro +// + +import Combine +import Foundation +import os +import TableProPluginKit +import TableProSQLGrammar + +struct OnDeviceRefusal: Equatable { + let turnID: UUID + let walkthroughBeforeSQL: String? + let walkthroughSource: QueryEditorAnchor? +} + +extension AIChatViewModel { + static let onDeviceReplyReserve = 1_024 + + var answersLastTurnOnDevice: Bool { + guard let onDeviceTurnID else { return false } + return messages.last?.id == onDeviceTurnID + } + + var fallbackProviderName: String? { + let settings = services.appSettings.ai + let config = selectedProviderId.flatMap { id in settings.providers.first { $0.id == id } } + ?? settings.activeProvider + return config?.displayName + } + + func beginOnDeviceTurn() { + guard let userTurn = messages.last, userTurn.role == .user else { return } + let model = onDeviceModel + if let unavailable = model.availability.unavailableExplanation { + clearPendingWalkthrough() + errorMessage = unavailable + return + } + AIProviderFactory.resetCopilotConversation(sessionId: sessionId) + let assistant = ChatTurn( + role: .assistant, + blocks: [], + modelId: AppleIntelligenceAvailability.displayName, + providerId: AppleIntelligenceAvailability.providerID + ) + messages.append(assistant) + trimMessagesIfNeeded() + let assistantID = assistant.id + streamingState = .streaming(assistantID: assistantID) + let instructions = AIPromptTemplates.onDeviceQueryActionInstructions( + engineName: connection.map { PluginMetadataRegistry.shared.snapshot(for: $0.type)?.displayName ?? $0.type.rawValue }, + connectionRules: connection?.aiRules + ) + streamingTask = Task { [weak self] in + await self?.answerOnDevice(userTurn: userTurn, assistantID: assistantID, model: model, instructions: instructions) + } + } + + func sendRefusedTurnToProvider() { + guard let refusal = onDeviceRefusal, messages.last?.id == refusal.turnID else { return } + onDeviceTurnID = nil + clearError() + pendingWalkthroughBeforeSQL = refusal.walkthroughBeforeSQL + pendingWalkthroughSource = refusal.walkthroughSource + startStreaming() + } + + private func answerOnDevice( + userTurn: ChatTurn, + assistantID: UUID, + model: any OnDeviceLanguageModel, + instructions: String + ) async { + let pending = unresolvedQueryContext(in: userTurn) + var snapshots: [QueryContextSnapshot] = [] + var candidates: [OnDeviceRequest] = [] + if let pending { + let full = await buildQueryContext(for: pending.attachment) + let grammar = connection.map { + SQLLexicalResolver.executionGrammar(for: $0.type, connectionId: pending.attachment.connectionId) + } ?? .ansi + let trimmed = full.keepingColumns(usedBy: pending.attachment.statement, grammar: grammar) + snapshots = trimmed == full ? [full] : [full, trimmed] + candidates = snapshots.map { snapshot in + OnDeviceRequest( + instructions: instructions, + prompt: Self.onDevicePrompt(typed: userTurn.plainText, context: QueryContextRenderer.render(snapshot)) + ) + } + } else { + let wire = await resolveTurnForWire(userTurn) + candidates = [OnDeviceRequest(instructions: instructions, prompt: wire.plainText)] + } + guard !Task.isCancelled else { return } + + let fit = await model.fit(candidates, reservingForReply: Self.onDeviceReplyReserve) + guard !Task.isCancelled else { return } + if connection != nil, resolveConnectionPolicy(settings: services.appSettings.ai) == .never { + abandonOnDeviceTurn(assistantID: assistantID, message: String(localized: "AI is disabled for this connection.")) + return + } + guard case .fits(let index, let allowance) = fit else { + refuseOnDevice(userTurnID: userTurn.id, assistantID: assistantID, trimmedStructure: pending != nil) + return + } + if let pending, snapshots.indices.contains(index) { + pending.block.setKind(.attachment(.queryContext(pending.attachment.resolved(with: snapshots[index])))) + userTurn.objectWillChange.send() + } + + do { + for try await delta in model.streamResponse(to: candidates[index], maximumResponseTokens: allowance) { + turn(withID: assistantID)?.appendStreamingToken(delta) + } + guard !Task.isCancelled else { return } + finalizeStreamingMessage(id: assistantID) + clearPendingWalkthrough() + streamingState = .idle + streamingTask = nil + persistCurrentConversation() + } catch { + guard !Task.isCancelled else { return } + Self.logger.error("On-device reply failed: \(error.publicLogShape, privacy: .public)") + clearPendingWalkthrough() + errorMessage = String( + format: String(localized: "Apple Intelligence could not answer: %@"), + error.localizedDescription + ) + streamingState = .failed(nil) + finalizeStreamingMessage(id: assistantID) + if let index = messages.firstIndex(where: { $0.id == assistantID }), messages[index].blocks.isEmpty { + messages.remove(at: index) + } + streamingTask = nil + } + } + + private func abandonOnDeviceTurn(assistantID: UUID, message: String) { + messages.removeAll { $0.id == assistantID } + clearPendingWalkthrough() + streamingState = .idle + streamingTask = nil + errorMessage = message + } + + private func refuseOnDevice(userTurnID: UUID, assistantID: UUID, trimmedStructure: Bool) { + messages.removeAll { $0.id == assistantID } + streamingState = .idle + streamingTask = nil + let refusal = OnDeviceRefusal( + turnID: userTurnID, + walkthroughBeforeSQL: pendingWalkthroughBeforeSQL, + walkthroughSource: pendingWalkthroughSource + ) + clearPendingWalkthrough() + errorMessage = trimmedStructure ? Self.tooLargeWithTrimmedStructure : Self.tooLarge + onDeviceRefusal = refusal + persistCurrentConversation() + } + + private static var tooLarge: String { + String(localized: "This request is too large for Apple Intelligence. Select a smaller part of the statement.") + } + + private static var tooLargeWithTrimmedStructure: String { + String( + localized: "This request is too large for Apple Intelligence, even with only the columns the statement uses. Select a smaller part of the statement." + ) + } + + private func unresolvedQueryContext(in turn: ChatTurn) -> (block: ChatContentBlock, attachment: QueryContextAttachment)? { + for block in turn.blocks { + if case .attachment(.queryContext(let attachment)) = block.kind, !attachment.isResolved { + return (block, attachment) + } + } + return nil + } + + private static func onDevicePrompt(typed: String, context: String) -> String { + typed.isEmpty ? context : typed + "\n\n---\n\n" + context + } +} diff --git a/apple/TablePro/ViewModels/AIChatViewModel+QueryActions.swift b/apple/TablePro/ViewModels/AIChatViewModel+QueryActions.swift index a3e44554c8..c0a8ff27ef 100644 --- a/apple/TablePro/ViewModels/AIChatViewModel+QueryActions.swift +++ b/apple/TablePro/ViewModels/AIChatViewModel+QueryActions.swift @@ -25,6 +25,11 @@ extension AIChatViewModel { errorMessage = rejection return } + let onDevice = services.appSettings.ai.useAppleIntelligence + if onDevice, let unavailable = onDeviceModel.availability.unavailableExplanation { + errorMessage = unavailable + return + } if let invocationText { messages.append(ChatTurn(role: .user, blocks: [.text(invocationText)])) @@ -39,7 +44,9 @@ extension AIChatViewModel { if let attachment { blocks.append(.attachment(.queryContext(attachment))) } - messages.append(ChatTurn(role: .user, blocks: blocks)) + let turn = ChatTurn(role: .user, blocks: blocks) + messages.append(turn) + onDeviceTurnID = onDevice ? turn.id : nil trimMessagesIfNeeded() clearError() pendingWalkthroughBeforeSQL = request.statement @@ -87,7 +94,7 @@ extension AIChatViewModel { } } - private func buildQueryContext(for attachment: QueryContextAttachment) async -> QueryContextSnapshot { + func buildQueryContext(for attachment: QueryContextAttachment) async -> QueryContextSnapshot { guard let connection, connection.id == attachment.connectionId else { return QueryContextSnapshot(engineName: "", databaseName: attachment.database, schemaName: attachment.schema) } diff --git a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift index a091a133cf..3c4fb970d3 100644 --- a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift +++ b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift @@ -47,12 +47,13 @@ extension AIChatViewModel { let settings = services.appSettings.ai - let resolved = AIProviderFactory.resolve( + let answersOnDevice = answersLastTurnOnDevice + let resolved = answersOnDevice ? nil : AIProviderFactory.resolve( settings: settings, overrideProviderId: selectedProviderId, overrideModel: selectedModel ) - guard let resolved else { + if !answersOnDevice, resolved == nil { clearPendingWalkthrough() errorMessage = String(localized: "No AI provider configured. Go to Settings > AI to add one.") return @@ -74,6 +75,10 @@ extension AIChatViewModel { } } + guard let resolved else { + beginOnDeviceTurn() + return + } beginStreamingTurn( resolved: resolved, settings: settings, diff --git a/apple/TablePro/ViewModels/AIChatViewModel.swift b/apple/TablePro/ViewModels/AIChatViewModel.swift index 2a118460af..be8e9d6341 100644 --- a/apple/TablePro/ViewModels/AIChatViewModel.swift +++ b/apple/TablePro/ViewModels/AIChatViewModel.swift @@ -101,6 +101,9 @@ final class AIChatViewModel: ObservableObject { var pendingWalkthroughBeforeSQL: String? var pendingWalkthroughSource: QueryEditorAnchor? + var onDeviceModel: any OnDeviceLanguageModel = SystemOnDeviceLanguageModel() + var onDeviceTurnID: UUID? + @Published var onDeviceRefusal: OnDeviceRefusal? var queryContextMetadata: any ScopedMetadataProviding var inFlightColumnFetches: [String: Task] = [:] var inFlightSchemaLoad: Task? @@ -296,14 +299,22 @@ final class AIChatViewModel: ObservableObject { let lastAssistantIndex = messages.lastIndex(where: { $0.role == .assistant }) else { return } - AIProviderFactory.copilotDeleteLastTurn(sessionId: sessionId) + let answeredOnDevice = messages[lastAssistantIndex].providerId == AppleIntelligenceAvailability.providerID + if !answeredOnDevice { + AIProviderFactory.copilotDeleteLastTurn(sessionId: sessionId) + } + let wasLast = lastAssistantIndex == messages.count - 1 messages.remove(at: lastAssistantIndex) + if wasLast, answeredOnDevice, let asked = messages.last, asked.role == .user { + onDeviceTurnID = asked.id + } clearError() startStreaming() } func clearError() { errorMessage = nil + onDeviceRefusal = nil if case .failed = streamingState { streamingState = .idle } diff --git a/apple/TablePro/Views/AIChat/AIChatPanelView.swift b/apple/TablePro/Views/AIChat/AIChatPanelView.swift index d8ae75da4b..8dddee0d30 100644 --- a/apple/TablePro/Views/AIChat/AIChatPanelView.swift +++ b/apple/TablePro/Views/AIChat/AIChatPanelView.swift @@ -52,11 +52,11 @@ struct AIChatPanelView: View { AIChatTranscriptView(viewModel: viewModel, contentWidth: contentWidth) } - if hasConfiguredProvider { - if let error = viewModel.errorMessage { - errorBanner(error) - } + if let error = viewModel.errorMessage { + errorBanner(error) + } + if hasConfiguredProvider { inputArea } else if !viewModel.messages.isEmpty { noProviderFooter @@ -122,6 +122,12 @@ struct AIChatPanelView: View { .foregroundStyle(.secondary) .fixedSize(horizontal: false, vertical: true) Spacer() + if viewModel.isStreaming { + Button(String(localized: "Stop Generating")) { + viewModel.cancelStream() + } + .controlSize(.small) + } Button(String(localized: "Settings…")) { WindowOpener.shared.openSettings(tab: .ai) } @@ -151,10 +157,19 @@ struct AIChatPanelView: View { HStack(spacing: 6) { Image(systemName: "exclamationmark.triangle.fill") .foregroundStyle(.yellow) - Text(message) - .font(.caption) - .foregroundStyle(.secondary) - .lineLimit(2) + VStack(alignment: .leading, spacing: 4) { + Text(message) + .font(.caption) + .foregroundStyle(.secondary) + .lineLimit(3) + .help(message) + if viewModel.onDeviceRefusal != nil, let provider = viewModel.fallbackProviderName { + Button(String(format: String(localized: "Send to %@"), provider)) { + viewModel.sendRefusedTurnToProvider() + } + .controlSize(.small) + } + } Spacer() Button { viewModel.clearError() diff --git a/apple/TablePro/Views/Editor/SQLEditorCoordinator.swift b/apple/TablePro/Views/Editor/SQLEditorCoordinator.swift index 5c71b5231c..45aa9dc426 100644 --- a/apple/TablePro/Views/Editor/SQLEditorCoordinator.swift +++ b/apple/TablePro/Views/Editor/SQLEditorCoordinator.swift @@ -39,6 +39,7 @@ final class SQLEditorCoordinator: ObservableObject, TextViewCoordinator, TextVie private var contextMenu: AIEditorContextMenu? private var inlineSuggestionManager: InlineSuggestionManager? private var aiChatInlineSource: AIChatInlineSource? + private var appleIntelligenceInlineSource: AppleIntelligenceInlineSource? private var copilotDocumentSync: CopilotDocumentSync? private var copilotInlineSource: CopilotInlineSource? private var editorSettingsCancellable: AnyCancellable? @@ -313,6 +314,7 @@ final class SQLEditorCoordinator: ObservableObject, TextViewCoordinator, TextVie copilotDocumentSync = nil copilotInlineSource = nil aiChatInlineSource = nil + appleIntelligenceInlineSource = nil // Release closure captures to break potential retain cycles releaseWindowSentinel() @@ -541,6 +543,17 @@ final class SQLEditorCoordinator: ObservableObject, TextViewCoordinator, TextVie aiChatInlineSource?.schemaProvider = schemaProvider aiChatInlineSource?.connectionId = connectionId return aiChatInlineSource + case .appleIntelligence: + if appleIntelligenceInlineSource == nil { + appleIntelligenceInlineSource = AppleIntelligenceInlineSource( + schemaProvider: schemaProvider, + connectionId: connectionId, + accessGate: inlineAccessGate + ) + } + appleIntelligenceInlineSource?.schemaProvider = schemaProvider + appleIntelligenceInlineSource?.connectionId = connectionId + return appleIntelligenceInlineSource } } @@ -589,6 +602,9 @@ final class SQLEditorCoordinator: ObservableObject, TextViewCoordinator, TextVie if kind != .chatCompletion { aiChatInlineSource = nil } + if kind != .appleIntelligence { + appleIntelligenceInlineSource = nil + } } // MARK: - Vim Mode diff --git a/apple/TablePro/Views/Main/Extensions/MainContentCoordinator+AIQueryActions.swift b/apple/TablePro/Views/Main/Extensions/MainContentCoordinator+AIQueryActions.swift index 5ef594ea84..c9998b929f 100644 --- a/apple/TablePro/Views/Main/Extensions/MainContentCoordinator+AIQueryActions.swift +++ b/apple/TablePro/Views/Main/Extensions/MainContentCoordinator+AIQueryActions.swift @@ -27,6 +27,7 @@ extension MainContentCoordinator { return AIQueryActionAvailability( aiEnabled: settings.enabled, hasActiveProvider: settings.hasActiveProvider, + appleIntelligence: settings.useAppleIntelligence ? SystemOnDeviceLanguageModel().availability : nil, connectionPolicy: policy, isQueryTab: tab?.tabType == .query, isConnected: MainWindowToolbar.hasLiveSession(toolbarState.connectionState), diff --git a/apple/TablePro/Views/Settings/AISettingsView.swift b/apple/TablePro/Views/Settings/AISettingsView.swift index e59b3f04a2..6fc07cb1ee 100644 --- a/apple/TablePro/Views/Settings/AISettingsView.swift +++ b/apple/TablePro/Views/Settings/AISettingsView.swift @@ -25,6 +25,7 @@ struct AISettingsView: View { if settings.enabled { activeProviderSection providersSection + appleIntelligenceSection inlineSuggestionsSection agentSection contextSection @@ -230,15 +231,41 @@ struct AISettingsView: View { } } + // MARK: - Apple Intelligence + + private var appleIntelligenceSection: some View { + let availability = SystemOnDeviceLanguageModel().availability + return Section { + Toggle("Use for inline suggestions and query actions", isOn: $settings.useAppleIntelligence) + .disabled(!availability.isAvailable && !settings.useAppleIntelligence) + } header: { + Text(AppleIntelligenceAvailability.displayName) + } footer: { + Text(Self.appleIntelligenceFooter(availability)) + .font(.caption) + .foregroundStyle(.secondary) + } + } + + private static func appleIntelligenceFooter(_ availability: AppleIntelligenceAvailability) -> String { + if let unavailable = availability.unavailableExplanation { return unavailable } + return String(localized: "Inline suggestions and Review, Explain, Optimize and Fix run on this Mac and send nothing over the network.") + + " " + String(localized: "The chat and Agent mode keep the active provider.") + } + // MARK: - Inline Suggestions + private var canSuggestInline: Bool { + settings.hasActiveProvider || settings.useAppleIntelligence + } + private var inlineSuggestionsSection: some View { Section { Toggle("Enable inline suggestions while typing", isOn: $settings.inlineSuggestionsEnabled) - .disabled(!settings.hasActiveProvider) - .help(settings.hasActiveProvider + .disabled(!canSuggestInline) + .help(canSuggestInline ? "" - : String(localized: "Configure an active provider to enable inline suggestions.")) + : String(localized: "Configure an active provider or turn on Apple Intelligence to enable inline suggestions.")) Stepper( String(format: String(localized: "Debounce: %d ms"), settings.inlineSuggestionDebounceMs), value: $settings.inlineSuggestionDebounceMs, diff --git a/apple/TableProTests/Core/AI/AIChatInlineSourceTests.swift b/apple/TableProTests/Core/AI/AIChatInlineSourceTests.swift index 8cd1fb14e1..443e7fdbc7 100644 --- a/apple/TableProTests/Core/AI/AIChatInlineSourceTests.swift +++ b/apple/TableProTests/Core/AI/AIChatInlineSourceTests.swift @@ -193,4 +193,21 @@ internal struct InlineSuggestionSourceKindTests { #expect(InlineSuggestionSourceKind.resolve(settings: disabled, accessAllowed: true) == .off) } + + @Test("Apple Intelligence takes inline suggestions over the active provider, and needs none") + func appleIntelligenceTakesInline() { + var copilot = settings(providerType: .copilot) + copilot.useAppleIntelligence = true + let alone = AISettings(enabled: true, inlineSuggestionsEnabled: true, useAppleIntelligence: true) + + #expect(InlineSuggestionSourceKind.resolve(settings: copilot, accessAllowed: true) == .appleIntelligence) + #expect(InlineSuggestionSourceKind.resolve(settings: alone, accessAllowed: true) == .appleIntelligence) + } + + @Test("Apple Intelligence follows the connection's AI policy too") + func appleIntelligenceFollowsTheGate() { + let alone = AISettings(enabled: true, inlineSuggestionsEnabled: true, useAppleIntelligence: true) + + #expect(InlineSuggestionSourceKind.resolve(settings: alone, accessAllowed: false) == .off) + } } diff --git a/apple/TableProTests/Core/AI/AIQueryActionSupportTests.swift b/apple/TableProTests/Core/AI/AIQueryActionSupportTests.swift index bef3a5ff97..9d1e15aba5 100644 --- a/apple/TableProTests/Core/AI/AIQueryActionSupportTests.swift +++ b/apple/TableProTests/Core/AI/AIQueryActionSupportTests.swift @@ -11,6 +11,7 @@ struct AIQueryActionAvailabilityTests { private func availability( aiEnabled: Bool = true, provider: Bool = true, + appleIntelligence: AppleIntelligenceAvailability? = nil, policy: AIConnectionPolicy? = .alwaysAllow, queryTab: Bool = true, connected: Bool = true, @@ -19,6 +20,7 @@ struct AIQueryActionAvailabilityTests { AIQueryActionAvailability( aiEnabled: aiEnabled, hasActiveProvider: provider, + appleIntelligence: appleIntelligence, connectionPolicy: policy, isQueryTab: queryTab, isConnected: connected, @@ -59,6 +61,30 @@ struct AIQueryActionAvailabilityTests { #expect(value.hint(base: "Review with AI").hasPrefix("Review with AI\n")) } } + + @Test("Apple Intelligence answers the actions with no provider set up") + func appleIntelligenceNeedsNoProvider() { + let value = availability(provider: false, appleIntelligence: .available) + #expect(value.isEnabled) + #expect(value.blockedReason == nil) + } + + @Test("Apple Intelligence that cannot run dims the actions with its reason, provider or not") + func appleIntelligenceUnavailable() { + let unavailable = AppleIntelligenceAvailability.unavailable(.appleIntelligenceNotEnabled) + for provider in [true, false] { + let value = availability(provider: provider, appleIntelligence: unavailable) + #expect(value.isVisible) + #expect(!value.isEnabled) + #expect(value.blockedReason == unavailable.unavailableExplanation) + } + } + + @Test("A Never connection hides the actions even with Apple Intelligence on") + func appleIntelligenceKeepsThePolicy() { + let value = availability(appleIntelligence: .available, policy: .never) + #expect(!value.isVisible) + } } struct AIQueryTargetTests { diff --git a/apple/TableProTests/Core/AI/AppleIntelligenceInlineSourceTests.swift b/apple/TableProTests/Core/AI/AppleIntelligenceInlineSourceTests.swift new file mode 100644 index 0000000000..baa860dff3 --- /dev/null +++ b/apple/TableProTests/Core/AI/AppleIntelligenceInlineSourceTests.swift @@ -0,0 +1,127 @@ +// +// AppleIntelligenceInlineSourceTests.swift +// TableProTests +// + +import Foundation +@testable import TablePro +import Testing + +private final class PolicySwitch: @unchecked Sendable { + private let lock = NSLock() + private var current: AIConnectionPolicy + + init(_ policy: AIConnectionPolicy) { + current = policy + } + + var policy: AIConnectionPolicy { + get { lock.withLock { current } } + set { lock.withLock { current = newValue } } + } +} + +@MainActor +internal struct AppleIntelligenceInlineSourceTests { + private let connectionId = UUID() + private let approvals = AIAccessApprovals() + + private let context = SuggestionContext( + textBefore: "SELECT * ", + fullText: "SELECT * ", + cursorOffset: 9, + cursorLine: 0, + cursorCharacter: 9 + ) + + private func makeSource( + _ model: FakeOnDeviceModel, + policy: AIConnectionPolicy = .alwaysAllow, + useAppleIntelligence: Bool = true + ) -> AppleIntelligenceInlineSource { + makeSource(model, policy: PolicySwitch(policy), useAppleIntelligence: useAppleIntelligence) + } + + private func makeSource( + _ model: FakeOnDeviceModel, + policy: PolicySwitch, + useAppleIntelligence: Bool = true + ) -> AppleIntelligenceInlineSource { + let settings = AISettings( + enabled: true, + inlineSuggestionsEnabled: true, + useAppleIntelligence: useAppleIntelligence, + includeSchema: false + ) + return AppleIntelligenceInlineSource( + schemaProvider: nil, + connectionId: connectionId, + accessGate: AIConnectionAccessGate(approvals: approvals) { _ in policy.policy }, + currentSettings: { settings }, + model: model + ) + } + + @Test("A suggestion comes from the on-device model, cleaned, with the reply kept short") + func suggestsOnDevice() async throws { + let model = FakeOnDeviceModel(reply: ["\nFROM ", "users \n"]) + let source = makeSource(model) + + let suggestion = try await source.requestSuggestion(context: context) + + #expect(source.isAvailable) + #expect(suggestion?.text == "FROM users") + let streamed = try #require(model.streamedRequests.first) + #expect(streamed.request.prompt.contains("SELECT * [CURSOR]")) + #expect(streamed.maximumResponseTokens == AppleIntelligenceInlineSource.replyTokenLimit) + } + + @Test("A model that cannot run, a Never connection or the setting off sends nothing", arguments: [ + (AppleIntelligenceAvailability.unavailable(.modelNotReady), AIConnectionPolicy.alwaysAllow, true), + (.available, .never, true), + (.available, .alwaysAllow, false) + ]) + func unavailableSendsNothing( + availability: AppleIntelligenceAvailability, + policy: AIConnectionPolicy, + enabled: Bool + ) async throws { + let model = FakeOnDeviceModel(availability: availability) + let source = makeSource(model, policy: policy, useAppleIntelligence: enabled) + + let suggestion = try await source.requestSuggestion(context: context) + + #expect(!source.isAvailable) + #expect(suggestion == nil) + #expect(model.countedRequests.isEmpty) + #expect(model.streamedRequests.isEmpty) + } + + @Test("A policy turned to Never while the request is measured stops it before generation") + func policyChangedWhileMeasuringSendsNothing() async throws { + let policy = PolicySwitch(.alwaysAllow) + let model = FakeOnDeviceModel { _ in + policy.policy = .never + return 10 + } + let source = makeSource(model, policy: policy) + + let suggestion = try await source.requestSuggestion(context: context) + + #expect(suggestion == nil) + #expect(model.countedRequests.count == 1) + #expect(model.streamedRequests.isEmpty) + } + + @Test("A request too large for the model is dropped before generation") + func tooLargeIsDropped() async throws { + let model = FakeOnDeviceModel(contextSize: 300) { _ in 200 } + let source = makeSource(model) + + let suggestion = try await source.requestSuggestion(context: context) + + #expect(suggestion == nil) + #expect(model.countedRequests.count == 1) + #expect(model.streamedRequests.isEmpty) + } +} diff --git a/apple/TableProTests/Core/AI/AppleIntelligenceTests.swift b/apple/TableProTests/Core/AI/AppleIntelligenceTests.swift new file mode 100644 index 0000000000..b466470905 --- /dev/null +++ b/apple/TableProTests/Core/AI/AppleIntelligenceTests.swift @@ -0,0 +1,188 @@ +// +// AppleIntelligenceTests.swift +// TableProTests +// + +import Foundation +@testable import TablePro +import TableProSQLGrammar +import Testing + +struct OnDeviceFitTests { + private let small = OnDeviceRequest(instructions: "i", prompt: "small") + private let large = OnDeviceRequest(instructions: "i", prompt: "large") + + private func model(context: Int = 4_096) -> FakeOnDeviceModel { + FakeOnDeviceModel(contextSize: context) { $0.prompt == "large" ? 3_500 : 900 } + } + + @Test("The first candidate that leaves room for the reply wins, and the reply gets the rest") + func firstFittingCandidate() async { + let fit = await model().fit([large, small], reservingForReply: 1_024) + #expect(fit == .fits(index: 1, maximumResponseTokens: 4_096 - OnDeviceTokenBudget.transcriptMargin - 900)) + } + + @Test("A candidate that fits stops the measuring") + func stopsAtFirstFit() async { + let fake = model() + _ = await fake.fit([small, large], reservingForReply: 1_024) + #expect(fake.countedRequests == [small]) + } + + @Test("Nothing that leaves room for the reply is too large, so nothing is generated") + func tooLarge() async { + let fake = model(context: 1_500) + let fit = await fake.fit([large, small], reservingForReply: 1_024) + #expect(fit == .tooLarge) + #expect(fake.streamedRequests.isEmpty) + } + + @Test("A reply that would only fit by eating the transcript margin does not fit") + func marginCounts() async { + let fake = FakeOnDeviceModel(contextSize: 2_000) { _ in 2_000 - 1_024 - OnDeviceTokenBudget.transcriptMargin + 1 } + #expect(await fake.fit([small], reservingForReply: 1_024) == .tooLarge) + } + + @Test("The estimate rounds up and adds the session's overhead") + func estimate() { + let request = OnDeviceRequest(instructions: "abc", prompt: "de") + #expect(OnDeviceTokenBudget.estimatedTokens(in: request) == 3 + OnDeviceTokenBudget.estimateOverhead) + } + + @Test("Text outside ASCII counts a token a character, above what the model measures") + func estimateNonASCII() { + let request = OnDeviceRequest(instructions: "", prompt: "注文テーブル") + #expect(OnDeviceTokenBudget.estimatedTokens(in: request) == 6 + OnDeviceTokenBudget.estimateOverhead) + } +} + +struct OnDeviceSnapshotTests { + @Test("A snapshot that extends the text yields only what is new") + func plainDelta() { + #expect(OnDeviceSnapshot.delta(from: "SELECT", to: "SELECT id") == " id") + #expect(OnDeviceSnapshot.delta(from: "", to: "SELECT") == "SELECT") + #expect(OnDeviceSnapshot.delta(from: "SELECT", to: "SELECT") == "") + } + + @Test("A snapshot that grows the last grapheme still yields the new scalars, and they join back up") + func extendsLastGrapheme() throws { + let emitted = "Done \u{1F469}" + let content = "Done \u{1F469}\u{200D}\u{1F4BB} ok" + let delta = try #require(OnDeviceSnapshot.delta(from: emitted, to: content)) + #expect(emitted + delta == content) + + let accented = try #require(OnDeviceSnapshot.delta(from: "cafe", to: "cafe\u{301}!")) + #expect("cafe" + accented == "caf\u{E9}!") + } + + @Test("A snapshot that rewrites earlier text is not a delta") + func rewriteIsRejected() { + #expect(OnDeviceSnapshot.delta(from: "SELECT id", to: "SELECT name") == nil) + } +} + +struct AppleIntelligenceAvailabilityTests { + @Test("Every reason the model is unavailable explains itself", arguments: [ + AppleIntelligenceAvailability.Reason.requiresNewerMacOS, + .deviceNotEligible, + .appleIntelligenceNotEnabled, + .modelNotReady, + .unknown + ]) + func unavailableExplains(reason: AppleIntelligenceAvailability.Reason) { + let availability = AppleIntelligenceAvailability.unavailable(reason) + #expect(!availability.isAvailable) + #expect(availability.unavailableExplanation?.isEmpty == false) + } + + @Test("An available model has nothing to explain") + func availableExplainsNothing() { + #expect(AppleIntelligenceAvailability.available.isAvailable) + #expect(AppleIntelligenceAvailability.available.unavailableExplanation == nil) + } +} + +struct QueryContextUsedColumnsTests { + private func snapshot() -> QueryContextSnapshot { + let orders = QueryContextTableStructure( + columns: [ + QueryContextColumn(name: "id", dataType: "bigint", isNullable: false, isPrimaryKey: true), + QueryContextColumn(name: "customer_id", dataType: "bigint", isNullable: false, isPrimaryKey: false), + QueryContextColumn(name: "Total", dataType: "numeric", isNullable: false, isPrimaryKey: false), + QueryContextColumn(name: "notes", dataType: "text", isNullable: true, isPrimaryKey: false) + ], + indexes: [ + QueryContextIndex(name: "orders_pkey", columns: ["id"], isUnique: true, isPrimary: true), + QueryContextIndex(name: "orders_customer_idx", columns: ["customer_id"], isUnique: false, isPrimary: false) + ], + foreignKeys: [ + QueryContextForeignKey(name: "fk_customer", columns: ["customer_id"], referencedTable: "customers", referencedColumns: ["id"]) + ], + approximateRowCount: 1_000 + ) + return QueryContextSnapshot( + engineName: "PostgreSQL", + tables: [ + QueryContextTable(name: "orders", schema: "public", kind: .table, content: .described(orders)), + QueryContextTable(name: "audit", schema: "public", kind: .table, content: .unavailable(reason: "denied")) + ] + ) + } + + @Test("Only the columns the statement names stay, matched without case, and the keys and indexes stay whole") + func keepsUsedColumns() throws { + let trimmed = snapshot().keepingColumns( + usedBy: "SELECT customer_id, total FROM orders WHERE customer_id = 7", + grammar: .ansi + ) + guard case .described(let orders) = try #require(trimmed.tables.first).content else { + Issue.record("orders lost its structure") + return + } + #expect(orders.columns.map(\.name) == ["customer_id", "Total"]) + #expect(orders.omittedColumnCount == 2) + #expect(orders.indexes.map(\.name) == ["orders_pkey", "orders_customer_idx"]) + #expect(orders.foreignKeys.map(\.name) == ["fk_customer"]) + #expect(orders.approximateRowCount == 1_000) + #expect(trimmed.tables.last?.content == .unavailable(reason: "denied")) + } + + @Test("The trimmed table says columns were left out rather than that it has none") + func rendersOmittedColumns() { + let some = QueryContextRenderer.render(snapshot().keepingColumns(usedBy: "SELECT total FROM orders", grammar: .ansi)) + #expect(some.contains("| Total |")) + #expect(!some.contains("| notes |")) + #expect(some.contains("… and 3 more columns not listed.")) + + let none = QueryContextRenderer.render(snapshot().keepingColumns(usedBy: "SELECT count(*) FROM orders", grammar: .ansi)) + #expect(none.contains("None of its 4 columns are named in the statement")) + #expect(!none.contains("No columns reported.")) + } + + @Test("A quoted name is matched whole in the connection's own quoting") + func quotedNamesMatchWhole() throws { + let structure = QueryContextTableStructure(columns: [ + QueryContextColumn(name: "order-id", dataType: "int", isNullable: false, isPrimaryKey: true), + QueryContextColumn(name: "order", dataType: "int", isNullable: false, isPrimaryKey: false) + ]) + let snapshot = QueryContextSnapshot( + engineName: "MySQL", + tables: [QueryContextTable(name: "orders", schema: nil, kind: .table, content: .described(structure))] + ) + + let trimmed = snapshot.keepingColumns(usedBy: "SELECT `order-id` FROM orders", grammar: TestGrammar.mysql) + + guard case .described(let orders) = try #require(trimmed.tables.first).content else { + Issue.record("orders lost its structure") + return + } + #expect(orders.columns.map(\.name) == ["order-id"]) + } + + @Test("A statement that names every column keeps the snapshot as it was") + func everyColumnUsed() { + let full = snapshot() + let trimmed = full.keepingColumns(usedBy: "SELECT id, customer_id, total, notes FROM orders", grammar: .ansi) + #expect(trimmed == full) + } +} diff --git a/apple/TableProTests/Helpers/FakeOnDeviceModel.swift b/apple/TableProTests/Helpers/FakeOnDeviceModel.swift new file mode 100644 index 0000000000..7b2ca1472f --- /dev/null +++ b/apple/TableProTests/Helpers/FakeOnDeviceModel.swift @@ -0,0 +1,54 @@ +// +// FakeOnDeviceModel.swift +// TableProTests +// + +import Foundation +@testable import TablePro + +final class FakeOnDeviceModel: OnDeviceLanguageModel, @unchecked Sendable { + let availability: AppleIntelligenceAvailability + let contextSize: Int + + private let lock = NSLock() + private let reply: [String] + private let tokens: @Sendable (OnDeviceRequest) -> Int + private var counted: [OnDeviceRequest] = [] + private var streamed: [(request: OnDeviceRequest, maximumResponseTokens: Int)] = [] + + init( + availability: AppleIntelligenceAvailability = .available, + contextSize: Int = 4_096, + reply: [String] = ["Looks ", "fine."], + tokens: @escaping @Sendable (OnDeviceRequest) -> Int = { ($0.instructions.count + $0.prompt.count) / 4 } + ) { + self.availability = availability + self.contextSize = contextSize + self.reply = reply + self.tokens = tokens + } + + var countedRequests: [OnDeviceRequest] { + lock.withLock { counted } + } + + var streamedRequests: [(request: OnDeviceRequest, maximumResponseTokens: Int)] { + lock.withLock { streamed } + } + + func tokenCount(for request: OnDeviceRequest) async -> Int { + lock.withLock { counted.append(request) } + return tokens(request) + } + + func streamResponse(to request: OnDeviceRequest, maximumResponseTokens: Int) -> AsyncThrowingStream { + lock.withLock { streamed.append((request, maximumResponseTokens)) } + let chunks = reply + return AsyncThrowingStream { continuation in + for chunk in chunks { + continuation.yield(chunk) + } + continuation.finish() + } + } +} diff --git a/apple/TableProTests/Models/AISettingsTests.swift b/apple/TableProTests/Models/AISettingsTests.swift index 114130ac1b..397b69d356 100644 --- a/apple/TableProTests/Models/AISettingsTests.swift +++ b/apple/TableProTests/Models/AISettingsTests.swift @@ -94,6 +94,17 @@ struct AISettingsTests { #expect(settings.maxToolRoundtripsEnabled == false) } + @Test("Apple Intelligence is off by default and in settings saved before it existed, and a stored choice survives") + func appleIntelligenceDecoding() throws { + #expect(AISettings.default.useAppleIntelligence == false) + let legacy = try JSONDecoder().decode(AISettings.self, from: Data("{}".utf8)) + #expect(legacy.useAppleIntelligence == false) + + let on = AISettings(useAppleIntelligence: true) + let roundTripped = try JSONDecoder().decode(AISettings.self, from: JSONEncoder().encode(on)) + #expect(roundTripped.useAppleIntelligence == true) + } + @Test("Turning the limit off makes it unlimited") func disabledLimitIsUnlimited() { let settings = AISettings(maxToolRoundtrips: 30, maxToolRoundtripsEnabled: false) diff --git a/apple/TableProTests/ViewModels/AIChatViewModelOnDeviceTests.swift b/apple/TableProTests/ViewModels/AIChatViewModelOnDeviceTests.swift new file mode 100644 index 0000000000..9d0c7ee50d --- /dev/null +++ b/apple/TableProTests/ViewModels/AIChatViewModelOnDeviceTests.swift @@ -0,0 +1,184 @@ +// +// AIChatViewModelOnDeviceTests.swift +// TableProTests +// + +import Foundation +@testable import TablePro +import TableProPluginKit +import Testing + +extension GlobalStateSuites { + @Suite(.appSettings { + $0.ai.enabled = true + $0.ai.providers = [] + $0.ai.activeProviderID = nil + $0.ai.useAppleIntelligence = true + $0.ai.includeSchema = false + $0.ai.defaultConnectionPolicy = .alwaysAllow + }) + @MainActor + struct AIChatViewModelOnDeviceTests { + private static let noProvider = String(localized: "No AI provider configured. Go to Settings > AI to add one.") + + private func makeViewModel(_ model: FakeOnDeviceModel) -> AIChatViewModel { + let viewModel = AIChatViewModel() + viewModel.connection = TestFixtures.makeConnection(type: .mysql) + viewModel.onDeviceModel = model + return viewModel + } + + private func request(_ statement: String = "SELECT id FROM orders") -> AIQueryRequest { + AIQueryRequest(action: .review, statement: statement, scope: nil, databaseType: .mysql) + } + + private func roles(_ viewModel: AIChatViewModel) -> [ChatRole] { + viewModel.messages.map { $0.role } + } + + @Test("A query action is answered on this Mac, with no provider set up") + func answersOnDevice() async throws { + let model = FakeOnDeviceModel() + let viewModel = makeViewModel(model) + + viewModel.runQueryAction(request()) + await viewModel.streamingTask?.value + + let reply = try #require(viewModel.messages.last) + #expect(roles(viewModel) == [.user, .assistant]) + #expect(reply.plainText == "Looks fine.") + #expect(reply.modelId == AppleIntelligenceAvailability.displayName) + #expect(viewModel.errorMessage == nil) + #expect(!viewModel.isStreaming) + let sent = try #require(model.streamedRequests.first) + #expect(sent.request.prompt.contains("SELECT id FROM orders")) + #expect(sent.request.instructions.contains("MySQL")) + #expect(!sent.request.instructions.contains(WalkthroughEnvelopeParser.openFence)) + } + + @Test("A request too large is refused before generation, keeps the action and offers the provider") + func tooLargeIsRefused() async throws { + let model = FakeOnDeviceModel(contextSize: 1_500) { _ in 1_000 } + let viewModel = makeViewModel(model) + + viewModel.runQueryAction(request()) + await viewModel.streamingTask?.value + + #expect(model.streamedRequests.isEmpty) + #expect(roles(viewModel) == [.user]) + #expect(viewModel.errorMessage?.contains("Apple Intelligence") == true) + #expect(viewModel.onDeviceRefusal?.turnID == viewModel.messages.last?.id) + #expect(!viewModel.isStreaming) + } + + @Test("Sending a refused action on answers the same turn through the provider path") + func refusedTurnGoesToProvider() async throws { + let model = FakeOnDeviceModel(contextSize: 1_500) { _ in 1_000 } + let viewModel = makeViewModel(model) + viewModel.runQueryAction(request()) + await viewModel.streamingTask?.value + + viewModel.sendRefusedTurnToProvider() + + #expect(model.countedRequests.count == 1) + #expect(viewModel.onDeviceTurnID == nil) + #expect(viewModel.onDeviceRefusal == nil) + #expect(viewModel.errorMessage == Self.noProvider) + } + + @Test("A follow-up in the same conversation goes to the provider, never the device") + func followUpLeavesTheDevice() async throws { + let model = FakeOnDeviceModel() + let viewModel = makeViewModel(model) + viewModel.runQueryAction(request()) + await viewModel.streamingTask?.value + + viewModel.inputText = "Why is that?" + viewModel.sendMessage() + + #expect(model.streamedRequests.count == 1) + #expect(viewModel.errorMessage == Self.noProvider) + } + + @Test("Regenerating an on-device answer asks the device again") + func regenerateStaysOnDevice() async throws { + let model = FakeOnDeviceModel() + let viewModel = makeViewModel(model) + viewModel.runQueryAction(request()) + await viewModel.streamingTask?.value + + viewModel.regenerate() + await viewModel.streamingTask?.value + + #expect(model.streamedRequests.count == 2) + #expect(roles(viewModel) == [.user, .assistant]) + } + + @Test("Regenerating a restored on-device answer stays on the device") + func regenerateRestoredAnswerStaysOnDevice() async throws { + let model = FakeOnDeviceModel() + let viewModel = makeViewModel(model) + viewModel.messages = [ + ChatTurn(role: .user, blocks: [.text("Review this SQL query.")]), + ChatTurn( + role: .assistant, + blocks: [.text("Looks fine.")], + modelId: AppleIntelligenceAvailability.displayName, + providerId: AppleIntelligenceAvailability.providerID + ) + ] + + viewModel.regenerate() + await viewModel.streamingTask?.value + + #expect(model.streamedRequests.count == 1) + #expect(viewModel.errorMessage == nil) + #expect(roles(viewModel) == [.user, .assistant]) + } + + @Test("With the model unavailable the action says why and sends nothing") + func unavailableExplains() { + let unavailable = AppleIntelligenceAvailability.unavailable(.appleIntelligenceNotEnabled) + let model = FakeOnDeviceModel(availability: unavailable) + let viewModel = makeViewModel(model) + + viewModel.runQueryAction(request()) + + #expect(viewModel.messages.isEmpty) + #expect(viewModel.errorMessage == unavailable.unavailableExplanation) + #expect(model.countedRequests.isEmpty) + } + + @Test("Ask Each Time asks before the device reads anything", .appSettings { + $0.ai.defaultConnectionPolicy = .askEachTime + }) + func askEachTimeAsksFirst() async throws { + let model = FakeOnDeviceModel() + let viewModel = makeViewModel(model) + + viewModel.runQueryAction(request()) + + #expect(viewModel.showAIAccessConfirmation) + #expect(model.countedRequests.isEmpty) + + viewModel.confirmAIAccess() + await viewModel.streamingTask?.value + + #expect(model.streamedRequests.count == 1) + #expect(roles(viewModel) == [.user, .assistant]) + } + + @Test("With the setting off the action goes to the provider path", .appSettings { + $0.ai.useAppleIntelligence = false + }) + func settingOffUsesProvider() { + let model = FakeOnDeviceModel() + let viewModel = makeViewModel(model) + + viewModel.runQueryAction(request()) + + #expect(model.countedRequests.isEmpty) + #expect(viewModel.errorMessage == Self.noProvider) + } + } +} diff --git a/apple/changelog.d/3749.added.md b/apple/changelog.d/3749.added.md new file mode 100644 index 0000000000..3e3af6fb77 --- /dev/null +++ b/apple/changelog.d/3749.added.md @@ -0,0 +1 @@ +Inline suggestions and Review, Explain, Optimize and Fix can run on Apple Intelligence, on the Mac itself. diff --git a/apple/scripts/probes/check-on-device-model.sh b/apple/scripts/probes/check-on-device-model.sh new file mode 100755 index 0000000000..05c5bdef3d --- /dev/null +++ b/apple/scripts/probes/check-on-device-model.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# +# Run the app's on-device model code (TablePro/Core/AI/OnDevice) against the Apple Intelligence model +# on this Mac. Prints availability, the context size, characters per token for SQL, a structure +# table and CJK text (OnDeviceTokenBudget.estimatedTokens rests on these), whether an oversized +# request is refused before generation, and a short streamed reply. +# +# Usage: +# scripts/probes/check-on-device-model.sh +# +# Needs macOS 26 or later with Apple Intelligence on and its model downloaded, and xcrun swiftc. +# Exits 0 when a reply streams, 1 when it does not, 3 when the model is unavailable. + +set -uo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +WORK="$(mktemp -d)" +trap 'rm -rf "$WORK"' EXIT + +xcrun swiftc -parse-as-library -target arm64-apple-macos13 -O \ + "$ROOT"/TablePro/Core/AI/OnDevice/*.swift \ + "$ROOT"/scripts/probes/on-device-model/main.swift \ + -o "$WORK/probe" || exit 3 + +"$WORK/probe" diff --git a/apple/scripts/probes/on-device-model/main.swift b/apple/scripts/probes/on-device-model/main.swift new file mode 100644 index 0000000000..96db1bd3dc --- /dev/null +++ b/apple/scripts/probes/on-device-model/main.swift @@ -0,0 +1,59 @@ +import Foundation + +extension Error { + var publicLogShape: String { String(describing: type(of: self)) } +} + +@main +struct OnDeviceModelProbe { + static func main() async { + let model = SystemOnDeviceLanguageModel() + print("availability: \(model.availability)") + guard model.availability.isAvailable else { exit(3) } + print("contextSize: \(model.contextSize)") + + let samples = [ + ("sql", "SELECT o.id, o.total, c.name FROM orders o JOIN customers c ON c.id = o.customer_id WHERE o.total > 100;"), + ("structure", """ + | Column | Type | Nullable | Key | Default | Comment | + | --- | --- | --- | --- | --- | --- | + | id | bigint | no | PK | | | + | customer_id | bigint | no | | | | + | created_at | timestamp with time zone | no | | now() | | + """), + ("cjk", "注文テーブルの合計金額を顧客ごとに集計してください。") + ] + let baseline = await model.tokenCount(for: OnDeviceRequest(instructions: "", prompt: "")) + for (label, text) in samples { + let request = OnDeviceRequest(instructions: "", prompt: text) + let tokens = await model.tokenCount(for: request) - baseline + let characters = (text as NSString).length + let ratio = Double(characters) / Double(max(tokens, 1)) + print("\(label): \(characters) characters, \(tokens) tokens, \(String(format: "%.2f", ratio)) per token") + } + + let request = OnDeviceRequest( + instructions: "You are the database assistant in TablePro. Answer briefly.", + prompt: "Explain this SQL query.\n\n```sql\n\(samples[0].1)\n```" + ) + let oversized = OnDeviceRequest( + instructions: request.instructions, + prompt: request.prompt + String(repeating: "\n| col | int | yes | | | |", count: 2_000) + ) + let fit = await model.fit([oversized, request], reservingForReply: 1_024) + print("fit: \(fit)") + guard case .fits(let index, let allowance) = fit, index == 1 else { exit(1) } + + var reply = "" + do { + for try await delta in model.streamResponse(to: request, maximumResponseTokens: min(allowance, 120)) { + reply += delta + } + } catch { + print("stream failed: \(error.localizedDescription)") + exit(1) + } + print("reply: \(reply)") + exit(reply.isEmpty ? 1 : 0) + } +}