feat: controller supports vocabulary bias, prompt-echo guard, corrector and result hook

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

Sources/SstCore/DictationController.swift

@@ -7,6 +7,8 @@ public final class DictationController {
     }
     public var onStateChange: ((DictationState) -> Void)?
     public var onError: ((String) -> Void)?
+    /// Called with every finished transcription, just before it is inserted.
+    public var onResult: ((DictationResult) -> Void)?
 
     private let recorder: AudioRecording
     private let transcriber: Transcriber
@@ -14,16 +16,25 @@ public final class DictationController {
     private let inserter: TextInserting
     private let minimumSamples: Int
     private let silenceThreshold: Float
+    private let corrector: TextCorrecting?
+    private let vocabularyProvider: () -> [String]
+    private let biasMinimumRMS: Float
 
     /// - Parameter minimumSamples: 16 kHz samples; default is 0.3 s.
     /// - Parameter silenceThreshold: RMS below which a recording counts as silence and is skipped.
+    /// - Parameter corrector: optional dictionary correction applied to the transcript.
+    /// - Parameter vocabularyProvider: words to bias the transcriber towards.
+    /// - Parameter biasMinimumRMS: below this RMS no vocabulary is passed (guards against invented text on quiet audio).
     public init(
         recorder: AudioRecording,
         transcriber: Transcriber,
         cleaner: Cleaner,
         inserter: TextInserting,
         minimumSamples: Int = 4_800,
-        silenceThreshold: Float = 0.005
+        silenceThreshold: Float = 0.005,
+        corrector: TextCorrecting? = nil,
+        vocabularyProvider: @escaping () -> [String] = { [] },
+        biasMinimumRMS: Float = 0.015
     ) {
         self.recorder = recorder
         self.transcriber = transcriber
@@ -31,6 +42,9 @@ public final class DictationController {
         self.inserter = inserter
         self.minimumSamples = minimumSamples
         self.silenceThreshold = silenceThreshold
+        self.corrector = corrector
+        self.vocabularyProvider = vocabularyProvider
+        self.biasMinimumRMS = biasMinimumRMS
     }
 
     public func handle(_ event: HotkeyEvent) async {
@@ -49,24 +63,35 @@ public final class DictationController {
                 return
             }
             let energy = samples.reduce(0) { $0 + $1 * $1 }
-            guard (energy / Float(samples.count)).squareRoot() >= silenceThreshold else {
+            let rms = (energy / Float(samples.count)).squareRoot()
+            guard rms >= silenceThreshold else {
                 state = .idle
                 return
             }
-            await process(samples)
+            await process(samples, rms: rms)
         default:
             break
         }
     }
 
-    private func process(_ samples: [Float]) async {
+    private func process(_ samples: [Float], rms: Float) async {
         defer { state = .idle }
         do {
             state = .transcribing
-            let raw = try await transcriber.transcribe(samples)
+            let vocabulary = rms >= biasMinimumRMS ? vocabularyProvider() : []
+            let raw = try await transcriber.transcribe(samples, vocabulary: vocabulary)
             let text = raw.trimmingCharacters(in: .whitespacesAndNewlines)
-            guard !text.isEmpty else { return }
+            guard !text.isEmpty, !PromptEcho.isEcho(text, vocabulary: vocabulary) else { return }
-            let cleaned = try await cleaner.clean(text, context: CleaningContext())
+            let corrected = corrector?.correct(text) ?? CorrectionResult(text: text, corrections: [])
+            let cleaned = try await cleaner.clean(corrected.text, context: CleaningContext())
+            onResult?(DictationResult(
+                date: Date(),
+                duration: Double(samples.count) / 16_000,
+                rawText: text,
+                finalText: cleaned,
+                // Offsets only stay valid if the cleaner left the corrected text alone.
+                corrections: cleaned == corrected.text ? corrected.corrections : []
+            ))
             state = .inserting
             try await inserter.insert(cleaned)
         } catch {

Sources/SstCore/PromptEcho.swift

@@ -0,0 +1,21 @@
+import Foundation
+
+/// Whisper sometimes answers a vocabulary prompt with the prompt itself when there is little speech.
+public enum PromptEcho {
+    static func normalize(_ text: String) -> String {
+        String(text.lowercased().map { $0.isLetter || $0.isNumber ? $0 : " " })
+            .split(separator: " ")
+            .joined(separator: " ")
+    }
+
+    /// True when `text` is nothing but the first k ≥ 2 terms of the prompt.
+    /// A single term on its own is valid speech ("Anthropic").
+    public static func isEcho(_ text: String, vocabulary: [String]) -> Bool {
+        guard vocabulary.count >= 2 else { return false }
+        let normalized = normalize(text)
+        guard !normalized.isEmpty else { return false }
+        return (2...vocabulary.count).contains { k in
+            normalize(vocabulary.prefix(k).joined(separator: " ")) == normalized
+        }
+    }
+}

Tests/SstCoreTests/DictationControllerHooksTests.swift

@@ -0,0 +1,124 @@
+import Foundation
+import Testing
+@testable import SstCore
+
+@MainActor
+struct DictationControllerHooksTests {
+    let recorder = FakeRecorder()
+    let transcriber = FakeVocabularyTranscriber()
+    let inserter = FakeInserter()
+    static let loud = [Float](repeating: 0.1, count: 16_000)
+    static let quiet = [Float](repeating: 0.01, count: 16_000)
+
+    func makeController(
+        cleaner: Cleaner = FakeCleaner(),
+        corrector: TextCorrecting? = nil,
+        vocabulary: [String] = []
+    ) -> DictationController {
+        DictationController(
+            recorder: recorder,
+            transcriber: transcriber,
+            cleaner: cleaner,
+            inserter: inserter,
+            corrector: corrector,
+            vocabularyProvider: { vocabulary }
+        )
+    }
+
+    func dictate(_ controller: DictationController) async {
+        await controller.handle(.pressed)
+        await controller.handle(.released)
+    }
+
+    @Test func loudRecordingPassesVocabularyToTranscriber() async {
+        recorder.samples = Self.loud
+        await dictate(makeController(vocabulary: ["Anthropic", "Vercel"]))
+        #expect(transcriber.vocabularies == [["Anthropic", "Vercel"]])
+    }
+
+    @Test func quietRecordingSkipsVocabularyButStillTranscribes() async {
+        recorder.samples = Self.quiet   // RMS 0.01: above silence (0.005), below bias minimum (0.015)
+        await dictate(makeController(vocabulary: ["Anthropic", "Vercel"]))
+        #expect(transcriber.vocabularies == [[]])
+        #expect(inserter.inserted == ["[hallo wereld]"])
+    }
+
+    @Test func promptEchoIsDiscarded() async {
+        recorder.samples = Self.loud
+        transcriber.result = "Anthropic, Vercel."
+        let controller = makeController(vocabulary: ["Anthropic", "Vercel", "Supabase"])
+        var results: [DictationResult] = []
+        controller.onResult = { results.append($0) }
+
+        await dictate(controller)
+
+        #expect(inserter.inserted.isEmpty)
+        #expect(results.isEmpty)
+        #expect(controller.state == .idle)
+    }
+
+    @Test func singleVocabularyTermIsKept() async {
+        recorder.samples = Self.loud
+        transcriber.result = "Anthropic"
+        await dictate(makeController(vocabulary: ["Anthropic", "Vercel", "Supabase"]))
+        #expect(inserter.inserted == ["[Anthropic]"])
+    }
+
+    @Test func correctorRunsBeforeCleanerAndCorrectionsSurviveAPassthroughCleaner() async {
+        recorder.samples = Self.loud
+        transcriber.result = "hallo"
+        let controller = makeController(cleaner: PassthroughCleaner(), corrector: UppercaseCorrector())
+        var results: [DictationResult] = []
+        controller.onResult = { results.append($0) }
+
+        await dictate(controller)
+
+        #expect(inserter.inserted == ["HALLO"])
+        #expect(results.count == 1)
+        #expect(results[0].rawText == "hallo")
+        #expect(results[0].finalText == "HALLO")
+        #expect(results[0].corrections.count == 1)
+    }
+
+    @Test func correctionsAreDroppedWhenTheCleanerChangesTheText() async {
+        recorder.samples = Self.loud
+        transcriber.result = "hallo"
+        let controller = makeController(cleaner: FakeCleaner(), corrector: UppercaseCorrector())
+        var results: [DictationResult] = []
+        controller.onResult = { results.append($0) }
+
+        await dictate(controller)
+
+        #expect(results[0].finalText == "[HALLO]")
+        #expect(results[0].corrections.isEmpty)
+    }
+
+    @Test func resultIsReportedEvenWhenInsertionFails() async {
+        recorder.samples = Self.loud
+        inserter.error = TestError()
+        let controller = makeController()
+        var results: [DictationResult] = []
+        var errors: [String] = []
+        controller.onResult = { results.append($0) }
+        controller.onError = { errors.append($0) }
+
+        await dictate(controller)
+
+        #expect(results.count == 1)
+        #expect(errors.count == 1)
+    }
+
+    @Test func resultCarriesDurationAndDate() async {
+        recorder.samples = Self.loud   // 16 000 samples = 1 s
+        let controller = makeController()
+        var results: [DictationResult] = []
+        controller.onResult = { results.append($0) }
+        let before = Date()
+
+        await dictate(controller)
+
+        #expect(results[0].duration == 1.0)
+        #expect(results[0].date >= before)
+        #expect(results[0].corrections.isEmpty)
+    }
+}

Tests/SstCoreTests/Fakes.swift

@@ -41,3 +41,31 @@ final class FakeInserter: TextInserting {
         inserted.append(text)
     }
 }
+
+final class FakeVocabularyTranscriber: Transcriber {
+    var result = "hallo wereld"
+    private(set) var vocabularies: [[String]] = []
+
+    func transcribe(_ samples: [Float]) async throws -> String { result }
+
+    func transcribe(_ samples: [Float], vocabulary: [String]) async throws -> String {
+        vocabularies.append(vocabulary)
+        return result
+    }
+}
+
+struct UppercaseCorrector: TextCorrecting {
+    func correct(_ text: String) -> CorrectionResult {
+        let upper = text.uppercased()
+        return CorrectionResult(
+            text: upper,
+            corrections: [AppliedCorrection(heard: text, wrote: upper, rule: "upper", start: 0, length: upper.utf16.count)]
+        )
+    }
+}
+
+final class PassthroughCleaner: Cleaner {
+    func clean(_ text: String, context: CleaningContext) async throws -> String {
+        text
+    }
+}

Tests/SstCoreTests/PromptEchoTests.swift

@@ -0,0 +1,31 @@
+import Testing
+@testable import SstCore
+
+struct PromptEchoTests {
+    let vocabulary = ["Anthropic", "Vercel", "Supabase"]
+
+    @Test func firstTwoTermsAreAnEcho() {
+        #expect(PromptEcho.isEcho("Anthropic, Vercel.", vocabulary: vocabulary))
+    }
+
+    @Test func wholeListIsAnEcho() {
+        #expect(PromptEcho.isEcho("anthropic vercel supabase", vocabulary: vocabulary))
+    }
+
+    @Test func singleTermIsRealSpeech() {
+        #expect(!PromptEcho.isEcho("Anthropic", vocabulary: vocabulary))
+    }
+
+    @Test func termsOutOfOrderAreNotAnEcho() {
+        #expect(!PromptEcho.isEcho("Vercel, Anthropic", vocabulary: vocabulary))
+    }
+
+    @Test func sentenceContainingTermsIsNotAnEcho() {
+        #expect(!PromptEcho.isEcho("Ik gebruik Anthropic en Vercel", vocabulary: vocabulary))
+    }
+
+    @Test func tooShortVocabularyNeverEchoes() {
+        #expect(!PromptEcho.isEcho("Anthropic", vocabulary: ["Anthropic"]))
+        #expect(!PromptEcho.isEcho("", vocabulary: vocabulary))
+    }
+}