[WIP] AWS Nova Sonic service - in our hacky direct manipulation of the context, aggregate assistant text rather than recording every chunk as a separate message

src/pipecat/services/aws_nova_sonic/aws.py

@@ -759,10 +759,8 @@ class AWSNovaSonicLLMService(LLMService):
         content_end = event_json["contentEnd"]
         stop_reason = content_end["stopReason"]
         # print(f"[pk] content end: {content}.\n  stop_reason: {stop_reason}")
-        if content.role == Role.ASSISTANT:
-            # print(f"[pk] assistant content end: {content}.\n  stop_reason: {stop_reason}")
-            if content.text_stage == TextStage.FINAL:
-                print(f"[pk] assistant FINAL text: {content.text_content}")
+        # if content.role == Role.ASSISTANT:
+        # print(f"[pk] assistant content end: {content}.\n  stop_reason: {stop_reason}")
 
         # Bookkeeping: clear current content being received
         self._content_being_received = None
@@ -814,7 +812,7 @@ class AWSNovaSonicLLMService(LLMService):
         # interspersed with audio. Note that when we move away from this hack, we need to make sure
         # that on an interruption we avoid sending LLMFullResponseEndFrame, which gets the
         # LLMAssistantContextAggregator into a bad state.
-        self._context.add_assistant_text_as_message(text)
+        self._context.buffer_assistant_text(text)
 
     async def _report_assistant_response_ended(self):
         # Report that the assistant has finished their response.
@@ -825,11 +823,14 @@ class AWSNovaSonicLLMService(LLMService):
         print("[pk] TTS stopped")
         await self.push_frame(TTSStoppedFrame())
 
+        # For an explanation of this hack, see _report_assistant_response_text_added.
+        self._context.flush_aggregated_assistant_text()
+
     async def _report_user_transcription_text_added(self, text):
         print(f"[pk] transcription: {text}")
         # Manually add new user transcription text to context.
         # We can't rely on the user context aggregator to do this since it's upstream from the LLM.
-        self._context.add_user_transcription_text_as_message(text)
+        self._context.add_user_transcription_text(text)
 
         # Report that some new user transcription text is available.
         if self._send_transcription_frames:

src/pipecat/services/aws_nova_sonic/context.py

@@ -53,12 +53,19 @@ class AWSNovaSonicConversationHistory:
     messages: list[AWSNovaSonicConversationHistoryMessage] = field(default_factory=list)
 
 
-@dataclass
 class AWSNovaSonicLLMContext(OpenAILLMContext):
+    def __init__(self, messages=None, tools=None, **kwargs):
+        super().__init__(messages=messages, tools=tools, **kwargs)
+        self.__setup_local()
+
+    def __setup_local(self):
+        self._assistant_text = ""
+
     @staticmethod
     def upgrade_to_nova_sonic(obj: OpenAILLMContext) -> "AWSNovaSonicLLMContext":
         if isinstance(obj, OpenAILLMContext) and not isinstance(obj, AWSNovaSonicLLMContext):
             obj.__class__ = AWSNovaSonicLLMContext
+            obj.__setup_local()
         return obj
 
     def get_messages_for_initializing_history(self) -> AWSNovaSonicConversationHistory:
@@ -110,7 +117,7 @@ class AWSNovaSonicLLMContext(OpenAILLMContext):
         # NOTE: we're ignoring messages with role "tool" since they can't be loaded into AWS Nova
         # Sonic conversation history
 
-    def add_user_transcription_text_as_message(self, text):
+    def add_user_transcription_text(self, text):
         message = {
             "role": "user",
             "content": [{"type": "text", "text": text}],
@@ -118,11 +125,16 @@ class AWSNovaSonicLLMContext(OpenAILLMContext):
         self.add_message(message)
         # print(f"[pk] context updated (user): {self.get_messages_for_logging()}")
 
-    def add_assistant_text_as_message(self, text):
+    def buffer_assistant_text(self, text):
+        self._assistant_text += text  # TODO: determine if we need to add space or something
+        # print(f"[pk] assistant text buffered: {self._assistant_text}")
+
+    def flush_aggregated_assistant_text(self):
         message = {
             "role": "assistant",
-            "content": [{"type": "text", "text": text}],
+            "content": [{"type": "text", "text": self._assistant_text}],
         }
+        self._assistant_text = ""
         self.add_message(message)
         # print(f"[pk] context updated (assistant): {self.get_messages_for_logging()}")