services(riva): first working version of ParakeetSTTService

src/pipecat/services/riva.py

@@ -5,7 +5,7 @@
 #
 
 import asyncio
-from typing import AsyncGenerator, List, Optional, Union, Iterator
+from typing import AsyncGenerator, Optional
 
 from loguru import logger
 from pydantic.main import BaseModel
@@ -13,7 +13,6 @@ from pydantic.main import BaseModel
 from pipecat.frames.frames import (
     CancelFrame,
     EndFrame,
-    ErrorFrame,
     Frame,
     InterimTranscriptionFrame,
     StartFrame,
@@ -23,7 +22,6 @@ from pipecat.frames.frames import (
     TTSStoppedFrame,
 )
 from pipecat.services.ai_services import STTService, TTSService
-from pipecat.transcriptions.language import Language
 from pipecat.utils.time import time_now_iso8601
 
 try:
@@ -74,12 +72,6 @@ class FastpitchTTSService(TTSService):
 
         self.service = riva.client.SpeechSynthesisService(auth)
 
-    async def stop(self, frame: EndFrame):
-        await super().stop(frame)
-
-    async def cancel(self, frame: CancelFrame):
-        await super().cancel(frame)
-
     async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
         logger.debug(f"Generating TTS: [{text}]")
 
@@ -130,17 +122,12 @@ class ParakeetSTTService(STTService):
 
         self.set_model_name("parakeet-ctc-1.1b-asr")
 
-        input_device = 0
-        list_devices = False
         profanity_filter = False
         automatic_punctuation = False
         no_verbatim_transcripts = False
         language_code = "en-US"
-        model_name = ""
         boosted_lm_words = None
         boosted_lm_score = 4.0
-        speaker_diarization = False
-        diarization_max_speakers = 3
         start_history = -1
         start_threshold = -1.0
         stop_history = -1
@@ -148,10 +135,7 @@ class ParakeetSTTService(STTService):
         stop_history_eou = -1
         stop_threshold_eou = -1.0
         custom_configuration = ""
-        ssl_cert = None
-        use_ssl = True
         sample_rate_hz: int = 16000
-        file_streaming_chunk = 1600
 
         metadata = [
             ["function-id", model],
@@ -189,56 +173,86 @@ class ParakeetSTTService(STTService):
         riva.client.add_custom_configuration_to_config(config, custom_configuration)
 
         # this doesn't work, but something like this perhaps? part 1
-        self.audio = []
+        self._queue = asyncio.Queue()
-        self.responses = self.asr_service.streaming_response_generator(
-            audio_chunks=[self.audio],
-            streaming_config=self.config,
-        )
 
     def can_generate_metrics(self) -> bool:
         return False
 
     async def start(self, frame: StartFrame):
         await super().start(frame)
+        self._thread_task = self.get_event_loop().create_task(self._thread_task_handler())
+        self._response_task = self.get_event_loop().create_task(self._response_task_handler())
+        self._response_queue = asyncio.Queue()
 
     async def stop(self, frame: EndFrame):
         await super().stop(frame)
+        await self._stop_tasks()
 
     async def cancel(self, frame: CancelFrame):
         await super().cancel(frame)
+        await self._stop_tasks()
 
-    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
+    async def _stop_tasks(self):
-        # this doesn't work, but something like this perhaps? part 2
+        self._thread_task.cancel()
-        self.audio.append(audio)
+        await self._thread_task
+        self._response_task.cancel()
+        await self._response_task
 
-        # need to start to run this generator only once somewhere...
+    def _response_handler(self):
-        # 'start' function doesn't work...
+        responses = self.asr_service.streaming_response_generator(
-        # something about the event loop...
+            audio_chunks=self,
-        # maybe an audio buffer... though my attempt at that didn't work either
+            streaming_config=self.config,
-        for response in self.responses:
+        )
+        for response in responses:
             if not response.results:
                 continue
-            partial_transcript = ""
+            asyncio.run_coroutine_threadsafe(
-            for result in response.results:
+                self._response_queue.put(response), self.get_event_loop()
-                if result:
+            )
-                    if not result.alternatives:
+
-                        continue
+    async def _thread_task_handler(self):
-                    transcript = result.alternatives[0].transcript
+        try:
-                    if transcript:
+            self._thread_running = True
-                        language = None
+            await asyncio.to_thread(self._response_handler)
-                        if len(transcript) > 0:
+        except asyncio.CancelledError:
-                            await self.stop_ttfb_metrics()
+            self._thread_running = False
-                            if result.is_final:
+            pass
-                                await self.stop_processing_metrics()
+
-                                yield TranscriptionFrame(
+    async def _handle_response(self, response):
-                                    transcript, "", time_now_iso8601(), language
+        for result in response.results:
-                                )
+            if result and not result.alternatives:
-                            else:
+                continue
-                                yield InterimTranscriptionFrame(
+
-                                    transcript, "", time_now_iso8601(), language
+            transcript = result.alternatives[0].transcript
-                                )
+            if transcript and len(transcript) > 0:
+                await self.stop_ttfb_metrics()
+                if result.is_final:
+                    await self.stop_processing_metrics()
+                    await self.push_frame(
+                        TranscriptionFrame(transcript, "", time_now_iso8601(), None)
+                    )
+                else:
+                    await self.push_frame(
+                        InterimTranscriptionFrame(transcript, "", time_now_iso8601(), None)
+                    )
+
+    async def _response_task_handler(self):
+        while True:
+            try:
+                response = await self._response_queue.get()
+                await self._handle_response(response)
+            except asyncio.CancelledError:
+                break
+
+    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
+        await self._queue.put(audio)
         yield None
 
-    async def _on_speech_started(self, *args, **kwargs):
+    def __next__(self) -> bytes:
-        await self.start_ttfb_metrics()
+        if not self._thread_running:
-        await self.start_processing_metrics()
+            raise StopIteration
+        future = asyncio.run_coroutine_threadsafe(self._queue.get(), self.get_event_loop())
+        return future.result()
+
+    def __iter__(self):
+        return self