add cancel and fix a bunch of async problems

charles_actor.py

@@ -32,7 +32,7 @@ class CharlesActor:
         self._respond_to_prompt_actor = RespondToPromptActor.remote()
         
         self._debug_queue = [
-            # "hello, how are you today?",
+            "hello, how are you today?",
             # "hmm, interesting, tell me more about that.",
         ]
         print("010")
@@ -55,7 +55,7 @@ class CharlesActor:
         while True:
             if len(self._debug_queue) > 0:
                 prompt = self._debug_queue.pop(0)
-                await self._respond_to_prompt_actor.enqueue(prompt)
+                self._respond_to_prompt_actor.enqueue_prompt.remote(prompt)
             audio_frames = await self._streamlit_av_queue.get_audio_frames_async()    
             if len(audio_frames) > 0:
                 total_audio_frames += len(audio_frames)

respond_to_prompt_actor.py

@@ -3,7 +3,7 @@ from ray.util.queue import Queue
 from dotenv import load_dotenv
 from audio_stream_processor import AudioStreamProcessor
 from streaming_chat_service import StreamingChatService
-
+import asyncio
 # from ray.actor import ActorHandle
 
 @ray.remote
@@ -14,15 +14,25 @@ class PromptToLLMActor:
         self.output_queue = output_queue
         self.audio_processor = AudioStreamProcessor()
         self.chat_service = StreamingChatService(self.audio_processor, voice_id=voice_id)
+        self.cancel_event = None
 
     async def run(self):
         while True:
-            prompt = self.input_queue.get()
-            async for sentence in self.chat_service.get_responses_as_sentances_async(prompt):
+            prompt = await self.input_queue.get_async()
+            self.cancel_event = asyncio.Event()
+            async for sentence in self.chat_service.get_responses_as_sentances_async(prompt, self.cancel_event):
                 if self.chat_service.ignore_sentence(sentence):
                     continue
                 print(f"{sentence}")
-                self.output_queue.put(sentence)
+                await self.output_queue.put_async(sentence)
+                
+    def cancel(self):
+        if self.cancel_event:
+            self.cancel_event.set()
+        while not self.input_queue.empty():
+            self.input_queue.get()
+        while not self.output_queue.empty():
+            self.output_queue.get()
 
 @ray.remote
 class LLMSentanceToSpeechActor:
@@ -32,12 +42,23 @@ class LLMSentanceToSpeechActor:
         self.output_queue = output_queue
         self.audio_processor = AudioStreamProcessor()
         self.chat_service = StreamingChatService(self.audio_processor, voice_id=voice_id)
+        self.cancel_event = None
 
     async def run(self):
         while True:
-            sentance = self.input_queue.get()
-            async for chunk in self.chat_service.get_speech_chunks_async(sentance):
-                self.output_queue.put(chunk)
+            sentance = await self.input_queue.get_async()
+            self.cancel_event = asyncio.Event()
+            async for chunk in self.chat_service.get_speech_chunks_async(sentance, self.cancel_event):
+                await self.output_queue.put_async(chunk)
+
+    def cancel(self):
+        if self.cancel_event:
+            self.cancel_event.set()
+        while not self.input_queue.empty():
+            self.input_queue.get()
+        while not self.output_queue.empty():
+            self.output_queue.get()
+
 
 @ray.remote
 class SpeechToSpeakerActor:
@@ -49,8 +70,12 @@ class SpeechToSpeakerActor:
 
     async def run(self):
         while True:
-            audio_chunk = self.input_queue.get()
+            audio_chunk = await self.input_queue.get_async()
             self.chat_service.enqueue_speech_bytes_to_play([audio_chunk])
+            
+    def cancel(self):
+        while not self.input_queue.empty():
+            self.input_queue.get()            
 
 @ray.remote
 class RespondToPromptActor:
@@ -65,13 +90,20 @@ class RespondToPromptActor:
         self.speech_to_speaker = SpeechToSpeakerActor.remote(self.speech_chunk_queue, voice_id)
 
         # Start the pipeline components.
-        print ("Starting pipeline components")
         self.prompt_to_llm.run.remote()
-        print ("prompt_to_llm running")
         self.llm_sentence_to_speech.run.remote()
-        print ("llm_sentence_to_speech running")
         self.speech_to_speaker.run.remote()
-        print ("speech_to_speaker running")
             
     def enqueue_prompt(self, prompt):
-        self.prompt_queue.put(prompt)
+        print("flush anything queued")
+        prompt_to_llm_future = self.prompt_to_llm.cancel.remote()
+        llm_sentence_to_speech_future = self.llm_sentence_to_speech.cancel.remote()
+        speech_to_speaker_future = self.speech_to_speaker.cancel.remote()
+        ray.get([
+            prompt_to_llm_future,
+            llm_sentence_to_speech_future,
+            speech_to_speaker_future,
+        ])
+        self.prompt_queue.put(prompt)
+        print("Enqueued prompt")
+        

streaming_chat_service.py

@@ -144,7 +144,7 @@ I fell off the pink step, and I had an accident.
         self._messages.append({"role": "assistant", "content": agent_response})
         return agent_response
 
-    async def get_responses_as_sentances_async(self, prompt):
+    async def get_responses_as_sentances_async(self, prompt, cancel_event):
         self._messages.append({"role": "user", "content": prompt})
         agent_response = ""
         current_sentence = ""
@@ -157,6 +157,8 @@ I fell off the pink step, and I had an accident.
         )
 
         async for chunk in response:
+            if cancel_event.is_set():
+                return
             chunk_message = chunk['choices'][0]['delta']
             if 'content' in chunk_message:
                 chunk_text = chunk_message['content']
@@ -167,11 +169,13 @@ I fell off the pink step, and I had an accident.
                     yield text_to_speak
                     current_sentence = current_sentence[len(text_to_speak):]
 
+        if cancel_event.is_set():
+            return
         if len(current_sentence) > 0:
             yield current_sentence
         self._messages.append({"role": "assistant", "content": agent_response})
      
-    async def get_speech_chunks_async(self, text_to_speak):
+    async def get_speech_chunks_async(self, text_to_speak, cancel_event):
         stream = self._speech_service.stream(text_to_speak)
         stream, stream_backup = itertools.tee(stream)
         while True:
@@ -183,6 +187,8 @@ I fell off the pink step, and I had an accident.
 
             # Run next(stream) in a separate thread to avoid blocking the event loop
             chunk = await asyncio.to_thread(next, stream)
+            if cancel_event.is_set():
+                return
             yield chunk
 
     def enqueue_speech_bytes_to_play(self, speech_bytes):