diff --git a/wyoming_satellite/satellite.py b/wyoming_satellite/satellite.py index eade272..d009531 100644 --- a/wyoming_satellite/satellite.py +++ b/wyoming_satellite/satellite.py @@ -42,6 +42,8 @@ class State(Enum): class SatelliteBase: + """Base class for satellites.""" + def __init__(self, settings: SatelliteSettings) -> None: self.settings = settings self.server_id: Optional[str] = None @@ -60,19 +62,23 @@ class SatelliteBase: @property def is_running(self) -> bool: + """True if not stopping/stopped.""" return self._state not in (State.STOPPING, State.STOPPED) @property def state(self) -> State: + """Get current state.""" return self._state @state.setter def state(self, value: State) -> None: + """Set state.""" self._state = value self._state_changed.set() self._state_changed.clear() async def run(self) -> None: + """Run main satellite loop.""" while self.is_running: try: @@ -93,36 +99,53 @@ class SatelliteBase: _LOGGER.exception("Unexpected error running satellite") self.state = State.RESTARTING - async def _start(self) -> None: - self.state = State.STARTING - await self._connect_to_services() - self.state = State.STARTED - await self.started() - - async def started(self) -> None: - pass - - async def _stop(self) -> None: - self.server_id = None - self._writer = None - - await self._disconnect_from_services() - self.state = State.STOPPED - async def stop(self) -> None: + """Stop the satellite""" self.state = State.STOPPING # Wait for STOPPED while self.state != State.STOPPED: await self._state_changed.wait() - async def stopped(self) -> None: - pass - def set_server(self, server_id: str, writer: asyncio.StreamWriter) -> None: + """Set event writer.""" self.server_id = server_id self._writer = writer + async def event_to_server(self, event: Event) -> None: + """Send an event to the server.""" + if self._writer is None: + return + + try: + await async_write_event(event, self._writer) + except Exception: + _LOGGER.exception("Unexpected error sending event to server") + self._writer = None + + # ------------------------------------------------------------------------- + + async def _start(self) -> None: + """Connect to services.""" + self.state = State.STARTING + await self._connect_to_services() + self.state = State.STARTED + await self.started() + + async def started(self) -> None: + """Called when satellite has started.""" + + async def _stop(self) -> None: + """Disconnect from services.""" + self.server_id = None + self._writer = None + + await self._disconnect_from_services() + self.state = State.STOPPED + + async def stopped(self) -> None: + """Called when satellite has stopped.""" + async def event_from_server(self, event: Event) -> None: """Called when an event is received from the server.""" if AudioChunk.is_type(event.type): @@ -165,18 +188,8 @@ class SatelliteBase: if not AudioChunk.is_type(event.type): await self.forward_event(event) - async def event_to_server(self, event: Event) -> None: - """Send an event to the server.""" - if self._writer is None: - return - - try: - await async_write_event(event, self._writer) - except Exception: - _LOGGER.exception("Unexpected error sending event to server") - self._writer = None - async def _send_run_pipeline(self) -> None: + """Sends a RunPipeline event with the correct stages.""" if self.settings.wake.enabled: # Local wake word detection start_stage = PipelineStage.ASR @@ -196,6 +209,7 @@ class SatelliteBase: await self.forward_event(run_pipeline) async def _restart(self) -> None: + """Disconnects from services and restarts loop.""" self.state = State.RESTARTING await self._disconnect_from_services() @@ -204,6 +218,7 @@ class SatelliteBase: self.state = State.NOT_STARTED async def _connect_to_services(self) -> None: + """Connects to configured services.""" if self.settings.mic.enabled: _LOGGER.debug( "Connecting to mic service: %s", @@ -235,6 +250,7 @@ class SatelliteBase: _LOGGER.info("Connected to services") async def _disconnect_from_services(self) -> None: + """Disconnects from running services.""" if self._mic_task is not None: _LOGGER.debug("Stopping microphone service") self._mic_task.cancel() @@ -264,9 +280,15 @@ class SatelliteBase: async def event_from_mic( self, event: Event, audio_bytes: Optional[bytes] = None ) -> None: - pass + """Called when an event is received from the mic service. + + For AudioChunk events, the audio_bytes may be set if the audio was + already unpacked for preprocessing. Use this to avoid unpacking the + event again. + """ def _make_mic_client(self) -> Optional[AsyncClient]: + """Create client for mic service.""" if self.settings.mic.command: program, *program_args = self.settings.mic.command return MicProcessAsyncClient( @@ -284,6 +306,7 @@ class SatelliteBase: return None async def _mic_task_proc(self) -> None: + """Mic service loop.""" mic_client: Optional[AsyncClient] = None audio_bytes: Optional[bytes] = None @@ -343,12 +366,14 @@ class SatelliteBase: await _disconnect() def _process_mic_audio(self, audio_bytes: bytes) -> bytes: + """Perform audio pre-processing on mic input.""" if self.settings.mic.volume_multiplier != 1.0: audio_bytes = multiply_volume( audio_bytes, self.settings.mic.volume_multiplier ) if self._mic_webrtc is not None: + # Noise suppression and auto gain audio_bytes = self._mic_webrtc(audio_bytes) return audio_bytes @@ -363,6 +388,7 @@ class SatelliteBase: self._snd_queue.put_nowait(event) def _make_snd_client(self) -> Optional[AsyncClient]: + """Create client for snd service.""" if self.settings.snd.command: program, *program_args = self.settings.snd.command return SndProcessAsyncClient( @@ -379,6 +405,7 @@ class SatelliteBase: return None async def _snd_task_proc(self) -> None: + """Snd service loop.""" snd_client: Optional[AsyncClient] = None async def _disconnect() -> None: @@ -401,7 +428,18 @@ class SatelliteBase: await snd_client.connect() _LOGGER.debug("Connected to snd service") - # TODO: audio processing + # Audio processing + if self.settings.snd.needs_processing and AudioChunk.is_type( + event.type + ): + chunk = AudioChunk.from_event(event) + audio_bytes = self._process_snd_audio(chunk.audio) + event = AudioChunk( + rate=chunk.rate, + width=chunk.width, + channels=chunk.channels, + audio=audio_bytes, + ).event() await snd_client.write_event(event) except asyncio.CancelledError: @@ -415,6 +453,15 @@ class SatelliteBase: await _disconnect() + def _process_snd_audio(self, audio_bytes: bytes) -> bytes: + """Perform audio pre-processing on snd output.""" + if self.settings.snd.volume_multiplier != 1.0: + audio_bytes = multiply_volume( + audio_bytes, self.settings.snd.volume_multiplier + ) + + return audio_bytes + async def _play_wav(self, wav_path: Optional[Union[str, Path]]) -> None: """Send WAV as events to sound service.""" if (not wav_path) or (not self.settings.snd.enabled): @@ -432,13 +479,15 @@ class SatelliteBase: # ------------------------------------------------------------------------- async def event_from_wake(self, event: Event) -> None: - pass + """Called when an event is received from the wake service.""" async def event_to_wake(self, event: Event) -> None: + """Send event to the wake service.""" if self._wake_queue is not None: self._wake_queue.put_nowait(event) def _make_wake_client(self) -> Optional[AsyncClient]: + """Create client for wake service.""" if self.settings.wake.command: program, *program_args = self.settings.wake.command return WakeProcessAsyncClient( @@ -455,6 +504,7 @@ class SatelliteBase: return None async def _wake_task_proc(self) -> None: + """Wake service loop.""" wake_client: Optional[AsyncClient] = None to_client_task: Optional[asyncio.Task] = None from_client_task: Optional[asyncio.Task] = None @@ -484,8 +534,10 @@ class SatelliteBase: pending = set() self._wake_queue = asyncio.Queue() + # Inform wake service of which wake word(s) to detect await self._send_wake_detect() + # Read/write in "parallel" if to_client_task is None: # From satellite to wake service to_client_task = asyncio.create_task(self._wake_queue.get()) @@ -501,12 +553,14 @@ class SatelliteBase: ) if to_client_task in done: + # Event to go to wake service (audio) assert to_client_task is not None event = to_client_task.result() to_client_task = None await wake_client.write_event(event) if from_client_task in done: + # Event from wake service (detection) assert from_client_task is not None event = from_client_task.result() from_client_task = None @@ -518,7 +572,6 @@ class SatelliteBase: await asyncio.sleep(self.settings.wake.reconnect_seconds) continue - # TODO: audio processing await self.event_from_wake(event) except asyncio.CancelledError: @@ -541,40 +594,51 @@ class SatelliteBase: # ------------------------------------------------------------------------- async def trigger_streaming_start(self) -> None: + """Called when audio streaming starts.""" await run_event_command(self.settings.event.streaming_start) await self.forward_event(StreamingStarted().event()) async def trigger_streaming_stop(self) -> None: + """Called when audio streaming stops.""" await run_event_command(self.settings.event.streaming_stop) await self.forward_event(StreamingStopped().event()) async def trigger_detect(self) -> None: + """Called when wake word detection starts.""" await run_event_command(self.settings.event.detect) async def trigger_detection(self, detection: Detection) -> None: + """Called when wake word is detected.""" await run_event_command(self.settings.event.detection, detection.name) await self._play_wav(self.settings.snd.awake_wav) async def trigger_transcript(self, transcript: Transcript) -> None: + """Called when speech-to-text text is received.""" await run_event_command(self.settings.event.transcript, transcript.text) await self._play_wav(self.settings.snd.done_wav) async def trigger_stt_start(self) -> None: + """Called when user starts speaking.""" await run_event_command(self.settings.event.stt_start) async def trigger_stt_stop(self) -> None: + """Called when user stops speaking.""" await run_event_command(self.settings.event.stt_stop) async def trigger_synthesize(self, synthesize: Synthesize) -> None: + """Called when text-to-speech text is received.""" await run_event_command(self.settings.event.synthesize, synthesize.text) async def trigger_tts_start(self) -> None: + """Called when text-to-speech audio starts.""" await run_event_command(self.settings.event.tts_start) async def trigger_tts_stop(self) -> None: + """Called when text-to-speech audio stops.""" await run_event_command(self.settings.event.tts_stop) async def trigger_error(self, error: Error) -> None: + """Called when an error occurs on the server.""" await run_event_command(self.settings.event.error, error.text) async def forward_event(self, event: Event) -> None: @@ -583,12 +647,14 @@ class SatelliteBase: self._event_queue.put_nowait(event) def _make_event_client(self) -> Optional[AsyncClient]: + """Create client for event service.""" if self.settings.event.uri: return AsyncClient.from_uri(self.settings.event.uri) return None async def _event_task_proc(self) -> None: + """Event service loop.""" event_client: Optional[AsyncClient] = None async def _disconnect() -> None: @@ -644,6 +710,9 @@ class AlwaysStreamingSatellite(SatelliteBase): await self.trigger_streaming_start() elif Transcript.is_type(event.type): _LOGGER.info("Streaming audio") + + # Re-trigger streaming start even though we technically don't stop + # so the event service can reset LEDs, etc. await self.trigger_streaming_start() async def event_from_mic( @@ -672,7 +741,12 @@ class VadStreamingSatellite(SatelliteBase): self.vad = SileroVad( threshold=settings.vad.threshold, trigger_level=settings.vad.trigger_level ) + + # Timestamp in the future when we will have timed out (set with + # time.monotonic()) self.timeout_seconds: Optional[float] = None + + # Audio from right before speech starts (circular buffer) self.vad_buffer: Optional[RingBuffer] = None if settings.vad.buffer_seconds > 0: @@ -729,14 +803,18 @@ class VadStreamingSatellite(SatelliteBase): await self.trigger_streaming_start() if self.settings.vad.wake_word_timeout is not None: + # Set future time when we'll stop streaming if the wake word + # hasn't been detected. self.timeout_seconds = ( time.monotonic() + self.settings.vad.wake_word_timeout ) else: + # No timeout self.timeout_seconds = None if self.vad_buffer is not None: - # Send contents of VAD buffer first + # Send contents of VAD buffer first. This is the audio that was + # recorded right before speech was detected. if chunk is None: chunk = AudioChunk.from_event(event) @@ -756,8 +834,11 @@ class VadStreamingSatellite(SatelliteBase): await self.event_to_server(event) def _reset_vad(self): + """Reset state of VAD.""" self.vad(None) + if self.vad_buffer is not None: + # Clear buffer self.vad_buffer.put(bytes(self.vad_buffer.maxlen)) @@ -775,6 +856,12 @@ class WakeStreamingSatellite(SatelliteBase): self.is_streaming = False async def event_from_server(self, event: Event) -> None: + if Transcript.is_type(event.type): + # Stop streaming before event_from_server is called because it will + # play the "done" WAV. + self.is_streaming = False + await self.trigger_streaming_stop() + await super().event_from_server(event) if RunSatellite.is_type(event.type) or Transcript.is_type(event.type): @@ -782,9 +869,6 @@ class WakeStreamingSatellite(SatelliteBase): await self._send_wake_detect() _LOGGER.info("Waiting for wake word") - if Transcript.is_type(event.type): - await self.trigger_streaming_stop() - async def event_from_mic( self, event: Event, audio_bytes: Optional[bytes] = None ) -> None: