Replace Espresso with a hardened ANE-LM backend

Package.swift

@@ -1,4 +1,4 @@
-// swift-tools-version: 6.0
+// swift-tools-version: 6.2
 import PackageDescription
 
 let package = Package(
@@ -13,6 +13,10 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/argmaxinc/argmax-oss-swift.git", from: "1.0.0"),
+        .package(
+            url: "https://github.com/IchenDEV/ANE-LM.git",
+            revision: "033472ec12ea796fc7ea4f8cefd7ed456f69900b"
+        ),
         .package(url: "https://github.com/Blaizzy/mlx-audio-swift.git", exact: "0.1.3"),
         .package(url: "https://github.com/huggingface/swift-transformers", from: "1.3.3"),
         .package(url: "https://github.com/ml-explore/mlx-swift-lm", exact: "3.31.4"),
@@ -22,6 +26,7 @@ let package = Package(
             name: "OpenType",
             dependencies: [
                 .product(name: "WhisperKit", package: "argmax-oss-swift"),
+                .product(name: "ANELMRuntime", package: "ANE-LM"),
                 .product(name: "MLXAudioCore", package: "mlx-audio-swift"),
                 .product(name: "MLXAudioSTT", package: "mlx-audio-swift"),
                 .product(name: "Hub", package: "swift-transformers"),

Sources/App/AppDelegate+Integrations.swift

@@ -117,6 +117,7 @@ extension AppDelegate {
     func makeIntegrationSessionCoordinator(service: OpenTypeService) -> InputSessionCoordinator {
         InputSessionCoordinator(
             service: service,
+            textProcessor: textProcessor,
             isUserWorkflowBusy: { [weak self] in
                 self?.appState.isBusy ?? false
             }

Sources/App/AppState.swift

@@ -13,6 +13,11 @@ enum AppPhase: Equatable {
     case error(String)
 }
 
+enum AppCompletionKind: Equatable {
+    case standard
+    case espressoFallback
+}
+
 @MainActor
 final class AppState: ObservableObject {
     @Published var phase: AppPhase = .idle
@@ -32,6 +37,7 @@ final class AppState: ObservableObject {
     @Published var lastFormattingDurationSeconds: Double = 0
     @Published var pendingReplacement: DeferredReplacement?
     @Published var activeInputMode: VoiceInputMode = .dictation
+    @Published var completionKind: AppCompletionKind = .standard
 
     let settings = AppSettings.shared
 
@@ -54,6 +60,7 @@ final class AppState: ObservableObject {
         resetDownloadProgress()
         pendingReplacement = nil
         activeInputMode = .dictation
+        completionKind = .standard
     }
 
     func clearPendingReplacement() {

Sources/App/OpenTypeApp.swift

@@ -24,6 +24,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
     private var settingsWindowDelegate: SettingsWindowDelegate?
     private var onboardingWindow: NSWindow?
     private let popoverOutsideClickMonitor = PopoverOutsideClickMonitor()
+    let textProcessor: TextProcessor
     var cancellables = Set<AnyCancellable>()
     var iconTimer: Timer?
     private var previousApp: NSRunningApplication?
@@ -39,6 +40,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
         let registry = IntegrationClientRegistry()
         integrationClientRegistry = registry
         integrationService = OpenTypeService(registry: registry)
+        textProcessor = TextProcessor()
         super.init()
         integrationSessionCoordinator = makeIntegrationSessionCoordinator(service: integrationService)
     }
@@ -95,7 +97,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
     }
 
     private func setupPipeline() {
-        pipeline = VoicePipeline(appState: appState)
+        pipeline = VoicePipeline(appState: appState, textProcessor: textProcessor)
         Task { await pipeline?.warmUp() }
     }
 
@@ -202,6 +204,10 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
             onUnloadWhisper: { [weak self] in self?.pipeline?.unloadWhisper() },
             onUnloadLLM: { [weak self] in self?.pipeline?.unloadLLM() },
             onLoadLLM: { [weak self] in self?.pipeline?.loadLLM() },
+            onBenchmarkLLM: { [weak self] modelID in
+                guard let self else { throw CancellationError() }
+                return try await self.textProcessor.benchmarkLLM(modelID: modelID)
+            },
             onUnloadLocalASR: { [weak self] in self?.pipeline?.unloadLocalASR() }
         )
         .environmentObject(appState)

Sources/App/VoicePipeline+EditCommandHelpers.swift

@@ -0,0 +1,103 @@
+import AppKit
+import Foundation
+
+@MainActor
+extension VoicePipeline {
+    func replacementInputContext(
+        settings: AppSettings,
+        targetApp: NSRunningApplication?
+    ) -> InputContext {
+        InputContext.capture(
+            targetApp: targetApp,
+            screenContext: "",
+            outputMode: .command,
+            inputLanguage: settings.inputLanguage,
+            source: .menuBar
+        )
+    }
+
+    func finalizedReplacementText(
+        _ text: String,
+        settings: AppSettings
+    ) -> String {
+        textProcessor.cleanCommandGeneratedOutput(
+            text,
+            inputLanguage: settings.inputLanguage
+        )
+    }
+
+    func rewriteSelectedText(
+        raw: String,
+        intent: SelectionRewriteIntent,
+        settings: AppSettings,
+        targetApp: NSRunningApplication?
+    ) async {
+        cancelScreenContextCapture()
+
+        guard let selectedText = await textInserter.selectedText(targetApp: targetApp),
+              !selectedText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
+            showErrorHint(L("pipeline.no_selected_text_to_replace"))
+            return
+        }
+
+        let context = InputContext.capture(
+            targetApp: targetApp,
+            screenContext: "",
+            selectedTextOverride: selectedText,
+            outputMode: .command,
+            inputLanguage: settings.inputLanguage,
+            source: .menuBar
+        )
+        var options = TextProcessingOptions(settings: settings)
+        options.llmModel = settings.llmModel
+        let memoryContext = VoicePipelinePolicy.memoryContext(
+            for: .command,
+            settings: settings,
+            currentContext: context
+        )
+
+        appState.phase = .processing
+        appState.statusMessage = L("pipeline.formatting")
+        let rewrittenText = await textProcessor.processSelectionEdit(
+            selectedText: selectedText,
+            intent: intent,
+            options: options,
+            spokenCommand: raw,
+            memoryContext: memoryContext,
+            inputContext: context
+        )
+        guard !rewrittenText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
+            showNoSpeechDetected(reason: "selection rewrite returned empty text")
+            return
+        }
+
+        appState.processedText = rewrittenText
+        appState.phase = .inserting
+        appState.statusMessage = L("pipeline.replacing")
+
+        let result = await textInserter.replaceSelectedText(text: rewrittenText, targetApp: targetApp)
+        appState.phase = .done
+        appState.statusMessage = L("status.done")
+        hideOverlayAfterDelay()
+
+        if case .probablyFailed(let reason) = result {
+            Log.info("[VoicePipeline] voice edit selection rewrite probably failed: \(reason)")
+            TextInserter.copyToClipboard(rewrittenText)
+            showInsertionFailedAlert(text: rewrittenText, reason: reason)
+            return
+        }
+
+        let recordID = InputHistory.shared.addRecord(
+            rawText: raw,
+            processedText: rewrittenText,
+            wasProcessed: true,
+            context: context
+        )
+        appState.lastInsertedText = rewrittenText
+        beginCorrectionCapture(
+            recordID: recordID,
+            insertedText: rewrittenText,
+            context: context
+        )
+    }
+}

Sources/App/VoicePipeline+EditCommands.swift

@@ -8,6 +8,7 @@ extension VoicePipeline {
         settings: AppSettings,
         targetApp: NSRunningApplication?
     ) async -> Bool {
+        let expectedEspressoModelPath = settings.espressoModelPath
         guard let command = await resolvedSpokenEditCommand(
             raw: raw,
             settings: settings,
@@ -24,38 +25,50 @@ extension VoicePipeline {
                 settings: settings,
                 targetApp: targetApp
             )
-            return true
         case .replaceSelection(let replacementRaw):
             await replaceSelectedText(
                 raw: raw,
                 replacementRaw: replacementRaw,
                 settings: settings,
                 targetApp: targetApp
             )
-            return true
         case .rewriteLast(let intent):
             await rewriteLastInsertion(
                 raw: raw,
                 intent: intent,
                 settings: settings,
                 targetApp: targetApp
             )
-            return true
         case .rewriteSelection(let intent):
             await rewriteSelectedText(
                 raw: raw,
                 intent: intent,
                 settings: settings,
                 targetApp: targetApp
             )
-            return true
         case .deleteSelection:
             await deleteSelectedText(targetApp: targetApp)
-            return true
         case .undoLastInsertion:
             await undoLastInsertion(targetApp: targetApp)
-            return true
         }
+
+        guard !Task.isCancelled else { return true }
+        if let espressoOutcome = await consumeEspressoOutcome(
+            settings: settings,
+            expectedEspressoModelPath: expectedEspressoModelPath
+        ) {
+            if case .error = appState.phase, espressoOutcome == .fallback {
+                Log.info("[VoicePipeline] preserving edit-command error after ANE-LM fallback")
+            } else if espressoOutcome != .fallback {
+                showErrorHint(espressoOutcome.message)
+            } else {
+                appState.completionKind = .espressoFallback
+                appState.statusMessage = espressoOutcome.message
+                showOverlay()
+                hideOverlayAfterDelay()
+            }
+        }
+        return true
     }
 
     private func replaceLastInsertion(
@@ -161,29 +174,6 @@ extension VoicePipeline {
         )
     }
 
-    private func replacementInputContext(
-        settings: AppSettings,
-        targetApp: NSRunningApplication?
-    ) -> InputContext {
-        InputContext.capture(
-            targetApp: targetApp,
-            screenContext: "",
-            outputMode: .command,
-            inputLanguage: settings.inputLanguage,
-            source: .menuBar
-        )
-    }
-
-    private func finalizedReplacementText(
-        _ text: String,
-        settings: AppSettings
-    ) -> String {
-        textProcessor.cleanCommandGeneratedOutput(
-            text,
-            inputLanguage: settings.inputLanguage
-        )
-    }
-
     private func deleteSelectedText(targetApp: NSRunningApplication?) async {
         cancelScreenContextCapture()
 
@@ -236,78 +226,4 @@ extension VoicePipeline {
         hideOverlayAfterDelay()
     }
 
-    private func rewriteSelectedText(
-        raw: String,
-        intent: SelectionRewriteIntent,
-        settings: AppSettings,
-        targetApp: NSRunningApplication?
-    ) async {
-        cancelScreenContextCapture()
-
-        guard let selectedText = await textInserter.selectedText(targetApp: targetApp),
-              !selectedText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
-            showErrorHint(L("pipeline.no_selected_text_to_replace"))
-            return
-        }
-
-        let context = InputContext.capture(
-            targetApp: targetApp,
-            screenContext: "",
-            selectedTextOverride: selectedText,
-            outputMode: .command,
-            inputLanguage: settings.inputLanguage,
-            source: .menuBar
-        )
-        var options = TextProcessingOptions(settings: settings)
-        options.llmModel = settings.llmModel
-        let memoryContext = VoicePipelinePolicy.memoryContext(
-            for: .command,
-            settings: settings,
-            currentContext: context
-        )
-
-        appState.phase = .processing
-        appState.statusMessage = L("pipeline.formatting")
-        let rewrittenText = await textProcessor.processSelectionEdit(
-            selectedText: selectedText,
-            intent: intent,
-            options: options,
-            spokenCommand: raw,
-            memoryContext: memoryContext,
-            inputContext: context
-        )
-        guard !rewrittenText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
-            showNoSpeechDetected(reason: "selection rewrite returned empty text")
-            return
-        }
-
-        appState.processedText = rewrittenText
-        appState.phase = .inserting
-        appState.statusMessage = L("pipeline.replacing")
-
-        let result = await textInserter.replaceSelectedText(text: rewrittenText, targetApp: targetApp)
-        appState.phase = .done
-        appState.statusMessage = L("status.done")
-        hideOverlayAfterDelay()
-
-        if case .probablyFailed(let reason) = result {
-            Log.info("[VoicePipeline] voice edit selection rewrite probably failed: \(reason)")
-            TextInserter.copyToClipboard(rewrittenText)
-            showInsertionFailedAlert(text: rewrittenText, reason: reason)
-            return
-        }
-
-        let recordID = InputHistory.shared.addRecord(
-            rawText: raw,
-            processedText: rewrittenText,
-            wasProcessed: true,
-            context: context
-        )
-        appState.lastInsertedText = rewrittenText
-        beginCorrectionCapture(
-            recordID: recordID,
-            insertedText: rewrittenText,
-            context: context
-        )
-    }
 }