Add a local spelling dictionary

Package.swift

@@ -16,5 +16,9 @@ let package = Package(
                 .product(name: "WhisperKit", package: "WhisperKit"),
             ]
         ),
+        .testTarget(
+            name: "parrotTests",
+            dependencies: ["parrot"]
+        ),
     ]
 )

README.md

@@ -40,6 +40,25 @@ parrot --hotkey right-option           # change the push-to-talk key
 parrot --no-overlay                    # disable the bottom-of-screen pill
 ```
 
+## Local dictionary
+
+On its first run, Parrot creates `~/.config/parrot/dictionary.json`. It never
+leaves the Mac and normalizes spelling variants after Whisper has transcribed a
+dictation. Add your own terms by editing the file:
+
+```json
+{
+  "entries": [
+    { "canonical": "AcmeAPI", "variants": ["acme api", "acme A.P.I."] }
+  ]
+}
+```
+
+The dictionary is applied at the end of transcription. Saved edits are picked
+up on the next dictation; no Parrot restart is needed. The easiest way to add
+a correction is from the menu-bar icon: **Add dictionary correction…**. Enter
+what Parrot wrote, then the spelling you want it to use.
+
 ## Stack
 
 - **Swift** — single SPM executable target

Sources/parrot/Dictionary/LocalDictionary.swift

@@ -0,0 +1,160 @@
+import Foundation
+
+/// A deliberately small, local-only spelling dictionary.
+///
+/// Whisper stays responsible for speech recognition. This layer only normalizes
+/// known variants after transcription, so a correction is deterministic and
+/// survives restarts without sending vocabulary or transcripts anywhere.
+struct LocalDictionary: Codable {
+    struct Entry: Codable, Equatable {
+        let canonical: String
+        let variants: [String]
+    }
+
+    let entries: [Entry]
+
+    static let defaultURL: URL = FileManager.default.homeDirectoryForCurrentUser
+        .appendingPathComponent(".config/parrot/dictionary.json")
+
+    static let starterEntries: [Entry] = []
+
+    static func loadOrCreateDefault() -> LocalDictionary {
+        let fileManager = FileManager.default
+        if !fileManager.fileExists(atPath: defaultURL.path) {
+            do {
+                try fileManager.createDirectory(
+                    at: defaultURL.deletingLastPathComponent(),
+                    withIntermediateDirectories: true
+                )
+                let starter = LocalDictionary(entries: starterEntries)
+                let data = try JSONEncoder.pretty.encode(starter)
+                try data.write(to: defaultURL, options: .atomic)
+                log("created local dictionary at \(defaultURL.path)")
+                return starter
+            } catch {
+                log("could not create local dictionary: \(error)")
+                return LocalDictionary(entries: [])
+            }
+        }
+
+        do {
+            let data = try Data(contentsOf: defaultURL)
+            return try JSONDecoder().decode(LocalDictionary.self, from: data)
+        } catch {
+            log("could not load local dictionary at \(defaultURL.path): \(error)")
+            return LocalDictionary(entries: [])
+        }
+    }
+
+    /// Adds one user-facing correction while keeping the on-disk file valid.
+    /// If the same mis-transcription already pointed at another spelling, the
+    /// new choice replaces it instead of leaving two competing rules behind.
+    static func addCorrection(transcribedAs: String, correctSpelling: String) throws {
+        let source = transcribedAs.trimmingCharacters(in: .whitespacesAndNewlines)
+        let canonical = correctSpelling.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard !source.isEmpty, !canonical.isEmpty else {
+            throw DictionaryError.emptyCorrection
+        }
+
+        let updated = loadOrCreateDefault().addingCorrection(
+            transcribedAs: source,
+            correctSpelling: canonical
+        )
+        let data = try JSONEncoder.pretty.encode(updated)
+        try data.write(to: defaultURL, options: .atomic)
+    }
+
+    func addingCorrection(transcribedAs: String, correctSpelling: String) -> LocalDictionary {
+        let sourceKey = transcribedAs.folding(
+            options: [.caseInsensitive, .diacriticInsensitive],
+            locale: .current
+        )
+        let canonicalKey = correctSpelling.folding(
+            options: [.caseInsensitive, .diacriticInsensitive],
+            locale: .current
+        )
+
+        var updatedEntries = entries.map { entry in
+            Entry(
+                canonical: entry.canonical,
+                variants: entry.variants.filter {
+                    $0.folding(options: [.caseInsensitive, .diacriticInsensitive], locale: .current) != sourceKey
+                }
+            )
+        }
+
+        if let index = updatedEntries.firstIndex(where: {
+            $0.canonical.folding(options: [.caseInsensitive, .diacriticInsensitive], locale: .current) == canonicalKey
+        }) {
+            var variants = updatedEntries[index].variants
+            if sourceKey != canonicalKey, !variants.contains(where: {
+                $0.folding(options: [.caseInsensitive, .diacriticInsensitive], locale: .current) == sourceKey
+            }) {
+                variants.append(transcribedAs)
+            }
+            updatedEntries[index] = Entry(canonical: correctSpelling, variants: variants)
+        } else {
+            updatedEntries.append(
+                Entry(
+                    canonical: correctSpelling,
+                    variants: sourceKey == canonicalKey ? [] : [transcribedAs]
+                )
+            )
+        }
+
+        return LocalDictionary(entries: updatedEntries)
+    }
+
+    func apply(to text: String) -> String {
+        let replacements = entries.flatMap { entry in
+            ([entry.canonical] + entry.variants).map { variant in
+                (variant: variant, canonical: entry.canonical)
+            }
+        }
+        .sorted { $0.variant.count > $1.variant.count }
+
+        return replacements.reduce(text) { output, replacement in
+            replaceWholeWord(
+                replacement.variant,
+                with: replacement.canonical,
+                in: output
+            )
+        }
+    }
+
+    private func replaceWholeWord(_ source: String, with replacement: String, in text: String) -> String {
+        let escapedSource = NSRegularExpression.escapedPattern(for: source)
+        let pattern = "(?i)(?<![\\p{L}\\p{N}_])\(escapedSource)(?![\\p{L}\\p{N}_])"
+        guard let regex = try? NSRegularExpression(pattern: pattern) else { return text }
+        let range = NSRange(text.startIndex..<text.endIndex, in: text)
+        let escapedReplacement = NSRegularExpression.escapedTemplate(for: replacement)
+        return regex.stringByReplacingMatches(
+            in: text,
+            range: range,
+            withTemplate: escapedReplacement
+        )
+    }
+
+    private static func log(_ message: String) {
+        FileHandle.standardError.write(Data("parrot dictionary: \(message)\n".utf8))
+    }
+}
+
+enum DictionaryError: LocalizedError {
+    case emptyCorrection
+
+    var errorDescription: String? {
+        switch self {
+        case .emptyCorrection:
+            return "Both dictionary fields need a value."
+        }
+    }
+}
+
+private extension JSONEncoder {
+    static var pretty: JSONEncoder {
+        let encoder = JSONEncoder()
+        encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
+        return encoder
+    }
+}

Sources/parrot/Transcription/WhisperKitTranscriber.swift

@@ -9,6 +9,7 @@ actor WhisperKitTranscriber: Transcriber {
     init(model: TranscriptionModel) {
         self.modelID = model.id
         self.model = model
+        _ = LocalDictionary.loadOrCreateDefault()
     }
 
     /// Loads the model into memory; downloads first if not already on disk.
@@ -31,7 +32,9 @@ actor WhisperKitTranscriber: Transcriber {
 
         let results = try await pipeline.transcribe(audioArray: audio)
         let raw = results.map(\.text).joined(separator: " ")
-        return Self.sanitize(raw)
+        // This file is tiny; reloading it lets a saved edit take effect for the
+        // very next dictation without restarting the daemon.
+        return LocalDictionary.loadOrCreateDefault().apply(to: Self.sanitize(raw))
     }
 
     /// Strip Whisper's non-speech bracket tokens ([BLANK_AUDIO], [MUSIC],

Sources/parrot/UI/MenuBarController.swift

@@ -27,6 +27,16 @@ final class MenuBarController {
 
         menu.addItem(.separator())
 
+        let dictionary = NSMenuItem(
+            title: "Add dictionary correction…",
+            action: #selector(addDictionaryCorrectionClicked),
+            keyEquivalent: ""
+        )
+        dictionary.target = self
+        menu.addItem(dictionary)
+
+        menu.addItem(.separator())
+
         let quit = NSMenuItem(
             title: "Quit parrot",
             action: #selector(quitClicked),
@@ -81,4 +91,74 @@ final class MenuBarController {
     @objc private func quitClicked() {
         NSApp.terminate(nil)
     }
+
+    @objc private func addDictionaryCorrectionClicked() {
+        let transcribedAs = NSTextField(string: "")
+        transcribedAs.placeholderString = "e.g. acme api"
+        transcribedAs.controlSize = .regular
+
+        let correctSpelling = NSTextField(string: "")
+        correctSpelling.placeholderString = "e.g. AcmeAPI"
+        correctSpelling.controlSize = .regular
+
+        let alert = NSAlert()
+        alert.messageText = "Add dictionary correction"
+        alert.informativeText = "Enter what Parrot wrote, then your preferred spelling."
+        alert.addButton(withTitle: "Add correction")
+        alert.addButton(withTitle: "Cancel")
+        alert.accessoryView = dictionaryForm(
+            transcribedAs: transcribedAs,
+            correctSpelling: correctSpelling
+        )
+        alert.window.initialFirstResponder = transcribedAs
+
+        NSApp.activate(ignoringOtherApps: true)
+        guard alert.runModal() == .alertFirstButtonReturn else { return }
+
+        do {
+            try LocalDictionary.addCorrection(
+                transcribedAs: transcribedAs.stringValue,
+                correctSpelling: correctSpelling.stringValue
+            )
+        } catch {
+            NSAlert(error: error).runModal()
+        }
+    }
+
+    private func dictionaryForm(
+        transcribedAs: NSTextField,
+        correctSpelling: NSTextField
+    ) -> NSView {
+        let stack = NSStackView(views: [
+            fieldGroup(title: "Parrot wrote", field: transcribedAs),
+            fieldGroup(title: "Use this spelling", field: correctSpelling),
+        ])
+        stack.orientation = .vertical
+        stack.alignment = .leading
+        stack.spacing = 12
+
+        // NSAlert sizes an accessory view from its frame rather than resolving
+        // external constraints. Give it a concrete, pre-laid-out size so the
+        // form stays inside the alert on every supported macOS version.
+        stack.layoutSubtreeIfNeeded()
+        let form = NSView(frame: NSRect(origin: .zero, size: stack.fittingSize))
+        stack.frame = form.bounds
+        stack.autoresizingMask = [.width, .height]
+        form.addSubview(stack)
+        return form
+    }
+
+    private func fieldGroup(title: String, field: NSTextField) -> NSStackView {
+        let label = NSTextField(labelWithString: title)
+        label.font = .systemFont(ofSize: NSFont.smallSystemFontSize, weight: .medium)
+
+        field.translatesAutoresizingMaskIntoConstraints = false
+        field.widthAnchor.constraint(equalToConstant: 360).isActive = true
+
+        let group = NSStackView(views: [label, field])
+        group.orientation = .vertical
+        group.alignment = .leading
+        group.spacing = 4
+        return group
+    }
 }

Tests/parrotTests/LocalDictionaryTests.swift

@@ -0,0 +1,45 @@
+import XCTest
+@testable import parrot
+
+final class LocalDictionaryTests: XCTestCase {
+    func testCanonicalizesVariantsWithoutChangingLongerWords() {
+        let dictionary = LocalDictionary(entries: [
+            .init(canonical: "AcmeAPI", variants: ["acme api", "acme A.P.I."]),
+            .init(canonical: "ExampleDB", variants: ["example db"]),
+        ])
+
+        XCTAssertEqual(
+            dictionary.apply(to: "acme api et EXAMPLE DB, mais acme apis reste inchangé."),
+            "AcmeAPI et ExampleDB, mais acme apis reste inchangé."
+        )
+    }
+
+    func testAddingCorrectionReplacesAnExistingConflictingVariant() {
+        let dictionary = LocalDictionary(entries: [
+            .init(canonical: "Old spelling", variants: ["parrot word"]),
+            .init(canonical: "AcmeAPI", variants: ["acme api"]),
+        ])
+
+        let updated = dictionary.addingCorrection(
+            transcribedAs: "parrot word",
+            correctSpelling: "Preferred spelling"
+        )
+
+        XCTAssertEqual(
+            updated.apply(to: "PARROT WORD puis acme api"),
+            "Preferred spelling puis AcmeAPI"
+        )
+        XCTAssertEqual(
+            updated.entries.first(where: { $0.canonical == "Old spelling" })?.variants,
+            []
+        )
+    }
+
+    func testCanonicalSpellingIsUsedLiterallyAsReplacementText() {
+        let dictionary = LocalDictionary(entries: [
+            .init(canonical: "C$", variants: ["c dollar"]),
+        ])
+
+        XCTAssertEqual(dictionary.apply(to: "c dollar"), "C$")
+    }
+}