dv-pipecat-ai 0.0.74.dev770__py3-none-any.whl → 0.0.82.dev776__py3-none-any.whl

pipecat/services/rime/tts.py

@@ -4,10 +4,16 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
+"""Rime text-to-speech service implementations.
+
+This module provides both WebSocket and HTTP-based text-to-speech services
+using Rime's API for streaming and batch audio synthesis.
+"""
+
 import base64
 import json
 import uuid
-from typing import AsyncGenerator, Optional
+from typing import Any, AsyncGenerator, Mapping, Optional
 
 import aiohttp
 from loguru import logger
@@ -33,7 +39,8 @@ from pipecat.utils.text.skip_tags_aggregator import SkipTagsAggregator
 from pipecat.utils.tracing.service_decorators import traced_tts
 
 try:
-    import websockets
+    from websockets.asyncio.client import connect as websocket_connect
+    from websockets.protocol import State
 except ModuleNotFoundError as e:
     logger.error(f"Exception: {e}")
     logger.error("In order to use Rime, you need to `pip install pipecat-ai[rime]`.")
@@ -47,7 +54,7 @@ def language_to_rime_language(language: Language) -> str:
         language: The pipecat Language enum value.
 
     Returns:
-        str: Three-letter language code used by Rime (e.g., 'eng' for English).
+        Three-letter language code used by Rime (e.g., 'eng' for English).
     """
     LANGUAGE_MAP = {
         Language.DE: "ger",
@@ -67,7 +74,15 @@ class RimeTTSService(AudioContextWordTTSService):
     """
 
     class InputParams(BaseModel):
-        """Configuration parameters for Rime TTS service."""
+        """Configuration parameters for Rime TTS service.
+
+        Parameters:
+            language: Language for synthesis. Defaults to English.
+            speed_alpha: Speech speed multiplier. Defaults to 1.0.
+            reduce_latency: Whether to reduce latency at potential quality cost.
+            pause_between_brackets: Whether to add pauses between bracketed content.
+            phonemize_between_brackets: Whether to phonemize bracketed content.
+        """
 
         language: Optional[Language] = Language.EN
         speed_alpha: Optional[float] = 1.0
@@ -85,6 +100,7 @@ class RimeTTSService(AudioContextWordTTSService):
         sample_rate: Optional[int] = None,
         params: Optional[InputParams] = None,
         text_aggregator: Optional[BaseTextAggregator] = None,
+        aggregate_sentences: Optional[bool] = True,
         **kwargs,
     ):
         """Initialize Rime TTS service.
@@ -96,10 +112,13 @@ class RimeTTSService(AudioContextWordTTSService):
             model: Model ID to use for synthesis.
             sample_rate: Audio sample rate in Hz.
             params: Additional configuration parameters.
+            text_aggregator: Custom text aggregator for processing input text.
+            aggregate_sentences: Whether to aggregate sentences within the TTSService.
+            **kwargs: Additional arguments passed to parent class.
         """
         # Initialize with parent class settings for proper frame handling
         super().__init__(
-            aggregate_sentences=True,
+            aggregate_sentences=aggregate_sentences,
             push_text_frames=False,
             push_stop_frames=True,
             pause_frame_processing=True,
@@ -135,17 +154,43 @@ class RimeTTSService(AudioContextWordTTSService):
         self._cumulative_time = 0  # Accumulates time across messages
 
     def can_generate_metrics(self) -> bool:
+        """Check if this service can generate processing metrics.
+
+        Returns:
+            True, as Rime service supports metrics generation.
+        """
         return True
 
     def language_to_service_language(self, language: Language) -> str | None:
-        """Convert pipecat language to Rime language code."""
+        """Convert pipecat language to Rime language code.
+
+        Args:
+            language: The language to convert.
+
+        Returns:
+            The Rime-specific language code, or None if not supported.
+        """
         return language_to_rime_language(language)
 
     async def set_model(self, model: str):
-        """Update the TTS model."""
+        """Update the TTS model.
+
+        Args:
+            model: The model name to use for synthesis.
+        """
         self._model = model
         await super().set_model(model)
 
+    async def _update_settings(self, settings: Mapping[str, Any]):
+        """Update service settings and reconnect if voice changed."""
+        prev_voice = self._voice_id
+        await super()._update_settings(settings)
+        if not prev_voice == self._voice_id:
+            self._settings["speaker"] = self._voice_id
+            logger.info(f"Switching TTS voice to: [{self._voice_id}]")
+            await self._disconnect()
+            await self._connect()
+
     def _build_msg(self, text: str = "") -> dict:
         """Build JSON message for Rime API."""
         return {"text": text, "contextId": self._context_id}
@@ -159,18 +204,30 @@ class RimeTTSService(AudioContextWordTTSService):
         return {"operation": "eos"}
 
     async def start(self, frame: StartFrame):
-        """Start the service and establish websocket connection."""
+        """Start the service and establish websocket connection.
+
+        Args:
+            frame: The start frame containing initialization parameters.
+        """
         await super().start(frame)
         self._settings["samplingRate"] = self.sample_rate
         await self._connect()
 
     async def stop(self, frame: EndFrame):
-        """Stop the service and close connection."""
+        """Stop the service and close connection.
+
+        Args:
+            frame: The end frame.
+        """
         await super().stop(frame)
         await self._disconnect()
 
     async def cancel(self, frame: CancelFrame):
-        """Cancel current operation and clean up."""
+        """Cancel current operation and clean up.
+
+        Args:
+            frame: The cancel frame.
+        """
         await super().cancel(frame)
         await self._disconnect()
 
@@ -192,13 +249,13 @@ class RimeTTSService(AudioContextWordTTSService):
     async def _connect_websocket(self):
         """Connect to Rime websocket API with configured settings."""
         try:
-            if self._websocket and self._websocket.open:
+            if self._websocket and self._websocket.state is State.OPEN:
                 return
 
             params = "&".join(f"{k}={v}" for k, v in self._settings.items())
             url = f"{self._url}?{params}"
             headers = {"Authorization": f"Bearer {self._api_key}"}
-            self._websocket = await websockets.connect(url, extra_headers=headers)
+            self._websocket = await websocket_connect(url, additional_headers=headers)
         except Exception as e:
             logger.error(f"{self} initialization error: {e}")
             self._websocket = None
@@ -261,6 +318,7 @@ class RimeTTSService(AudioContextWordTTSService):
         return word_pairs
 
     async def flush_audio(self):
+        """Flush any pending audio synthesis."""
         if not self._context_id or not self._websocket:
             return
 
@@ -310,7 +368,12 @@ class RimeTTSService(AudioContextWordTTSService):
                 self._context_id = None
 
     async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
-        """Push frame and handle end-of-turn conditions."""
+        """Push frame and handle end-of-turn conditions.
+
+        Args:
+            frame: The frame to push.
+            direction: The direction to push the frame.
+        """
         await super().push_frame(frame, direction)
         if isinstance(frame, (TTSStoppedFrame, StartInterruptionFrame)):
             if isinstance(frame, TTSStoppedFrame):
@@ -318,17 +381,17 @@ class RimeTTSService(AudioContextWordTTSService):
 
     @traced_tts
     async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
-        """Generate speech from text.
+        """Generate speech from text using Rime's streaming API.
 
         Args:
             text: The text to convert to speech.
 
         Yields:
-            Frames containing audio data and timing information.
+            Frame: Audio frames containing the synthesized speech.
         """
         logger.debug(f"{self}: Generating TTS [{text}]")
         try:
-            if not self._websocket or self._websocket.closed:
+            if not self._websocket or self._websocket.state is State.CLOSED:
                 await self._connect()
 
             try:
@@ -354,7 +417,24 @@ class RimeTTSService(AudioContextWordTTSService):
 
 
 class RimeHttpTTSService(TTSService):
+    """Rime HTTP-based text-to-speech service.
+
+    Provides text-to-speech synthesis using Rime's HTTP API for batch processing.
+    Suitable for use cases where streaming is not required.
+    """
+
     class InputParams(BaseModel):
+        """Configuration parameters for Rime HTTP TTS service.
+
+        Parameters:
+            language: Language for synthesis. Defaults to English.
+            pause_between_brackets: Whether to add pauses between bracketed content.
+            phonemize_between_brackets: Whether to phonemize bracketed content.
+            inline_speed_alpha: Inline speed control markup.
+            speed_alpha: Speech speed multiplier. Defaults to 1.0.
+            reduce_latency: Whether to reduce latency at potential quality cost.
+        """
+
         language: Optional[Language] = Language.EN
         pause_between_brackets: Optional[bool] = False
         phonemize_between_brackets: Optional[bool] = False
@@ -373,6 +453,17 @@ class RimeHttpTTSService(TTSService):
         params: Optional[InputParams] = None,
         **kwargs,
     ):
+        """Initialize Rime HTTP TTS service.
+
+        Args:
+            api_key: Rime API key for authentication.
+            voice_id: ID of the voice to use.
+            aiohttp_session: Shared aiohttp session for HTTP requests.
+            model: Model ID to use for synthesis.
+            sample_rate: Audio sample rate in Hz.
+            params: Additional configuration parameters.
+            **kwargs: Additional arguments passed to parent TTSService.
+        """
         super().__init__(sample_rate=sample_rate, **kwargs)
 
         params = params or RimeHttpTTSService.InputParams()
@@ -396,14 +487,34 @@ class RimeHttpTTSService(TTSService):
             self._settings["inlineSpeedAlpha"] = params.inline_speed_alpha
 
     def can_generate_metrics(self) -> bool:
+        """Check if this service can generate processing metrics.
+
+        Returns:
+            True, as Rime HTTP service supports metrics generation.
+        """
         return True
 
     def language_to_service_language(self, language: Language) -> str | None:
-        """Convert pipecat language to Rime language code."""
+        """Convert pipecat language to Rime language code.
+
+        Args:
+            language: The language to convert.
+
+        Returns:
+            The Rime-specific language code, or None if not supported.
+        """
         return language_to_rime_language(language)
 
     @traced_tts
     async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
+        """Generate speech from text using Rime's HTTP API.
+
+        Args:
+            text: The text to synthesize into speech.
+
+        Yields:
+            Frame: Audio frames containing the synthesized speech.
+        """
         logger.debug(f"{self}: Generating TTS [{text}]")
 
         headers = {

pipecat/services/riva/stt.py

@@ -4,7 +4,10 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
+"""NVIDIA Riva Speech-to-Text service implementations for real-time and batch transcription."""
+
 import asyncio
+from concurrent.futures import CancelledError as FuturesCancelledError
 from typing import AsyncGenerator, List, Mapping, Optional
 
 from loguru import logger
@@ -21,7 +24,6 @@ from pipecat.frames.frames import (
 )
 from pipecat.services.stt_service import SegmentedSTTService, STTService
 from pipecat.transcriptions.language import Language
-from pipecat.utils.asyncio.watchdog_queue import WatchdogQueue
 from pipecat.utils.time import time_now_iso8601
 from pipecat.utils.tracing.service_decorators import traced_stt
 
@@ -87,7 +89,20 @@ def language_to_riva_language(language: Language) -> Optional[str]:
 
 
 class RivaSTTService(STTService):
+    """Real-time speech-to-text service using NVIDIA Riva streaming ASR.
+
+    Provides real-time transcription capabilities using NVIDIA's Riva ASR models
+    through streaming recognition. Supports interim results and continuous audio
+    processing for low-latency applications.
+    """
+
     class InputParams(BaseModel):
+        """Configuration parameters for Riva STT service.
+
+        Parameters:
+            language: Target language for transcription. Defaults to EN_US.
+        """
+
         language: Optional[Language] = Language.EN_US
 
     def __init__(
@@ -103,6 +118,16 @@ class RivaSTTService(STTService):
         params: Optional[InputParams] = None,
         **kwargs,
     ):
+        """Initialize the Riva STT service.
+
+        Args:
+            api_key: NVIDIA API key for authentication.
+            server: Riva server address. Defaults to NVIDIA Cloud Function endpoint.
+            model_function_map: Mapping containing 'function_id' and 'model_name' for the ASR model.
+            sample_rate: Audio sample rate in Hz. If None, uses pipeline default.
+            params: Additional configuration parameters for Riva.
+            **kwargs: Additional arguments passed to STTService.
+        """
         super().__init__(sample_rate=sample_rate, **kwargs)
 
         params = params or RivaSTTService.InputParams()
@@ -142,15 +167,29 @@ class RivaSTTService(STTService):
 
         self._asr_service = riva.client.ASRService(auth)
 
-        self._queue = asyncio.Queue()
+        self._queue = None
         self._config = None
         self._thread_task = None
         self._response_task = None
 
     def can_generate_metrics(self) -> bool:
+        """Check if this service can generate processing metrics.
+
+        Returns:
+            False - this service does not support metrics generation.
+        """
         return False
 
     async def set_model(self, model: str):
+        """Set the ASR model for transcription.
+
+        Args:
+            model: Model name to set.
+
+        Note:
+            Model cannot be changed after initialization. Use model_function_map
+            parameter in constructor instead.
+        """
         logger.warning(f"Cannot set model after initialization. Set model and function id like so:")
         example = {"function_id": "<UUID>", "model_name": "<model_name>"}
         logger.warning(
@@ -158,6 +197,11 @@ class RivaSTTService(STTService):
         )
 
     async def start(self, frame: StartFrame):
+        """Start the Riva STT service and initialize streaming configuration.
+
+        Args:
+            frame: StartFrame indicating pipeline start.
+        """
         await super().start(frame)
 
         if self._config:
@@ -194,19 +238,30 @@ class RivaSTTService(STTService):
         riva.client.add_custom_configuration_to_config(config, self._custom_configuration)
 
         self._config = config
+        self._queue = asyncio.Queue()
 
         if not self._thread_task:
             self._thread_task = self.create_task(self._thread_task_handler())
 
         if not self._response_task:
-            self._response_queue = WatchdogQueue(self.task_manager)
+            self._response_queue = asyncio.Queue()
             self._response_task = self.create_task(self._response_task_handler())
 
     async def stop(self, frame: EndFrame):
+        """Stop the Riva STT service and clean up resources.
+
+        Args:
+            frame: EndFrame indicating pipeline stop.
+        """
         await super().stop(frame)
         await self._stop_tasks()
 
     async def cancel(self, frame: CancelFrame):
+        """Cancel the Riva STT service operation.
+
+        Args:
+            frame: CancelFrame indicating operation cancellation.
+        """
         await super().cancel(frame)
         await self._stop_tasks()
 
@@ -225,7 +280,6 @@ class RivaSTTService(STTService):
             streaming_config=self._config,
         )
         for response in responses:
-            self.reset_watchdog()
             if not response.results:
                 continue
             asyncio.run_coroutine_threadsafe(
@@ -260,7 +314,7 @@ class RivaSTTService(STTService):
                     await self.push_frame(
                         TranscriptionFrame(
                             transcript,
-                            "",
+                            self._user_id,
                             time_now_iso8601(),
                             self._language_code,
                             result=result,
@@ -275,7 +329,7 @@ class RivaSTTService(STTService):
                     await self.push_frame(
                         InterimTranscriptionFrame(
                             transcript,
-                            "",
+                            self._user_id,
                             time_now_iso8601(),
                             self._language_code,
                             result=result,
@@ -289,18 +343,43 @@ class RivaSTTService(STTService):
             self._response_queue.task_done()
 
     async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
+        """Process audio data for speech-to-text transcription.
+
+        Args:
+            audio: Raw audio bytes to transcribe.
+
+        Yields:
+            None - transcription results are pushed to the pipeline via frames.
+        """
         await self.start_ttfb_metrics()
         await self.start_processing_metrics()
         await self._queue.put(audio)
         yield None
 
     def __next__(self) -> bytes:
+        """Get the next audio chunk for Riva processing.
+
+        Returns:
+            Audio bytes from the queue.
+
+        Raises:
+            StopIteration: When the thread is no longer running.
+        """
         if not self._thread_running:
             raise StopIteration
-        future = asyncio.run_coroutine_threadsafe(self._queue.get(), self.get_event_loop())
-        return future.result()
+
+        try:
+            future = asyncio.run_coroutine_threadsafe(self._queue.get(), self.get_event_loop())
+            return future.result()
+        except FuturesCancelledError:
+            raise StopIteration
 
     def __iter__(self):
+        """Return iterator for audio chunk processing.
+
+        Returns:
+            Self as iterator.
+        """
         return self
 
 
@@ -310,17 +389,20 @@ class RivaSegmentedSTTService(SegmentedSTTService):
     By default, his service uses NVIDIA's Riva Canary ASR API to perform speech-to-text
     transcription on audio segments. It inherits from SegmentedSTTService to handle
     audio buffering and speech detection.
-
-    Args:
-        api_key: NVIDIA API key for authentication
-        server: Riva server address (defaults to NVIDIA Cloud Function endpoint)
-        model_function_map: Mapping of model name and its corresponding NVIDIA Cloud Function ID
-        sample_rate: Audio sample rate in Hz. If not provided, uses the pipeline's rate
-        params: Additional configuration parameters for Riva
-        **kwargs: Additional arguments passed to SegmentedSTTService
     """
 
     class InputParams(BaseModel):
+        """Configuration parameters for Riva segmented STT service.
+
+        Parameters:
+            language: Target language for transcription. Defaults to EN_US.
+            profanity_filter: Whether to filter profanity from results.
+            automatic_punctuation: Whether to add automatic punctuation.
+            verbatim_transcripts: Whether to return verbatim transcripts.
+            boosted_lm_words: List of words to boost in language model.
+            boosted_lm_score: Score boost for specified words.
+        """
+
         language: Optional[Language] = Language.EN_US
         profanity_filter: bool = False
         automatic_punctuation: bool = True
@@ -341,6 +423,16 @@ class RivaSegmentedSTTService(SegmentedSTTService):
         params: Optional[InputParams] = None,
         **kwargs,
     ):
+        """Initialize the Riva segmented STT service.
+
+        Args:
+            api_key: NVIDIA API key for authentication
+            server: Riva server address (defaults to NVIDIA Cloud Function endpoint)
+            model_function_map: Mapping of model name and its corresponding NVIDIA Cloud Function ID
+            sample_rate: Audio sample rate in Hz. If not provided, uses the pipeline's rate
+            params: Additional configuration parameters for Riva
+            **kwargs: Additional arguments passed to SegmentedSTTService
+        """
         super().__init__(sample_rate=sample_rate, **kwargs)
 
         params = params or RivaSegmentedSTTService.InputParams()
@@ -380,7 +472,14 @@ class RivaSegmentedSTTService(SegmentedSTTService):
         self._settings = {"language": self._language_enum}
 
     def language_to_service_language(self, language: Language) -> Optional[str]:
-        """Convert pipecat Language enum to Riva's language code."""
+        """Convert pipecat Language enum to Riva's language code.
+
+        Args:
+            language: Language enum value.
+
+        Returns:
+            Riva language code or None if not supported.
+        """
         return language_to_riva_language(language)
 
     def _initialize_client(self):
@@ -435,10 +534,23 @@ class RivaSegmentedSTTService(SegmentedSTTService):
         return config
 
     def can_generate_metrics(self) -> bool:
-        """Indicates whether this service can generate processing metrics."""
+        """Check if this service can generate processing metrics.
+
+        Returns:
+            True - this service supports metrics generation.
+        """
         return True
 
     async def set_model(self, model: str):
+        """Set the ASR model for transcription.
+
+        Args:
+            model: Model name to set.
+
+        Note:
+            Model cannot be changed after initialization. Use model_function_map
+            parameter in constructor instead.
+        """
         logger.warning(f"Cannot set model after initialization. Set model and function id like so:")
         example = {"function_id": "<UUID>", "model_name": "<model_name>"}
         logger.warning(
@@ -446,13 +558,21 @@ class RivaSegmentedSTTService(SegmentedSTTService):
         )
 
     async def start(self, frame: StartFrame):
-        """Initialize the service when the pipeline starts."""
+        """Initialize the service when the pipeline starts.
+
+        Args:
+            frame: StartFrame indicating pipeline start.
+        """
         await super().start(frame)
         self._initialize_client()
         self._config = self._create_recognition_config()
 
     async def set_language(self, language: Language):
-        """Set the language for the STT service."""
+        """Set the language for the STT service.
+
+        Args:
+            language: Target language for transcription.
+        """
         logger.info(f"Switching STT language to: [{language}]")
         self._language_enum = language
         self._language = self.language_to_service_language(language) or "en-US"
@@ -520,7 +640,10 @@ class RivaSegmentedSTTService(SegmentedSTTService):
                         if text:
                             logger.debug(f"Transcription: [{text}]")
                             yield TranscriptionFrame(
-                                text, "", time_now_iso8601(), self._language_enum
+                                text,
+                                self._user_id,
+                                time_now_iso8601(),
+                                self._language_enum,
                             )
                             transcription_found = True
 
@@ -539,7 +662,12 @@ class RivaSegmentedSTTService(SegmentedSTTService):
 
 
 class ParakeetSTTService(RivaSTTService):
-    """Deprecated: Use RivaSTTService instead."""
+    """Deprecated speech-to-text service using NVIDIA Parakeet models.
+
+    .. deprecated:: 0.0.66
+        This class is deprecated. Use `RivaSTTService` instead for equivalent functionality
+        with Parakeet models by specifying the appropriate model_function_map.
+    """
 
     def __init__(
         self,
@@ -554,6 +682,16 @@ class ParakeetSTTService(RivaSTTService):
         params: Optional[RivaSTTService.InputParams] = None,  # Use parent class's type
         **kwargs,
     ):
+        """Initialize the Parakeet STT service.
+
+        Args:
+            api_key: NVIDIA API key for authentication.
+            server: Riva server address. Defaults to NVIDIA Cloud Function endpoint.
+            model_function_map: Mapping containing 'function_id' and 'model_name' for Parakeet model.
+            sample_rate: Audio sample rate in Hz. If None, uses pipeline default.
+            params: Additional configuration parameters for Riva.
+            **kwargs: Additional arguments passed to RivaSTTService.
+        """
         super().__init__(
             api_key=api_key,
             server=server,