[WIP] AWS Nova Sonic service - add more accurate typing

This commit is contained in:
Paul Kompfner
2025-05-07 09:52:05 -04:00
parent 27bff7a759
commit 4ba9a42861

View File

@@ -148,30 +148,34 @@ class AWSNovaSonicLLMService(LLMService):
self._access_key_id = access_key_id self._access_key_id = access_key_id
self._region = region self._region = region
self._model = model self._model = model
self._client: BedrockRuntimeClient = None self._client: Optional[BedrockRuntimeClient] = None
self._voice_id = voice_id self._voice_id = voice_id
self._params = params self._params = params
self._system_instruction = system_instruction self._system_instruction = system_instruction
self._tools = tools self._tools = tools
self._send_transcription_frames = send_transcription_frames self._send_transcription_frames = send_transcription_frames
self._context: AWSNovaSonicLLMContext = None self._context: Optional[AWSNovaSonicLLMContext] = None
self._stream: DuplexEventStream[ self._stream: Optional[
InvokeModelWithBidirectionalStreamInput, DuplexEventStream[
InvokeModelWithBidirectionalStreamOutput, InvokeModelWithBidirectionalStreamInput,
InvokeModelWithBidirectionalStreamOperationOutput, InvokeModelWithBidirectionalStreamOutput,
InvokeModelWithBidirectionalStreamOperationOutput,
]
] = None ] = None
self._receive_task = None self._receive_task: Optional[asyncio.Task] = None
self._prompt_name = None self._prompt_name: Optional[str] = None
self._input_audio_content_name = None self._input_audio_content_name: Optional[str] = None
self._content_being_received = None self._content_being_received: Optional[CurrentContent] = None
self._assistant_is_responding = False self._assistant_is_responding = False
self._context_available = False self._context_available = False
self._ready_to_send_context = False self._ready_to_send_context = False
self._handling_bot_stopped_speaking = False self._handling_bot_stopped_speaking = False
self._triggering_assistant_response = False self._triggering_assistant_response = False
self._assistant_response_trigger_audio: bytes = None # Not cleared on _disconnect() self._assistant_response_trigger_audio: Optional[bytes] = (
None # Not cleared on _disconnect()
)
self._disconnecting = False self._disconnecting = False
self._connected_time: float = None self._connected_time: Optional[float] = None
self._wants_connection = False self._wants_connection = False
# #
@@ -437,6 +441,9 @@ class AWSNovaSonicLLMService(LLMService):
await self._send_client_event(session_start) await self._send_client_event(session_start)
async def _send_prompt_start_event(self, tools: List[Any]): async def _send_prompt_start_event(self, tools: List[Any]):
if not self._prompt_name:
return
tools_config = ( tools_config = (
f""", f""",
"toolUseOutputConfiguration": {{ "toolUseOutputConfiguration": {{
@@ -474,6 +481,9 @@ class AWSNovaSonicLLMService(LLMService):
await self._send_client_event(prompt_start) await self._send_client_event(prompt_start)
async def _send_audio_input_start_event(self): async def _send_audio_input_start_event(self):
if not self._prompt_name:
return
audio_content_start = f''' audio_content_start = f'''
{{ {{
"event": {{ "event": {{
@@ -498,7 +508,7 @@ class AWSNovaSonicLLMService(LLMService):
await self._send_client_event(audio_content_start) await self._send_client_event(audio_content_start)
async def _send_text_event(self, text: str, role: Role): async def _send_text_event(self, text: str, role: Role):
if not self._stream or not text: if not self._stream or not self._prompt_name or not text:
return return
content_name = str(uuid.uuid4()) content_name = str(uuid.uuid4())
@@ -566,7 +576,7 @@ class AWSNovaSonicLLMService(LLMService):
await self._send_client_event(audio_event) await self._send_client_event(audio_event)
async def _send_session_end_events(self): async def _send_session_end_events(self):
if not self._stream: if not self._stream or not self._prompt_name:
return return
prompt_end = f''' prompt_end = f'''
@@ -590,7 +600,7 @@ class AWSNovaSonicLLMService(LLMService):
await self._send_client_event(session_end) await self._send_client_event(session_end)
async def _send_tool_result(self, tool_call_id, result): async def _send_tool_result(self, tool_call_id, result):
if not self._stream: if not self._stream or not self._prompt_name:
return return
content_name = str(uuid.uuid4()) content_name = str(uuid.uuid4())
@@ -643,6 +653,9 @@ class AWSNovaSonicLLMService(LLMService):
await self._send_client_event(result_content_end) await self._send_client_event(result_content_end)
async def _send_client_event(self, event_json: str): async def _send_client_event(self, event_json: str):
if not self._stream: # should never happen
return
event = InvokeModelWithBidirectionalStreamInputChunk( event = InvokeModelWithBidirectionalStreamInputChunk(
value=BidirectionalInputPayloadPart(bytes_=event_json.encode("utf-8")) value=BidirectionalInputPayloadPart(bytes_=event_json.encode("utf-8"))
) )
@@ -732,8 +745,7 @@ class AWSNovaSonicLLMService(LLMService):
await self._report_assistant_response_started() await self._report_assistant_response_started()
async def _handle_text_output_event(self, event_json): async def _handle_text_output_event(self, event_json):
# This should never happen if not self._content_being_received: # should never happen
if not self._content_being_received:
return return
content = self._content_being_received content = self._content_being_received
@@ -744,8 +756,7 @@ class AWSNovaSonicLLMService(LLMService):
content.text_content = text_content content.text_content = text_content
async def _handle_audio_output_event(self, event_json): async def _handle_audio_output_event(self, event_json):
# This should never happen if not self._content_being_received: # should never happen
if not self._content_being_received:
return return
# Get audio # Get audio
@@ -761,8 +772,7 @@ class AWSNovaSonicLLMService(LLMService):
await self.push_frame(frame) await self.push_frame(frame)
async def _handle_tool_use_event(self, event_json): async def _handle_tool_use_event(self, event_json):
# This should never happen if not self._content_being_received or not self._context: # should never happen
if not self._content_being_received:
return return
# Get tool use details # Get tool use details
@@ -793,8 +803,7 @@ class AWSNovaSonicLLMService(LLMService):
) )
async def _handle_content_end_event(self, event_json): async def _handle_content_end_event(self, event_json):
# This should never happen if not self._content_being_received: # should never happen
if not self._content_being_received:
return return
content = self._content_being_received content = self._content_being_received
@@ -817,8 +826,6 @@ class AWSNovaSonicLLMService(LLMService):
# User transcription text added # User transcription text added
await self._report_user_transcription_text_added(content.text_content) await self._report_user_transcription_text_added(content.text_content)
self._content_being_received = False
async def _handle_completion_end_event(self, event_json): async def _handle_completion_end_event(self, event_json):
pass pass
@@ -832,6 +839,9 @@ class AWSNovaSonicLLMService(LLMService):
await self.push_frame(TTSStartedFrame()) await self.push_frame(TTSStartedFrame())
async def _report_assistant_response_text_added(self, text): async def _report_assistant_response_text_added(self, text):
if not self._context: # should never happen
return
logger.debug(f"Assistant response text added: {text}") logger.debug(f"Assistant response text added: {text}")
# Report some text added to the ongoing assistant response # Report some text added to the ongoing assistant response
@@ -853,6 +863,9 @@ class AWSNovaSonicLLMService(LLMService):
self._context.buffer_assistant_text(text) self._context.buffer_assistant_text(text)
async def _report_assistant_response_ended(self): async def _report_assistant_response_ended(self):
if not self._context: # should never happen
return
logger.debug("Assistant response ended") logger.debug("Assistant response ended")
# Report that the assistant has finished their response. # Report that the assistant has finished their response.
@@ -865,6 +878,9 @@ class AWSNovaSonicLLMService(LLMService):
self._context.flush_aggregated_assistant_text() self._context.flush_aggregated_assistant_text()
async def _report_user_transcription_text_added(self, text): async def _report_user_transcription_text_added(self, text):
if not self._context: # should never happen
return
logger.debug(f"User transcription text added: {text}") logger.debug(f"User transcription text added: {text}")
# Manually add new user transcription text to context. # Manually add new user transcription text to context.
@@ -918,12 +934,16 @@ class AWSNovaSonicLLMService(LLMService):
self._assistant_response_trigger_audio = wav_file.readframes(wav_file.getnframes()) self._assistant_response_trigger_audio = wav_file.readframes(wav_file.getnframes())
# Send the trigger audio, if we're fully connected and set up # Send the trigger audio, if we're fully connected and set up
# NOTE: maybe there's a better way to determine whether we're done setting up? if self._connected_time is not None:
if self._receive_task:
await self._send_assistant_response_trigger() await self._send_assistant_response_trigger()
self._triggering_assistant_response = False self._triggering_assistant_response = False
async def _send_assistant_response_trigger(self): async def _send_assistant_response_trigger(self):
if (
not self._assistant_response_trigger_audio or self._connected_time is None
): # should never happen
return
logger.debug("Sending assistant response trigger...") logger.debug("Sending assistant response trigger...")
chunk_duration = 0.02 # what we might get from InputAudioRawFrame chunk_duration = 0.02 # what we might get from InputAudioRawFrame