Skip to content

Commit

Permalink
Add get/set/config messages for voice assistant (#956)
Browse files Browse the repository at this point in the history
  • Loading branch information
synesthesiam committed Sep 13, 2024
1 parent bf6aff7 commit 709cb3f
Show file tree
Hide file tree
Showing 7 changed files with 332 additions and 109 deletions.
30 changes: 30 additions & 0 deletions aioesphomeapi/api.proto
Original file line number Diff line number Diff line change
Expand Up @@ -1589,6 +1589,36 @@ message VoiceAssistantAnnounceFinished {
bool success = 1;
}

message VoiceAssistantWakeWord {
uint32 id = 1;
string wake_word = 2;
repeated string trained_languages = 3;
}

message VoiceAssistantConfigurationRequest {
option (id) = 121;
option (source) = SOURCE_CLIENT;
option (ifdef) = "USE_VOICE_ASSISTANT";
}

message VoiceAssistantConfigurationResponse {
option (id) = 122;
option (source) = SOURCE_SERVER;
option (ifdef) = "USE_VOICE_ASSISTANT";

repeated VoiceAssistantWakeWord available_wake_words = 1;
repeated uint32 active_wake_words = 2;
uint32 max_active_wake_words = 3;
}

message VoiceAssistantSetConfiguration {
option (id) = 123;
option (source) = SOURCE_CLIENT;
option (ifdef) = "USE_VOICE_ASSISTANT";

repeated uint32 active_wake_words = 1;
}

// ==================== ALARM CONTROL PANEL ====================
enum AlarmControlPanelState {
ALARM_STATE_DISARMED = 0;
Expand Down
264 changes: 155 additions & 109 deletions aioesphomeapi/api_pb2.py

Large diffs are not rendered by default.

20 changes: 20 additions & 0 deletions aioesphomeapi/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,10 +73,13 @@
VoiceAssistantAnnounceFinished,
VoiceAssistantAnnounceRequest,
VoiceAssistantAudio,
VoiceAssistantConfigurationRequest,
VoiceAssistantConfigurationResponse,
VoiceAssistantEventData,
VoiceAssistantEventResponse,
VoiceAssistantRequest,
VoiceAssistantResponse,
VoiceAssistantSetConfiguration,
VoiceAssistantTimerEventResponse,
)
from .client_callbacks import (
Expand Down Expand Up @@ -133,6 +136,7 @@
VoiceAssistantAudioData,
VoiceAssistantAudioSettings as VoiceAssistantAudioSettingsModel,
VoiceAssistantCommand,
VoiceAssistantConfigurationResponse as VoiceAssistantConfigurationResponseModel,
VoiceAssistantEventType,
VoiceAssistantSubscriptionFlag,
VoiceAssistantTimerEventType,
Expand Down Expand Up @@ -1459,6 +1463,22 @@ async def send_voice_assistant_announcement_await_response(
)
return VoiceAssistantAnnounceFinishedModel.from_pb(resp)

async def get_voice_assistant_configuration(
self, timeout: float
) -> VoiceAssistantConfigurationResponseModel:
resp = await self._get_connection().send_message_await_response(
VoiceAssistantConfigurationRequest(),
VoiceAssistantConfigurationResponse,
timeout,
)
return VoiceAssistantConfigurationResponseModel.from_pb(resp)

async def set_voice_assistant_configuration(
self, active_wake_words: list[int]
) -> None:
req = VoiceAssistantSetConfiguration(active_wake_words=active_wake_words)
self._get_connection().send_message(req)

def alarm_control_panel_command(
self,
key: int,
Expand Down
6 changes: 6 additions & 0 deletions aioesphomeapi/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,9 +121,12 @@
VoiceAssistantAnnounceFinished,
VoiceAssistantAnnounceRequest,
VoiceAssistantAudio,
VoiceAssistantConfigurationRequest,
VoiceAssistantConfigurationResponse,
VoiceAssistantEventResponse,
VoiceAssistantRequest,
VoiceAssistantResponse,
VoiceAssistantSetConfiguration,
VoiceAssistantTimerEventResponse,
)

Expand Down Expand Up @@ -396,6 +399,9 @@ def __init__(self, error: BluetoothGATTError) -> None:
118: UpdateCommandRequest,
119: VoiceAssistantAnnounceRequest,
120: VoiceAssistantAnnounceFinished,
121: VoiceAssistantConfigurationRequest,
122: VoiceAssistantConfigurationResponse,
123: VoiceAssistantSetConfiguration,
}

MESSAGE_NUMBER_TO_PROTO = tuple(MESSAGE_TYPE_TO_PROTO.values())
36 changes: 36 additions & 0 deletions aioesphomeapi/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -1298,6 +1298,42 @@ class VoiceAssistantAnnounceFinished(APIModelBase):
success: bool = False


@_frozen_dataclass_decorator
class VoiceAssistantWakeWord(APIModelBase):
id: int
wake_word: str
trained_languages: list[str]

@classmethod
def convert_list(cls, value: list[Any]) -> list[VoiceAssistantWakeWord]:
ret = []
for x in value:
if isinstance(x, dict):
ret.append(VoiceAssistantWakeWord.from_dict(x))
else:
ret.append(VoiceAssistantWakeWord.from_pb(x))
return ret


@_frozen_dataclass_decorator
class VoiceAssistantConfigurationResponse(APIModelBase):
available_wake_words: list[VoiceAssistantWakeWord] = converter_field(
default_factory=list, converter=VoiceAssistantWakeWord.convert_list
)
active_wake_words: list[int] = converter_field(default_factory=list, converter=list)
max_active_wake_words: int = 0


@_frozen_dataclass_decorator
class VoiceAssistantConfigurationRequest(APIModelBase):
pass


@_frozen_dataclass_decorator
class VoiceAssistantSetConfiguration(APIModelBase):
active_wake_words: list[int] = converter_field(default_factory=list, converter=list)


class LogLevel(APIIntEnum):
LOG_LEVEL_NONE = 0
LOG_LEVEL_ERROR = 1
Expand Down
56 changes: 56 additions & 0 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,11 +71,15 @@
VoiceAssistantAnnounceRequest,
VoiceAssistantAudio,
VoiceAssistantAudioSettings,
VoiceAssistantConfigurationRequest,
VoiceAssistantConfigurationResponse,
VoiceAssistantEventData,
VoiceAssistantEventResponse,
VoiceAssistantRequest,
VoiceAssistantResponse,
VoiceAssistantSetConfiguration,
VoiceAssistantTimerEventResponse,
VoiceAssistantWakeWord,
)
from aioesphomeapi.client import APIClient, BluetoothConnectionDroppedError
from aioesphomeapi.connection import APIConnection
Expand Down Expand Up @@ -113,6 +117,7 @@
UserServiceArgType,
VoiceAssistantAnnounceFinished as VoiceAssistantAnnounceFinishedModel,
VoiceAssistantAudioSettings as VoiceAssistantAudioSettingsModel,
VoiceAssistantConfigurationResponse as VoiceAssistantConfigurationResponseModel,
VoiceAssistantEventType as VoiceAssistantEventModelType,
VoiceAssistantTimerEventType as VoiceAssistantTimerEventModelType,
)
Expand Down Expand Up @@ -2660,6 +2665,57 @@ async def handle_announcement_finished(
assert len(send.mock_calls) == 0


@pytest.mark.asyncio
async def test_get_voice_assistant_configuration(
api_client: tuple[
APIClient, APIConnection, asyncio.Transport, APIPlaintextFrameHelper
],
) -> None:
client, connection, _transport, protocol = api_client
original_send_message = connection.send_message

def send_message(msg):
assert msg == VoiceAssistantConfigurationRequest()
original_send_message(msg)

with patch.object(connection, "send_message", new=send_message):
config_task = asyncio.create_task(
client.get_voice_assistant_configuration(timeout=1.0)
)
await asyncio.sleep(0)
response: message.Message = VoiceAssistantConfigurationResponse(
available_wake_words=[
VoiceAssistantWakeWord(
id=1,
wake_word="okay nabu",
trained_languages=["en"],
)
],
active_wake_words=[1],
max_active_wake_words=1,
)
mock_data_received(protocol, generate_plaintext_packet(response))
config = await config_task
assert isinstance(config, VoiceAssistantConfigurationResponseModel)


@pytest.mark.asyncio
async def test_set_voice_assistant_configuration(
api_client: tuple[
APIClient, APIConnection, asyncio.Transport, APIPlaintextFrameHelper
],
) -> None:
client, connection, _transport, protocol = api_client
original_send_message = connection.send_message

def send_message(msg):
assert msg == VoiceAssistantSetConfiguration(active_wake_words=[1])
original_send_message(msg)

with patch.object(connection, "send_message", new=send_message):
await client.set_voice_assistant_configuration([1])


@pytest.mark.asyncio
async def test_api_version_after_connection_closed(
api_client: tuple[
Expand Down
29 changes: 29 additions & 0 deletions tests/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,9 @@
UserServiceArgType,
ValveInfo,
ValveState,
VoiceAssistantConfigurationResponse,
VoiceAssistantFeature,
VoiceAssistantWakeWord,
build_unique_id,
converter_field,
)
Expand Down Expand Up @@ -681,3 +683,30 @@ def test_media_player_supported_format_convert_list() -> None:
)
],
)


def test_voice_assistant_wake_word_convert_list() -> None:
"""Test list conversion for VoiceAssistantWakeWord."""
assert VoiceAssistantConfigurationResponse.from_dict(
{
"available_wake_words": [
{
"id": 1,
"wake_word": "okay nabu",
"trained_languages": ["en"],
}
],
"active_wake_words": [1],
"max_active_wake_words": 1,
}
) == VoiceAssistantConfigurationResponse(
available_wake_words=[
VoiceAssistantWakeWord(
id=1,
wake_word="okay nabu",
trained_languages=["en"],
)
],
active_wake_words=[1],
max_active_wake_words=1,
)

0 comments on commit 709cb3f

Please sign in to comment.