diff --git a/src/pipecat/services/riva.py b/src/pipecat/services/riva.py index d5365df47..d6c5f85f5 100644 --- a/src/pipecat/services/riva.py +++ b/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.responses = self.asr_service.streaming_response_generator( - audio_chunks=[self.audio], - streaming_config=self.config, - ) + self._queue = asyncio.Queue() 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]: - # this doesn't work, but something like this perhaps? part 2 - self.audio.append(audio) + async def _stop_tasks(self): + self._thread_task.cancel() + await self._thread_task + self._response_task.cancel() + await self._response_task - # need to start to run this generator only once somewhere... - # 'start' function doesn't work... - # something about the event loop... - # maybe an audio buffer... though my attempt at that didn't work either - for response in self.responses: + def _response_handler(self): + responses = self.asr_service.streaming_response_generator( + audio_chunks=self, + streaming_config=self.config, + ) + for response in responses: if not response.results: continue - partial_transcript = "" - for result in response.results: - if result: - if not result.alternatives: - continue - transcript = result.alternatives[0].transcript - if transcript: - language = None - if len(transcript) > 0: - await self.stop_ttfb_metrics() - if result.is_final: - await self.stop_processing_metrics() - yield TranscriptionFrame( - transcript, "", time_now_iso8601(), language - ) - else: - yield InterimTranscriptionFrame( - transcript, "", time_now_iso8601(), language - ) + asyncio.run_coroutine_threadsafe( + self._response_queue.put(response), self.get_event_loop() + ) + + async def _thread_task_handler(self): + try: + self._thread_running = True + await asyncio.to_thread(self._response_handler) + except asyncio.CancelledError: + self._thread_running = False + pass + + async def _handle_response(self, response): + for result in response.results: + if result and not result.alternatives: + continue + + 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): - await self.start_ttfb_metrics() - await self.start_processing_metrics() + def __next__(self) -> bytes: + if not self._thread_running: + raise StopIteration + future = asyncio.run_coroutine_threadsafe(self._queue.get(), self.get_event_loop()) + return future.result() + + def __iter__(self): + return self