Small refactor

linux_voice_assistant/__main__.py

@@ -5,45 +5,17 @@ import asyncio
 import json
 import logging
 import threading
-import time
-from collections.abc import Iterable
-from dataclasses import dataclass
 from pathlib import Path
 from queue import Queue
-from typing import Dict, List, Optional
+from typing import Dict
 
-# pylint: disable=no-name-in-module
 import sounddevice as sd
-from aioesphomeapi.api_pb2 import (  # type: ignore[attr-defined]
-    DeviceInfoRequest,
-    DeviceInfoResponse,
-    ListEntitiesDoneResponse,
-    ListEntitiesRequest,
-    MediaPlayerCommandRequest,
-    SubscribeHomeAssistantStatesRequest,
-    VoiceAssistantAnnounceFinished,
-    VoiceAssistantAnnounceRequest,
-    VoiceAssistantAudio,
-    VoiceAssistantConfigurationRequest,
-    VoiceAssistantConfigurationResponse,
-    VoiceAssistantEventResponse,
-    VoiceAssistantRequest,
-    VoiceAssistantSetConfiguration,
-    VoiceAssistantTimerEventResponse,
-    VoiceAssistantWakeWord,
-)
-from aioesphomeapi.model import (
-    VoiceAssistantEventType,
-    VoiceAssistantFeature,
-    VoiceAssistantTimerEventType,
-)
-from google.protobuf import message
 
-from .api_server import APIServer
-from .entity import ESPHomeEntity, MediaPlayerEntity
 from .microwakeword import MicroWakeWord
+from .models import AvailableWakeWord, ServerState
 from .mpv_player import MpvMediaPlayer
-from .util import call_all, get_mac, is_arm
+from .satellite import VoiceSatelliteProtocol
+from .util import get_mac, is_arm
 
 _LOGGER = logging.getLogger(__name__)
 _MODULE_DIR = Path(__file__).parent
@@ -57,306 +29,6 @@ else:
     _LIB_DIR = _REPO_DIR / "lib" / "linux_amd64"
 
 
-@dataclass
-class AvailableWakeWord:
-    id: str
-    wake_word: str
-    trained_languages: List[str]
-    config_path: Path
-
-
-@dataclass
-class ServerState:
-    name: str
-    mac_address: str
-    audio_queue: "Queue[Optional[bytes]]"
-    entities: List[ESPHomeEntity]
-    available_wake_words: Dict[str, AvailableWakeWord]
-    wake_word: MicroWakeWord
-    stop_word: MicroWakeWord
-    music_player: MpvMediaPlayer
-    tts_player: MpvMediaPlayer
-    wakeup_sound: str
-    timer_finished_sound: str
-    media_player_entity: Optional[MediaPlayerEntity] = None
-    satellite: "Optional[VoiceSatelliteProtocol]" = None
-
-
-# -----------------------------------------------------------------------------
-
-
-class VoiceSatelliteProtocol(APIServer):
-
-    def __init__(self, state: ServerState) -> None:
-        super().__init__(state.name)
-
-        self.state = state
-        self.state.satellite = self
-
-        if self.state.media_player_entity is None:
-            self.state.media_player_entity = MediaPlayerEntity(
-                server=self,
-                key=len(state.entities),
-                name="Media Player",
-                object_id="linux_voice_assistant_media_player",
-                music_player=state.music_player,
-                announce_player=state.tts_player,
-            )
-            self.state.entities.append(self.state.media_player_entity)
-
-        self._is_streaming_audio = False
-        self._tts_url: Optional[str] = None
-        self._tts_played = False
-        self._continue_conversation = False
-        self._timer_finished = False
-
-    def handle_voice_event(
-        self, event_type: VoiceAssistantEventType, data: Dict[str, str]
-    ) -> None:
-        _LOGGER.debug("Voice event: type=%s, data=%s", event_type.name, data)
-
-        if event_type == VoiceAssistantEventType.VOICE_ASSISTANT_RUN_START:
-            self._tts_url = data.get("url")
-            self._tts_played = False
-            self._continue_conversation = False
-        elif event_type in (
-            VoiceAssistantEventType.VOICE_ASSISTANT_STT_VAD_END,
-            VoiceAssistantEventType.VOICE_ASSISTANT_STT_END,
-        ):
-            self._is_streaming_audio = False
-        elif event_type == VoiceAssistantEventType.VOICE_ASSISTANT_INTENT_PROGRESS:
-            if data.get("tts_start_streaming") == "1":
-                # Start streaming early
-                self.play_tts()
-        elif event_type == VoiceAssistantEventType.VOICE_ASSISTANT_INTENT_END:
-            if data.get("continue_conversation") == "1":
-                self._continue_conversation = True
-        elif event_type == VoiceAssistantEventType.VOICE_ASSISTANT_TTS_END:
-            self._tts_url = data.get("url")
-            self.play_tts()
-        elif event_type == VoiceAssistantEventType.VOICE_ASSISTANT_RUN_END:
-            self._is_streaming_audio = False
-            if not self._tts_played:
-                self._tts_finished()
-
-            self._tts_played = False
-
-        # TODO: handle error
-
-    def handle_timer_event(
-        self,
-        event_type: VoiceAssistantTimerEventType,
-        msg: VoiceAssistantTimerEventResponse,
-    ) -> None:
-        _LOGGER.debug("Timer event: type=%s", event_type.name)
-        if event_type == VoiceAssistantTimerEventType.VOICE_ASSISTANT_TIMER_FINISHED:
-            if not self._timer_finished:
-                self.state.stop_word.is_active = True
-                self._timer_finished = True
-                self.duck()
-                self._play_timer_finished()
-
-    def handle_message(self, msg: message.Message) -> Iterable[message.Message]:
-        if isinstance(msg, VoiceAssistantEventResponse):
-            # Pipeline event
-            data: Dict[str, str] = {}
-            for arg in msg.data:
-                data[arg.name] = arg.value
-
-            self.handle_voice_event(VoiceAssistantEventType(msg.event_type), data)
-        elif isinstance(msg, VoiceAssistantAnnounceRequest):
-            _LOGGER.debug("Announcing: %s", msg.text)
-
-            assert self.state.media_player_entity is not None
-
-            urls = []
-            if msg.preannounce_media_id:
-                urls.append(msg.preannounce_media_id)
-
-            urls.append(msg.media_id)
-
-            self.state.stop_word.is_active = True
-            self._continue_conversation = msg.start_conversation
-
-            self.duck()
-            yield from self.state.media_player_entity.play(
-                urls, announcement=True, done_callback=self._tts_finished
-            )
-        elif isinstance(msg, VoiceAssistantTimerEventResponse):
-            self.handle_timer_event(VoiceAssistantTimerEventType(msg.event_type), msg)
-        elif isinstance(msg, DeviceInfoRequest):
-            yield DeviceInfoResponse(
-                uses_password=False,
-                name=self.state.name,
-                mac_address=self.state.mac_address,
-                voice_assistant_feature_flags=(
-                    VoiceAssistantFeature.VOICE_ASSISTANT
-                    | VoiceAssistantFeature.API_AUDIO
-                    | VoiceAssistantFeature.ANNOUNCE
-                    | VoiceAssistantFeature.START_CONVERSATION
-                    | VoiceAssistantFeature.TIMERS
-                ),
-            )
-        elif isinstance(
-            msg,
-            (
-                ListEntitiesRequest,
-                SubscribeHomeAssistantStatesRequest,
-                MediaPlayerCommandRequest,
-            ),
-        ):
-            for entity in self.state.entities:
-                yield from entity.handle_message(msg)
-
-            if isinstance(msg, ListEntitiesRequest):
-                yield ListEntitiesDoneResponse()
-        elif isinstance(msg, VoiceAssistantConfigurationRequest):
-            yield VoiceAssistantConfigurationResponse(
-                available_wake_words=[
-                    VoiceAssistantWakeWord(
-                        id=ww.id,
-                        wake_word=ww.wake_word,
-                        trained_languages=ww.trained_languages,
-                    )
-                    for ww in self.state.available_wake_words.values()
-                ],
-                active_wake_words=[self.state.wake_word.id],
-                max_active_wake_words=1,
-            )
-            _LOGGER.info("Connected to Home Assistant")
-        elif isinstance(msg, VoiceAssistantSetConfiguration):
-            # TODO: support multiple wake words
-            for wake_word_id in msg.active_wake_words:
-                if wake_word_id == self.state.wake_word.id:
-                    # Already active
-                    break
-
-                model_info = self.state.available_wake_words.get(wake_word_id)
-                if not model_info:
-                    continue
-
-                _LOGGER.debug("Loading wake word: %s", model_info.config_path)
-                self.state.wake_word = MicroWakeWord.from_config(
-                    model_info.config_path,
-                    self.state.wake_word.libtensorflowlite_c_path,
-                )
-
-                _LOGGER.info("Wake word set: %s", self.state.wake_word.wake_word)
-                break
-
-    def handle_audio(self, audio_chunk: bytes) -> None:
-
-        if not self._is_streaming_audio:
-            return
-
-        self.send_messages([VoiceAssistantAudio(data=audio_chunk)])
-
-    def wakeup(self) -> None:
-        if self._timer_finished:
-            # Stop timer instead
-            self._timer_finished = False
-            self.state.tts_player.stop()
-            _LOGGER.debug("Stopping timer finished sound")
-            return
-
-        wake_word_phrase = self.state.wake_word.wake_word
-        _LOGGER.debug("Detected wake word: %s", wake_word_phrase)
-        self.send_messages(
-            [VoiceAssistantRequest(start=True, wake_word_phrase=wake_word_phrase)]
-        )
-        self.duck()
-        self._is_streaming_audio = True
-        self.state.tts_player.play(self.state.wakeup_sound)
-
-    def stop(self) -> None:
-        self.state.stop_word.is_active = False
-        self.state.tts_player.stop()
-
-        if self._timer_finished:
-            self._timer_finished = False
-            _LOGGER.debug("Stopping timer finished sound")
-        else:
-            _LOGGER.debug("TTS response stopped manually")
-            self._tts_finished()
-
-    def play_tts(self) -> None:
-        if (not self._tts_url) or self._tts_played:
-            return
-
-        self._tts_played = True
-        _LOGGER.debug("Playing TTS response: %s", self._tts_url)
-
-        self.state.stop_word.is_active = True
-        self.state.tts_player.play(self._tts_url, done_callback=self._tts_finished)
-
-    def duck(self) -> None:
-        _LOGGER.debug("Ducking music")
-        self.state.music_player.duck()
-
-    def unduck(self) -> None:
-        _LOGGER.debug("Unducking music")
-        self.state.music_player.unduck()
-
-    def _tts_finished(self) -> None:
-        self.state.stop_word.is_active = False
-        self.send_messages([VoiceAssistantAnnounceFinished()])
-
-        if self._continue_conversation:
-            self.send_messages([VoiceAssistantRequest(start=True)])
-            self._is_streaming_audio = True
-            _LOGGER.debug("Continuing conversation")
-        else:
-            self.unduck()
-
-        _LOGGER.debug("TTS response finished")
-
-    def _play_timer_finished(self) -> None:
-        if not self._timer_finished:
-            self.unduck()
-            return
-
-        self.state.tts_player.play(
-            self.state.timer_finished_sound,
-            done_callback=lambda: call_all(
-                lambda: time.sleep(1.0), self._play_timer_finished
-            ),
-        )
-
-    def connection_lost(self, exc):
-        super().connection_lost(exc)
-        _LOGGER.info("Disconnected from Home Assistant")
-
-
-def process_audio(state: ServerState):
-
-    try:
-        while True:
-            audio_chunk = state.audio_queue.get()
-            if audio_chunk is None:
-                break
-
-            if state.satellite is None:
-                continue
-
-            try:
-                state.satellite.handle_audio(audio_chunk)
-
-                if state.wake_word.is_active and state.wake_word.process_streaming(
-                    audio_chunk
-                ):
-                    state.satellite.wakeup()
-
-                if state.stop_word.is_active and state.stop_word.process_streaming(
-                    audio_chunk
-                ):
-                    state.satellite.stop()
-            except Exception:
-                _LOGGER.exception("Unexpected error handling audio")
-
-    except Exception:
-        _LOGGER.exception("Unexpected error processing audio")
-
-
 # -----------------------------------------------------------------------------
 
 
@@ -482,5 +154,40 @@ async def main() -> None:
     _LOGGER.debug("Server stopped")
 
 
+# -----------------------------------------------------------------------------
+
+
+def process_audio(state: ServerState):
+
+    try:
+        while True:
+            audio_chunk = state.audio_queue.get()
+            if audio_chunk is None:
+                break
+
+            if state.satellite is None:
+                continue
+
+            try:
+                state.satellite.handle_audio(audio_chunk)
+
+                if state.wake_word.is_active and state.wake_word.process_streaming(
+                    audio_chunk
+                ):
+                    state.satellite.wakeup()
+
+                if state.stop_word.is_active and state.stop_word.process_streaming(
+                    audio_chunk
+                ):
+                    state.satellite.stop()
+            except Exception:
+                _LOGGER.exception("Unexpected error handling audio")
+
+    except Exception:
+        _LOGGER.exception("Unexpected error processing audio")
+
+
+# -----------------------------------------------------------------------------
+
 if __name__ == "__main__":
     asyncio.run(main())

linux_voice_assistant/models.py

@@ -0,0 +1,37 @@
+"""Shared models."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from queue import Queue
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+if TYPE_CHECKING:
+    from .entity import ESPHomeEntity, MediaPlayerEntity
+    from .microwakeword import MicroWakeWord
+    from .mpv_player import MpvMediaPlayer
+    from .satellite import VoiceSatelliteProtocol
+
+
+@dataclass
+class AvailableWakeWord:
+    id: str
+    wake_word: str
+    trained_languages: List[str]
+    config_path: Path
+
+
+@dataclass
+class ServerState:
+    name: str
+    mac_address: str
+    audio_queue: "Queue[Optional[bytes]]"
+    entities: "List[ESPHomeEntity]"
+    available_wake_words: "Dict[str, AvailableWakeWord]"
+    wake_word: "MicroWakeWord"
+    stop_word: "MicroWakeWord"
+    music_player: "MpvMediaPlayer"
+    tts_player: "MpvMediaPlayer"
+    wakeup_sound: str
+    timer_finished_sound: str
+    media_player_entity: "Optional[MediaPlayerEntity]" = None
+    satellite: "Optional[VoiceSatelliteProtocol]" = None

linux_voice_assistant/satellite.py

@@ -0,0 +1,282 @@
+"""Voice satellite protocol."""
+
+import logging
+import time
+from collections.abc import Iterable
+from typing import Dict, Optional
+
+# pylint: disable=no-name-in-module
+from aioesphomeapi.api_pb2 import (  # type: ignore[attr-defined]
+    DeviceInfoRequest,
+    DeviceInfoResponse,
+    ListEntitiesDoneResponse,
+    ListEntitiesRequest,
+    MediaPlayerCommandRequest,
+    SubscribeHomeAssistantStatesRequest,
+    VoiceAssistantAnnounceFinished,
+    VoiceAssistantAnnounceRequest,
+    VoiceAssistantAudio,
+    VoiceAssistantConfigurationRequest,
+    VoiceAssistantConfigurationResponse,
+    VoiceAssistantEventResponse,
+    VoiceAssistantRequest,
+    VoiceAssistantSetConfiguration,
+    VoiceAssistantTimerEventResponse,
+    VoiceAssistantWakeWord,
+)
+from aioesphomeapi.model import (
+    VoiceAssistantEventType,
+    VoiceAssistantFeature,
+    VoiceAssistantTimerEventType,
+)
+from google.protobuf import message
+
+from .api_server import APIServer
+from .entity import MediaPlayerEntity
+from .microwakeword import MicroWakeWord
+from .models import ServerState
+from .util import call_all
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class VoiceSatelliteProtocol(APIServer):
+
+    def __init__(self, state: ServerState) -> None:
+        super().__init__(state.name)
+
+        self.state = state
+        self.state.satellite = self
+
+        if self.state.media_player_entity is None:
+            self.state.media_player_entity = MediaPlayerEntity(
+                server=self,
+                key=len(state.entities),
+                name="Media Player",
+                object_id="linux_voice_assistant_media_player",
+                music_player=state.music_player,
+                announce_player=state.tts_player,
+            )
+            self.state.entities.append(self.state.media_player_entity)
+
+        self._is_streaming_audio = False
+        self._tts_url: Optional[str] = None
+        self._tts_played = False
+        self._continue_conversation = False
+        self._timer_finished = False
+
+    def handle_voice_event(
+        self, event_type: VoiceAssistantEventType, data: Dict[str, str]
+    ) -> None:
+        _LOGGER.debug("Voice event: type=%s, data=%s", event_type.name, data)
+
+        if event_type == VoiceAssistantEventType.VOICE_ASSISTANT_RUN_START:
+            self._tts_url = data.get("url")
+            self._tts_played = False
+            self._continue_conversation = False
+        elif event_type in (
+            VoiceAssistantEventType.VOICE_ASSISTANT_STT_VAD_END,
+            VoiceAssistantEventType.VOICE_ASSISTANT_STT_END,
+        ):
+            self._is_streaming_audio = False
+        elif event_type == VoiceAssistantEventType.VOICE_ASSISTANT_INTENT_PROGRESS:
+            if data.get("tts_start_streaming") == "1":
+                # Start streaming early
+                self.play_tts()
+        elif event_type == VoiceAssistantEventType.VOICE_ASSISTANT_INTENT_END:
+            if data.get("continue_conversation") == "1":
+                self._continue_conversation = True
+        elif event_type == VoiceAssistantEventType.VOICE_ASSISTANT_TTS_END:
+            self._tts_url = data.get("url")
+            self.play_tts()
+        elif event_type == VoiceAssistantEventType.VOICE_ASSISTANT_RUN_END:
+            self._is_streaming_audio = False
+            if not self._tts_played:
+                self._tts_finished()
+
+            self._tts_played = False
+
+        # TODO: handle error
+
+    def handle_timer_event(
+        self,
+        event_type: VoiceAssistantTimerEventType,
+        msg: VoiceAssistantTimerEventResponse,
+    ) -> None:
+        _LOGGER.debug("Timer event: type=%s", event_type.name)
+        if event_type == VoiceAssistantTimerEventType.VOICE_ASSISTANT_TIMER_FINISHED:
+            if not self._timer_finished:
+                self.state.stop_word.is_active = True
+                self._timer_finished = True
+                self.duck()
+                self._play_timer_finished()
+
+    def handle_message(self, msg: message.Message) -> Iterable[message.Message]:
+        if isinstance(msg, VoiceAssistantEventResponse):
+            # Pipeline event
+            data: Dict[str, str] = {}
+            for arg in msg.data:
+                data[arg.name] = arg.value
+
+            self.handle_voice_event(VoiceAssistantEventType(msg.event_type), data)
+        elif isinstance(msg, VoiceAssistantAnnounceRequest):
+            _LOGGER.debug("Announcing: %s", msg.text)
+
+            assert self.state.media_player_entity is not None
+
+            urls = []
+            if msg.preannounce_media_id:
+                urls.append(msg.preannounce_media_id)
+
+            urls.append(msg.media_id)
+
+            self.state.stop_word.is_active = True
+            self._continue_conversation = msg.start_conversation
+
+            self.duck()
+            yield from self.state.media_player_entity.play(
+                urls, announcement=True, done_callback=self._tts_finished
+            )
+        elif isinstance(msg, VoiceAssistantTimerEventResponse):
+            self.handle_timer_event(VoiceAssistantTimerEventType(msg.event_type), msg)
+        elif isinstance(msg, DeviceInfoRequest):
+            yield DeviceInfoResponse(
+                uses_password=False,
+                name=self.state.name,
+                mac_address=self.state.mac_address,
+                voice_assistant_feature_flags=(
+                    VoiceAssistantFeature.VOICE_ASSISTANT
+                    | VoiceAssistantFeature.API_AUDIO
+                    | VoiceAssistantFeature.ANNOUNCE
+                    | VoiceAssistantFeature.START_CONVERSATION
+                    | VoiceAssistantFeature.TIMERS
+                ),
+            )
+        elif isinstance(
+            msg,
+            (
+                ListEntitiesRequest,
+                SubscribeHomeAssistantStatesRequest,
+                MediaPlayerCommandRequest,
+            ),
+        ):
+            for entity in self.state.entities:
+                yield from entity.handle_message(msg)
+
+            if isinstance(msg, ListEntitiesRequest):
+                yield ListEntitiesDoneResponse()
+        elif isinstance(msg, VoiceAssistantConfigurationRequest):
+            yield VoiceAssistantConfigurationResponse(
+                available_wake_words=[
+                    VoiceAssistantWakeWord(
+                        id=ww.id,
+                        wake_word=ww.wake_word,
+                        trained_languages=ww.trained_languages,
+                    )
+                    for ww in self.state.available_wake_words.values()
+                ],
+                active_wake_words=[self.state.wake_word.id],
+                max_active_wake_words=1,
+            )
+            _LOGGER.info("Connected to Home Assistant")
+        elif isinstance(msg, VoiceAssistantSetConfiguration):
+            # TODO: support multiple wake words
+            for wake_word_id in msg.active_wake_words:
+                if wake_word_id == self.state.wake_word.id:
+                    # Already active
+                    break
+
+                model_info = self.state.available_wake_words.get(wake_word_id)
+                if not model_info:
+                    continue
+
+                _LOGGER.debug("Loading wake word: %s", model_info.config_path)
+                self.state.wake_word = MicroWakeWord.from_config(
+                    model_info.config_path,
+                    self.state.wake_word.libtensorflowlite_c_path,
+                )
+
+                _LOGGER.info("Wake word set: %s", self.state.wake_word.wake_word)
+                break
+
+    def handle_audio(self, audio_chunk: bytes) -> None:
+
+        if not self._is_streaming_audio:
+            return
+
+        self.send_messages([VoiceAssistantAudio(data=audio_chunk)])
+
+    def wakeup(self) -> None:
+        if self._timer_finished:
+            # Stop timer instead
+            self._timer_finished = False
+            self.state.tts_player.stop()
+            _LOGGER.debug("Stopping timer finished sound")
+            return
+
+        wake_word_phrase = self.state.wake_word.wake_word
+        _LOGGER.debug("Detected wake word: %s", wake_word_phrase)
+        self.send_messages(
+            [VoiceAssistantRequest(start=True, wake_word_phrase=wake_word_phrase)]
+        )
+        self.duck()
+        self._is_streaming_audio = True
+        self.state.tts_player.play(self.state.wakeup_sound)
+
+    def stop(self) -> None:
+        self.state.stop_word.is_active = False
+        self.state.tts_player.stop()
+
+        if self._timer_finished:
+            self._timer_finished = False
+            _LOGGER.debug("Stopping timer finished sound")
+        else:
+            _LOGGER.debug("TTS response stopped manually")
+            self._tts_finished()
+
+    def play_tts(self) -> None:
+        if (not self._tts_url) or self._tts_played:
+            return
+
+        self._tts_played = True
+        _LOGGER.debug("Playing TTS response: %s", self._tts_url)
+
+        self.state.stop_word.is_active = True
+        self.state.tts_player.play(self._tts_url, done_callback=self._tts_finished)
+
+    def duck(self) -> None:
+        _LOGGER.debug("Ducking music")
+        self.state.music_player.duck()
+
+    def unduck(self) -> None:
+        _LOGGER.debug("Unducking music")
+        self.state.music_player.unduck()
+
+    def _tts_finished(self) -> None:
+        self.state.stop_word.is_active = False
+        self.send_messages([VoiceAssistantAnnounceFinished()])
+
+        if self._continue_conversation:
+            self.send_messages([VoiceAssistantRequest(start=True)])
+            self._is_streaming_audio = True
+            _LOGGER.debug("Continuing conversation")
+        else:
+            self.unduck()
+
+        _LOGGER.debug("TTS response finished")
+
+    def _play_timer_finished(self) -> None:
+        if not self._timer_finished:
+            self.unduck()
+            return
+
+        self.state.tts_player.play(
+            self.state.timer_finished_sound,
+            done_callback=lambda: call_all(
+                lambda: time.sleep(1.0), self._play_timer_finished
+            ),
+        )
+
+    def connection_lost(self, exc):
+        super().connection_lost(exc)
+        _LOGGER.info("Disconnected from Home Assistant")

mypy.ini

@@ -18,3 +18,6 @@ ignore_missing_imports = True
 
 [mypy-google.protobuf.*]
 ignore_missing_imports = True
+
+[mypy-aioesphomeapi.*]
+ignore_missing_imports = True