Skip to main content

aria_kit/audio/capture/microphone/
wake_word.rs

1//! Listening for the wake word while the microphone is idle.
2
3use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, SyncSender, TrySendError};
4use std::time::Duration;
5
6use crate::audio::canonical_samples;
7use crate::primitives::error_chain;
8use crate::voice_activity::{SPEECH_THRESHOLD, VoiceActivityDetector, VoiceActivityError};
9use crate::wake_word::{WakeWordDetector, WakeWordMatch};
10
11use super::worker::{Command, IdleMonitor, IdleObservation};
12use super::{MicrophoneError, Result};
13
14#[derive(Debug)]
15pub enum WakeWordEvent {
16    Detected,
17    /// Monitoring stopped and stays stopped until the microphone reopens.
18    Failed(WakeWordFailure),
19}
20
21#[derive(Debug, thiserror::Error)]
22pub enum WakeWordFailure {
23    #[error("voice-activity detection failed")]
24    VoiceActivity(#[source] VoiceActivityError),
25
26    #[error("wake-word detection failed")]
27    Detector(#[source] Box<dyn std::error::Error + Send + Sync>),
28}
29
30pub struct WakeWordMonitor {
31    control: WakeWordControl,
32    events: Receiver<WakeWordEvent>,
33}
34
35impl WakeWordMonitor {
36    pub(super) fn new(commands: Sender<Command>, events: Receiver<WakeWordEvent>) -> Self {
37        Self {
38            control: WakeWordControl { commands },
39            events,
40        }
41    }
42
43    pub fn control(&self) -> WakeWordControl {
44        self.control.clone()
45    }
46
47    pub fn wait_timeout(&self, timeout: Duration) -> Result<Option<WakeWordEvent>> {
48        match self.events.recv_timeout(timeout) {
49            Ok(event) => Ok(Some(event)),
50            Err(RecvTimeoutError::Timeout) => Ok(None),
51            Err(RecvTimeoutError::Disconnected) => Err(MicrophoneError::WorkerGone.into()),
52        }
53    }
54
55    pub fn resume(&self) -> Result<()> {
56        self.control.resume()
57    }
58}
59
60#[derive(Clone)]
61pub struct WakeWordControl {
62    commands: Sender<Command>,
63}
64
65impl WakeWordControl {
66    pub fn suspend(&self) -> Result<WakeWordSuspension> {
67        self.commands
68            .send(Command::SuspendWakeWord)
69            .map_err(|_command| MicrophoneError::WorkerGone)?;
70        Ok(WakeWordSuspension {
71            control: self.clone(),
72        })
73    }
74
75    pub fn resume(&self) -> Result<()> {
76        self.commands
77            .send(Command::ResumeWakeWord)
78            .map_err(|_command| MicrophoneError::WorkerGone.into())
79    }
80}
81
82pub struct WakeWordSuspension {
83    control: WakeWordControl,
84}
85
86impl Drop for WakeWordSuspension {
87    fn drop(&mut self) {
88        if let Err(error) = self.control.resume() {
89            tracing::error!(%error, "wake-word monitoring could not resume");
90        }
91    }
92}
93
94/// How recently speech must have been heard for a detection to count. Wake
95/// classifiers score a window of a second or more and fire up to a few hundred
96/// milliseconds after the word ends, often on a silent frame.
97const SPEECH_LOOKBACK_SAMPLES: usize = canonical_samples(Duration::from_secs(1)) as usize;
98
99pub(super) struct WakeMonitor<W> {
100    detector: W,
101    events: SyncSender<WakeWordEvent>,
102    state: WakeMonitorState,
103    vad_frame: Vec<i16>,
104    samples_since_speech: Option<usize>,
105}
106
107#[derive(Clone, Copy, PartialEq, Eq)]
108enum WakeMonitorState {
109    Armed,
110    Suspended,
111    Triggered,
112    Failed,
113}
114
115enum Frame {
116    Quiet,
117    Detected,
118    Failed,
119}
120
121impl<W: WakeWordDetector> WakeMonitor<W> {
122    pub(super) fn new(
123        detector: W,
124        events: SyncSender<WakeWordEvent>,
125        frame_samples: usize,
126    ) -> Self {
127        Self {
128            detector,
129            events,
130            state: WakeMonitorState::Armed,
131            vad_frame: Vec::with_capacity(frame_samples),
132            samples_since_speech: None,
133        }
134    }
135
136    fn forget_audio(&mut self) {
137        self.detector.reset();
138        self.vad_frame.clear();
139        self.samples_since_speech = None;
140    }
141
142    fn track_speech(&mut self, probability: f32) {
143        self.samples_since_speech = if probability >= SPEECH_THRESHOLD {
144            Some(0)
145        } else {
146            self.samples_since_speech
147                .map(|samples| samples + self.vad_frame.len())
148        };
149    }
150
151    fn heard_speech_recently(&self) -> bool {
152        self.samples_since_speech
153            .is_some_and(|samples| samples <= SPEECH_LOOKBACK_SAMPLES)
154    }
155
156    fn publish(&self, event: WakeWordEvent) {
157        match self.events.try_send(event) {
158            Ok(()) | Err(TrySendError::Full(WakeWordEvent::Detected)) => {}
159            Err(TrySendError::Full(WakeWordEvent::Failed(failure))) => {
160                tracing::warn!(
161                    error = %error_chain(&failure),
162                    "wake-word monitoring failed; the event queue was full"
163                );
164            }
165            Err(TrySendError::Disconnected(_)) => {
166                tracing::info!("wake-word monitor stopped; event receiver closed");
167            }
168        }
169    }
170
171    fn fail(&mut self, failure: WakeWordFailure) {
172        self.state = WakeMonitorState::Failed;
173        self.publish(WakeWordEvent::Failed(failure));
174    }
175
176    fn observe_frame<V: VoiceActivityDetector>(&mut self, vad: &mut V) -> Frame {
177        let probability = match vad.probability(&self.vad_frame) {
178            Ok(probability) => probability,
179            Err(error) => {
180                self.fail(WakeWordFailure::VoiceActivity(error));
181                return Frame::Failed;
182            }
183        };
184        let detection = match self.detector.detect(&self.vad_frame) {
185            Ok(detection) => detection,
186            Err(error) => {
187                self.fail(WakeWordFailure::Detector(Box::new(error)));
188                return Frame::Failed;
189            }
190        };
191        self.track_speech(probability);
192        if detection == WakeWordMatch::Detected && self.heard_speech_recently() {
193            self.state = WakeMonitorState::Triggered;
194            self.publish(WakeWordEvent::Detected);
195            return Frame::Detected;
196        }
197        Frame::Quiet
198    }
199}
200
201impl<W: WakeWordDetector> IdleMonitor for WakeMonitor<W> {
202    fn observe<V: VoiceActivityDetector>(
203        &mut self,
204        vad: &mut V,
205        samples: &[i16],
206    ) -> IdleObservation {
207        if self.state != WakeMonitorState::Armed {
208            return IdleObservation::Continue;
209        }
210
211        for (index, &sample) in samples.iter().enumerate() {
212            self.vad_frame.push(sample);
213            if self.vad_frame.len() == V::FRAME_SAMPLES {
214                let frame = self.observe_frame(vad);
215                self.vad_frame.clear();
216                match frame {
217                    Frame::Quiet => {}
218                    Frame::Detected => {
219                        return IdleObservation::Detected {
220                            command_start: index + 1,
221                        };
222                    }
223                    Frame::Failed => return IdleObservation::Continue,
224                }
225            }
226        }
227        IdleObservation::Continue
228    }
229
230    fn reset(&mut self) {
231        self.forget_audio();
232    }
233
234    fn suspend(&mut self) {
235        self.forget_audio();
236        self.state = match self.state {
237            WakeMonitorState::Armed | WakeMonitorState::Suspended | WakeMonitorState::Triggered => {
238                WakeMonitorState::Suspended
239            }
240            WakeMonitorState::Failed => WakeMonitorState::Failed,
241        };
242    }
243
244    fn resume(&mut self) {
245        self.forget_audio();
246        self.state = match self.state {
247            WakeMonitorState::Armed | WakeMonitorState::Suspended | WakeMonitorState::Triggered => {
248                WakeMonitorState::Armed
249            }
250            WakeMonitorState::Failed => WakeMonitorState::Failed,
251        };
252    }
253
254    fn recording_finished(&mut self) {
255        self.forget_audio();
256    }
257
258    fn requires_live_stream(&self) -> bool {
259        self.state != WakeMonitorState::Failed
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use std::convert::Infallible;
266    use std::result::Result;
267    use std::sync::mpsc::sync_channel;
268
269    use super::*;
270
271    const FRAME: usize = 4;
272
273    struct FixedVad(f32);
274
275    impl VoiceActivityDetector for FixedVad {
276        const FRAME_SAMPLES: usize = FRAME;
277
278        fn probability(
279            &mut self,
280            samples: &[i16],
281        ) -> Result<f32, crate::voice_activity::VoiceActivityError> {
282            assert_eq!(samples.len(), Self::FRAME_SAMPLES);
283            Ok(self.0)
284        }
285    }
286
287    struct SpeechThenSilence {
288        frames: usize,
289    }
290
291    impl VoiceActivityDetector for SpeechThenSilence {
292        const FRAME_SAMPLES: usize = FRAME;
293
294        fn probability(
295            &mut self,
296            _samples: &[i16],
297        ) -> Result<f32, crate::voice_activity::VoiceActivityError> {
298            self.frames += 1;
299            Ok(if self.frames == 1 { 1.0 } else { 0.0 })
300        }
301    }
302
303    struct DetectsOnFrame {
304        frame: usize,
305        calls: usize,
306    }
307
308    impl WakeWordDetector for DetectsOnFrame {
309        type Error = Infallible;
310
311        fn detect(&mut self, _canonical_samples: &[i16]) -> Result<WakeWordMatch, Self::Error> {
312            self.calls += 1;
313            Ok(if self.calls == self.frame {
314                WakeWordMatch::Detected
315            } else {
316                WakeWordMatch::NoMatch
317            })
318        }
319
320        fn reset(&mut self) {}
321    }
322
323    struct FailingVad;
324
325    impl VoiceActivityDetector for FailingVad {
326        const FRAME_SAMPLES: usize = FRAME;
327
328        fn probability(
329            &mut self,
330            _samples: &[i16],
331        ) -> Result<f32, crate::voice_activity::VoiceActivityError> {
332            Err(crate::voice_activity::VoiceActivityError::InvalidProbability(f32::NAN))
333        }
334    }
335
336    fn monitor(frame: usize) -> (WakeMonitor<DetectsOnFrame>, Receiver<WakeWordEvent>) {
337        let (events, detections) = sync_channel(1);
338        let monitor = WakeMonitor::new(DetectsOnFrame { frame, calls: 0 }, events, FRAME);
339        (monitor, detections)
340    }
341
342    #[test]
343    fn recording_completion_does_not_rearm_a_triggered_monitor() {
344        let (mut monitor, _detections) = monitor(1);
345        let mut vad = FixedVad(1.0);
346
347        monitor.observe(&mut vad, &[0; FRAME]);
348        monitor.recording_finished();
349
350        assert!(matches!(monitor.state, WakeMonitorState::Triggered));
351    }
352
353    #[test]
354    fn idle_vad_rejects_a_wake_detection_without_speech() {
355        let (mut monitor, detections) = monitor(1);
356        let mut vad = FixedVad(SPEECH_THRESHOLD - 0.1);
357
358        let observation = monitor.observe(&mut vad, &[0; FRAME]);
359
360        assert!(matches!(observation, IdleObservation::Continue));
361        assert!(matches!(monitor.state, WakeMonitorState::Armed));
362        assert_eq!(monitor.detector.calls, 1);
363        assert!(detections.try_recv().is_err());
364    }
365
366    #[test]
367    fn idle_vad_buffers_partial_frames() {
368        let (mut monitor, detections) = monitor(1);
369        let mut vad = FixedVad(SPEECH_THRESHOLD);
370
371        let partial = monitor.observe(&mut vad, &[0; FRAME - 1]);
372        let complete = monitor.observe(&mut vad, &[0]);
373
374        assert!(matches!(partial, IdleObservation::Continue));
375        assert!(matches!(
376            complete,
377            IdleObservation::Detected { command_start: 1 }
378        ));
379        assert_eq!(monitor.detector.calls, 1);
380        assert!(matches!(detections.try_recv(), Ok(WakeWordEvent::Detected)));
381    }
382
383    #[test]
384    fn idle_vad_failure_suspends_wake_monitoring() {
385        let (mut monitor, detections) = monitor(1);
386        let mut vad = FailingVad;
387
388        let observation = monitor.observe(&mut vad, &[0; FRAME]);
389
390        assert!(matches!(observation, IdleObservation::Continue));
391        assert!(matches!(monitor.state, WakeMonitorState::Failed));
392        assert!(!monitor.requires_live_stream());
393        assert_eq!(monitor.detector.calls, 0);
394        assert!(matches!(
395            detections.try_recv(),
396            Ok(WakeWordEvent::Failed(WakeWordFailure::VoiceActivity(_)))
397        ));
398    }
399
400    #[test]
401    fn a_detection_on_silence_counts_when_speech_was_just_heard() {
402        let (mut monitor, detections) = monitor(3);
403        let mut vad = SpeechThenSilence { frames: 0 };
404
405        let observation = monitor.observe(&mut vad, &[0; 5 * FRAME]);
406
407        assert!(matches!(
408            observation,
409            IdleObservation::Detected { command_start: 12 }
410        ));
411        assert!(matches!(detections.try_recv(), Ok(WakeWordEvent::Detected)));
412    }
413
414    #[test]
415    fn a_detection_long_after_the_last_speech_is_rejected() {
416        let frames_past_lookback = SPEECH_LOOKBACK_SAMPLES / FRAME + 2;
417        let (mut monitor, detections) = monitor(frames_past_lookback);
418        let mut vad = SpeechThenSilence { frames: 0 };
419
420        let observation = monitor.observe(&mut vad, &vec![0; frames_past_lookback * FRAME]);
421
422        assert!(matches!(observation, IdleObservation::Continue));
423        assert!(detections.try_recv().is_err());
424    }
425}