aria_kit/audio/capture/microphone/
wake_word.rs1use 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 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
94const 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}