send llm response with eos

llm_service.py

@@ -202,13 +202,16 @@ class MistralTensorRTLLM:
 
             prompt = transcription_output['prompt'].strip()
             input_text=[self.format_prompt_qa(prompt)]
-            self.eos = transcription_output["eos"]
-            
+                    
+            # if prompt is same but EOS is True, we need that to send outputs to websockets
             if self.last_prompt == prompt:
-                if self.last_output is not None:
-                    # logging.info(f"[LLM info:] Same prompt, adding last llm output to audio queue.")
+                if self.last_output is not None and transcription_output["eos"]:
+                    self.eos = transcription_output["eos"]
+                    llm_queue.put({"uid": transcription_output["uid"], "llm_output": self.last_output, "eos": self.eos})
                     audio_queue.put({"llm_output": self.last_output, "eos": self.eos})
                     continue
+            
+            self.eos = transcription_output["eos"]
 
             logging.info(f"[LLM INFO:] WhisperLive prompt: {prompt}, eos: {self.eos}")
             batch_input_ids = self.parse_input(
@@ -278,11 +281,12 @@ class MistralTensorRTLLM:
             if output is not None:
                 self.last_output = output
                 self.last_prompt = prompt
-                llm_queue.put({"uid": transcription_output["uid"], "llm_output": output})
+                llm_queue.put({"uid": transcription_output["uid"], "llm_output": output, "eos": self.eos})
                 audio_queue.put({"llm_output": output, "eos": self.eos})
             
             if self.eos:
                 self.last_prompt = None
+                self.last_output = None
 
 
 if __name__=="__main__":

whisper_live/trt_server.py

@@ -350,16 +350,19 @@ class ServeClient:
 
         """
         while True:
-            if self.eos:
-                try:
-                    llm_output = None
-                    if self.llm_queue is not None:
-                        while not self.llm_queue.empty():
-                            llm_output = self.llm_queue.get_nowait()
-                        if llm_output:
-                            self.websocket.send(json.dumps(llm_output))
-                except queue.Empty:
-                    pass
+            # send the LLM outputs
+            try:
+                llm_response = None
+                if self.llm_queue is not None:
+                    while not self.llm_queue.empty():
+                        llm_response = self.llm_queue.get()
+                    
+                    if llm_response:
+                        # eos = llm_response["eos"]
+                        # if eos:
+                        self.websocket.send(json.dumps(llm_response))
+            except queue.Empty:
+                pass
             
             if self.exit:
                 logging.info("Exiting speech to text thread")
@@ -400,12 +403,12 @@ class ServeClient:
                                     "eos": self.eos
                                 })
                             )
-                            logging.info(f"[INFO]: {segments}, eos: {self.eos}")
-                        
+                            
                         self.transcription_queue.put({"uid": self.client_uid, "prompt": self.prompt, "eos": self.eos})
                         if self.eos:
                             # self.append_segment(last_segment)
                             self.timestamp_offset += duration
+                            logging.info(f"[INFO]: {segments}, eos: {self.eos}")
                             logging.info(
                                 f"[INFO:] Processed : {self.timestamp_offset} seconds / {self.frames_np.shape[0] / self.RATE} seconds"
                             )