[WIP] AWS Nova Sonic service - do hacky direct manipulation of the context for now, since I can't seem to get assistant context aggregation working properly with frames, grr

src/pipecat/services/aws_nova_sonic/aws.py

@@ -147,7 +147,6 @@ class AWSNovaSonicLLMService(LLMService):
         self._ready_to_send_context = False
         self._handling_bot_stopped_speaking = False
 
-
     #
     # standard AIService frame handling
     #
@@ -760,8 +759,10 @@ 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.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}")
 
         # Bookkeeping: clear current content being received
         self._content_being_received = None
@@ -803,6 +804,18 @@ class AWSNovaSonicLLMService(LLMService):
         print(f"[pk] TTS text: {text}")
         await self.push_frame(TTSTextFrame(text))
 
+        # TODO: this is a (hopefully temporary) HACK. Here we directly manipulate the context rather
+        # than relying on the frames pushed to the assistant context aggregator. The pattern of
+        # receiving full-sentence text after the assistant has spoken does not easily fit with the
+        # Pipecat expectation of chunks of text streaming in while the assistant is speaking.
+        # Interruption handling was especially challenging. Rather than spend days trying to fit a
+        # square peg in a round hole, I decided on this hack for the time being. We can most cleanly
+        # abandon this hack if/when AWS Nova Sonic implements streaming smaller text chunks
+        # 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)
+
     async def _report_assistant_response_ended(self):
         # Report that the assistant has finished their response.
         print("[pk] LLM full response ended")

src/pipecat/services/aws_nova_sonic/context.py

@@ -11,13 +11,19 @@ from enum import Enum
 from loguru import logger
 
 from pipecat.frames.frames import (
+    BotStoppedSpeakingFrame,
     DataFrame,
     Frame,
     FunctionCallResultFrame,
+    LLMFullResponseEndFrame,
+    LLMFullResponseStartFrame,
+    LLMMessagesAppendFrame,
     LLMMessagesUpdateFrame,
+    LLMSetToolChoiceFrame,
     LLMSetToolsFrame,
-    LLMTextFrame,
-    TranscriptionFrame,
+    StartInterruptionFrame,
+    TextFrame,
+    UserImageRawFrame,
 )
 from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext
 from pipecat.processors.frame_processor import FrameDirection
@@ -110,6 +116,15 @@ class AWSNovaSonicLLMContext(OpenAILLMContext):
             "content": [{"type": "text", "text": text}],
         }
         self.add_message(message)
+        # print(f"[pk] context updated (user): {self.get_messages_for_logging()}")
+
+    def add_assistant_text_as_message(self, text):
+        message = {
+            "role": "assistant",
+            "content": [{"type": "text", "text": text}],
+        }
+        self.add_message(message)
+        # print(f"[pk] context updated (assistant): {self.get_messages_for_logging()}")
 
 
 @dataclass
@@ -134,21 +149,28 @@ class AWSNovaSonicUserContextAggregator(OpenAIUserContextAggregator):
 
 
 class AWSNovaSonicAssistantContextAggregator(OpenAIAssistantContextAggregator):
-    # AWS Nova Sonic is a speech-to-speech model.
-    # It behaves like a combined STT + LLM + TTS service, emitting all of:
-    # - TranscriptionFrame (for user text)
-    # - LLMTextFrame (for assistant text)
-    # - TTSTextFrame (for assistant text)
-    # In a "standard" pipeline (with separate STT + LLM + TTS services):
-    # - The TranscriptionFrame is swallowed by the LLMUserContextAggregator
-    # - The LLMTextFrame is swallowed by the TTS service
-    # Meaning the LLMAssistantContextAggregator only receives the TTSTextFrames. It actually
-    # implicitly  assumes it will receive only *non-duplicate* *assistant-related* text frames, and
-    # will misbehave otherwise (double-counting assistant text, or mis-categorizing user text as
-    # assistant text).
-    # So, let's override process_frame here to ignore TranscriptionFrames and LLMTextFrames.
     async def process_frame(self, frame: Frame, direction: FrameDirection):
-        if not isinstance(frame, (LLMTextFrame, TranscriptionFrame)):
+        # HACK: For now, disable the context aggregator by making it just pass through all frames
+        # that the parent handles (except the function call stuff, which we still need).
+        # For an explanation of this hack, see
+        # AWSNovaSonicLLMService._report_assistant_response_text_added.
+        if isinstance(
+            frame,
+            (
+                StartInterruptionFrame,
+                LLMFullResponseStartFrame,
+                LLMFullResponseEndFrame,
+                TextFrame,
+                LLMMessagesAppendFrame,
+                LLMMessagesUpdateFrame,
+                LLMSetToolsFrame,
+                LLMSetToolChoiceFrame,
+                UserImageRawFrame,
+                BotStoppedSpeakingFrame,
+            ),
+        ):
+            await self.push_frame(frame, direction)
+        else:
             await super().process_frame(frame, direction)
 
     async def handle_function_call_result(self, frame: FunctionCallResultFrame):