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

pipecat/services/ultravox/stt.py

@@ -44,13 +44,12 @@ except ModuleNotFoundError as e:
 class AudioBuffer:
     """Buffer to collect audio frames before processing.
 
-    Attributes:
-        frames: List of AudioRawFrames to process
-        started_at: Timestamp when speech started
-        is_processing: Flag to prevent concurrent processing
+    Manages the collection and state of audio frames during speech
+    recording sessions, including timing and processing flags.
     """
 
     def __init__(self):
+        """Initialize the audio buffer."""
         self.frames: List[AudioRawFrame] = []
         self.started_at: Optional[float] = None
         self.is_processing: bool = False
@@ -59,19 +58,17 @@ class AudioBuffer:
 class UltravoxModel:
     """Model wrapper for the Ultravox multimodal model.
 
-    This class handles loading and running the Ultravox model for speech-to-text.
-
-    Args:
-        model_name: The name or path of the Ultravox model to load
-
-    Attributes:
-        model_name: The name of the loaded model
-        engine: The vLLM engine for model inference
-        tokenizer: The tokenizer for the model
-        stop_token_ids: Optional token IDs to stop generation
+    This class handles loading and running the Ultravox model for speech-to-text
+    transcription using vLLM for efficient inference.
     """
 
     def __init__(self, model_name: str = "fixie-ai/ultravox-v0_5-llama-3_1-8b"):
+        """Initialize the Ultravox model.
+
+        Args:
+            model_name: The name or path of the Ultravox model to load.
+                Defaults to "fixie-ai/ultravox-v0_5-llama-3_1-8b".
+        """
         self.model_name = model_name
         self._initialize_engine()
         self._initialize_tokenizer()
@@ -95,10 +92,10 @@ class UltravoxModel:
         """Format chat messages into a prompt for the model.
 
         Args:
-            messages: List of message dictionaries with 'role' and 'content'
+            messages: List of message dictionaries with 'role' and 'content'.
 
         Returns:
-            str: Formatted prompt string
+            str: Formatted prompt string ready for model input.
         """
         return self.tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
@@ -114,13 +111,13 @@ class UltravoxModel:
         """Generate text from audio input using the model.
 
         Args:
-            messages: List of message dictionaries
-            temperature: Sampling temperature
-            max_tokens: Maximum tokens to generate
-            audio: Audio data as numpy array
+            messages: List of message dictionaries for conversation context.
+            temperature: Sampling temperature for generation randomness.
+            max_tokens: Maximum number of tokens to generate.
+            audio: Audio data as numpy array in float32 format.
 
         Yields:
-            str: JSON chunks of the generated response
+            str: JSON chunks of the generated response in OpenAI format.
         """
         sampling_params = SamplingParams(
             temperature=temperature, max_tokens=max_tokens, stop_token_ids=self.stop_token_ids
@@ -173,22 +170,9 @@ class UltravoxModel:
 class UltravoxSTTService(AIService):
     """Service to transcribe audio using the Ultravox multimodal model.
 
-    This service collects audio frames and processes them with Ultravox
-    to generate text transcriptions.
-
-    Args:
-        model_name: The Ultravox model to use (ModelSize enum or string)
-        hf_token: Hugging Face token for model access
-        temperature: Sampling temperature for generation
-        max_tokens: Maximum tokens to generate
-        **kwargs: Additional arguments passed to AIService
-
-    Attributes:
-        model: The UltravoxModel instance
-        buffer: Buffer to collect audio frames
-        temperature: Temperature for text generation
-        max_tokens: Maximum tokens to generate
-        _connection_active: Flag indicating if service is active
+    This service collects audio frames during speech and processes them with
+    Ultravox to generate text transcriptions. It handles real-time audio
+    buffering, model warm-up, and streaming text generation.
     """
 
     def __init__(
@@ -200,6 +184,17 @@ class UltravoxSTTService(AIService):
         max_tokens: int = 100,
         **kwargs,
     ):
+        """Initialize the UltravoxSTTService.
+
+        Args:
+            model_name: The Ultravox model to use. Defaults to
+                "fixie-ai/ultravox-v0_5-llama-3_1-8b".
+            hf_token: Hugging Face token for model access. If None, will try
+                to use HF_TOKEN environment variable.
+            temperature: Sampling temperature for generation. Defaults to 0.7.
+            max_tokens: Maximum tokens to generate. Defaults to 100.
+            **kwargs: Additional arguments passed to AIService.
+        """
         super().__init__(**kwargs)
 
         # Authenticate with Hugging Face if token provided
@@ -283,8 +278,11 @@ class UltravoxSTTService(AIService):
     async def start(self, frame: StartFrame):
         """Handle service start.
 
+        Starts the service, marks it as active, and performs model warm-up
+        to ensure optimal performance for the first inference.
+
         Args:
-            frame: StartFrame that triggered this method
+            frame: StartFrame that triggered this method.
         """
         await super().start(frame)
         self._connection_active = True
@@ -296,8 +294,10 @@ class UltravoxSTTService(AIService):
     async def stop(self, frame: EndFrame):
         """Handle service stop.
 
+        Stops the service and marks it as inactive.
+
         Args:
-            frame: EndFrame that triggered this method
+            frame: EndFrame that triggered this method.
         """
         await super().stop(frame)
         self._connection_active = False
@@ -306,8 +306,10 @@ class UltravoxSTTService(AIService):
     async def cancel(self, frame: CancelFrame):
         """Handle service cancellation.
 
+        Cancels the service, clears any buffered audio, and marks it as inactive.
+
         Args:
-            frame: CancelFrame that triggered this method
+            frame: CancelFrame that triggered this method.
         """
         await super().cancel(frame)
         self._connection_active = False
@@ -317,11 +319,12 @@ class UltravoxSTTService(AIService):
     async def process_frame(self, frame: Frame, direction: FrameDirection):
         """Process incoming frames.
 
-        This method collects audio frames and processes them when speech ends.
+        This method collects audio frames during speech and processes them
+        when speech ends to generate text transcriptions.
 
         Args:
-            frame: The frame to process
-            direction: Direction of the frame (input/output)
+            frame: The frame to process.
+            direction: Direction of the frame (input/output).
         """
         await super().process_frame(frame, direction)
 

pipecat/services/vision_service.py

@@ -25,12 +25,14 @@ class VisionService(AIService):
     Provides common functionality for vision services that process images and
     generate textual responses. Handles image frame processing and integrates
     with the AI service infrastructure for metrics and lifecycle management.
-
-    Args:
-        **kwargs: Additional arguments passed to the parent AIService.
     """
 
     def __init__(self, **kwargs):
+        """Initialize the vision service.
+
+        Args:
+            **kwargs: Additional arguments passed to the parent AIService.
+        """
         super().__init__(**kwargs)
         self._describe_text = None
 

pipecat/services/websocket_service.py

@@ -12,6 +12,7 @@ from typing import Awaitable, Callable, Optional
 
 import websockets
 from loguru import logger
+from websockets.exceptions import ConnectionClosedOK
 from websockets.protocol import State
 
 from pipecat.frames.frames import ErrorFrame
@@ -24,13 +25,15 @@ class WebsocketService(ABC):
     Provides websocket connection management, automatic reconnection with
     exponential backoff, connection verification, and error handling.
     Subclasses implement service-specific connection and message handling logic.
-
-    Args:
-        reconnect_on_error: Whether to automatically reconnect on connection errors.
-        **kwargs: Additional arguments (unused, for compatibility).
     """
 
     def __init__(self, *, reconnect_on_error: bool = True, **kwargs):
+        """Initialize the websocket service.
+
+        Args:
+            reconnect_on_error: Whether to automatically reconnect on connection errors.
+            **kwargs: Additional arguments (unused, for compatibility).
+        """
         self._websocket: Optional[websockets.WebSocketClientProtocol] = None
         self._reconnect_on_error = reconnect_on_error
 
@@ -41,7 +44,7 @@ class WebsocketService(ABC):
             True if connection is verified working, False otherwise.
         """
         try:
-            if not self._websocket or self._websocket.closed:
+            if not self._websocket or self._websocket.state is State.CLOSED:
                 return False
             await self._websocket.ping()
             return True
@@ -80,12 +83,10 @@ class WebsocketService(ABC):
             try:
                 await self._receive_messages()
                 retry_count = 0  # Reset counter on successful message receive
-                if self._websocket and self._websocket.state == State.CLOSED:
-                    raise websockets.ConnectionClosedOK(
-                        self._websocket.close_rcvd,
-                        self._websocket.close_sent,
-                        self._websocket.close_rcvd_then_sent,
-                    )
+            except ConnectionClosedOK as e:
+                # Normal closure, don't retry
+                logger.debug(f"{self} connection closed normally: {e}")
+                break
             except Exception as e:
                 message = f"{self} error receiving messages: {e}"
                 logger.error(message)

pipecat/services/whisper/base_stt.py

@@ -4,6 +4,12 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
+"""Base class for Whisper-based speech-to-text services.
+
+This module provides common functionality for services implementing the Whisper API
+interface, including language mapping, metrics generation, and error handling.
+"""
+
 from typing import AsyncGenerator, Optional
 
 from loguru import logger
@@ -18,9 +24,16 @@ from pipecat.utils.tracing.service_decorators import traced_stt
 
 
 def language_to_whisper_language(language: Language) -> Optional[str]:
-    """Language support for Whisper API.
+    """Maps pipecat Language enum to Whisper API language codes.
 
+    Language support for Whisper API.
     Docs: https://platform.openai.com/docs/guides/speech-to-text#supported-languages
+
+    Args:
+        language: A Language enum value representing the input language.
+
+    Returns:
+        str or None: The corresponding Whisper language code, or None if not supported.
     """
     BASE_LANGUAGES = {
         Language.AF: "af",
@@ -98,15 +111,6 @@ class BaseWhisperSTTService(SegmentedSTTService):
 
     Provides common functionality for services implementing the Whisper API interface,
     including metrics generation and error handling.
-
-    Args:
-        model: Name of the Whisper model to use.
-        api_key: Service API key. Defaults to None.
-        base_url: Service API base URL. Defaults to None.
-        language: Language of the audio input. Defaults to English.
-        prompt: Optional text to guide the model's style or continue a previous segment.
-        temperature: Sampling temperature between 0 and 1. Defaults to 0.0.
-        **kwargs: Additional arguments passed to SegmentedSTTService.
     """
 
     def __init__(
@@ -120,6 +124,17 @@ class BaseWhisperSTTService(SegmentedSTTService):
         temperature: Optional[float] = None,
         **kwargs,
     ):
+        """Initialize the Whisper STT service.
+
+        Args:
+            model: Name of the Whisper model to use.
+            api_key: Service API key. Defaults to None.
+            base_url: Service API base URL. Defaults to None.
+            language: Language of the audio input. Defaults to English.
+            prompt: Optional text to guide the model's style or continue a previous segment.
+            temperature: Sampling temperature between 0 and 1. Defaults to 0.0.
+            **kwargs: Additional arguments passed to SegmentedSTTService.
+        """
         super().__init__(**kwargs)
         self.set_model_name(model)
         self._client = self._create_client(api_key, base_url)
@@ -138,12 +153,30 @@ class BaseWhisperSTTService(SegmentedSTTService):
         return AsyncOpenAI(api_key=api_key, base_url=base_url)
 
     async def set_model(self, model: str):
+        """Set the model name for transcription.
+
+        Args:
+            model: The name of the model to use.
+        """
         self.set_model_name(model)
 
     def can_generate_metrics(self) -> bool:
+        """Indicates whether this service can generate metrics.
+
+        Returns:
+            bool: True, as this service supports metric generation.
+        """
         return True
 
     def language_to_service_language(self, language: Language) -> Optional[str]:
+        """Convert from pipecat Language to service language code.
+
+        Args:
+            language: The Language enum value to convert.
+
+        Returns:
+            str or None: The corresponding service language code, or None if not supported.
+        """
         return language_to_whisper_language(language)
 
     async def set_language(self, language: Language):
@@ -153,7 +186,7 @@ class BaseWhisperSTTService(SegmentedSTTService):
             language: The Language enum value to use for transcription.
         """
         logger.info(f"Switching STT language to: [{language}]")
-        self._language = language
+        self._language = self.language_to_service_language(language)
 
     @traced_stt
     async def _handle_transcription(
@@ -163,6 +196,15 @@ class BaseWhisperSTTService(SegmentedSTTService):
         pass
 
     async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
+        """Transcribe audio data to text.
+
+        Args:
+            audio: Raw audio data to transcribe.
+
+        Yields:
+            Frame: Either a TranscriptionFrame containing the transcribed text
+                  or an ErrorFrame if transcription fails.
+        """
         try:
             await self.start_processing_metrics()
             await self.start_ttfb_metrics()
@@ -177,7 +219,11 @@ class BaseWhisperSTTService(SegmentedSTTService):
             if text:
                 await self._handle_transcription(text, True, self._language)
                 logger.debug(f"Transcription: [{text}]")
-                yield TranscriptionFrame(text, "", time_now_iso8601())
+                yield TranscriptionFrame(
+                    text,
+                    self._user_id,
+                    time_now_iso8601(),
+                )
             else:
                 logger.warning("Received empty transcription from API")
 

pipecat/services/whisper/stt.py

@@ -4,7 +4,11 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
-"""This module implements Whisper transcription with a locally-downloaded model."""
+"""Whisper speech-to-text services with locally-downloaded models.
+
+This module implements Whisper transcription using locally-downloaded models,
+supporting both Faster Whisper and MLX Whisper backends for efficient inference.
+"""
 
 import asyncio
 from enum import Enum
@@ -37,25 +41,29 @@ if TYPE_CHECKING:
 
 
 class Model(Enum):
-    """Class of basic Whisper model selection options.
-
-    Available models:
-        Multilingual models:
-            TINY: Smallest multilingual model
-            BASE: Basic multilingual model
-            MEDIUM: Good balance for multilingual
-            LARGE: Best quality multilingual
-            DISTIL_LARGE_V2: Fast multilingual
-
-        English-only models:
-            DISTIL_MEDIUM_EN: Fast English-only
+    """Whisper model selection options for Faster Whisper.
+
+    Provides various model sizes and specializations for speech recognition,
+    balancing quality and performance based on use case requirements.
+
+    Parameters:
+        TINY: Smallest multilingual model, fastest inference.
+        BASE: Basic multilingual model, good speed/quality balance.
+        SMALL: Small multilingual model, better speed/quality balance than BASE.
+        MEDIUM: Medium-sized multilingual model, better quality.
+        LARGE: Best quality multilingual model, slower inference.
+        LARGE_V3_TURBO: Fast multilingual model, slightly lower quality than LARGE.
+        DISTIL_LARGE_V2: Fast multilingual distilled model.
+        DISTIL_MEDIUM_EN: Fast English-only distilled model.
     """
 
     # Multilingual models
     TINY = "tiny"
     BASE = "base"
+    SMALL = "small"
     MEDIUM = "medium"
     LARGE = "large-v3"
+    LARGE_V3_TURBO = "deepdml/faster-whisper-large-v3-turbo-ct2"
     DISTIL_LARGE_V2 = "Systran/faster-distil-whisper-large-v2"
 
     # English-only models
@@ -63,16 +71,18 @@ class Model(Enum):
 
 
 class MLXModel(Enum):
-    """Class of MLX Whisper model selection options.
-
-    Available models:
-        Multilingual models:
-            TINY: Smallest multilingual model
-            MEDIUM: Good balance for multilingual
-            LARGE_V3: Best quality multilingual
-            LARGE_V3_TURBO: Finetuned, pruned Whisper large-v3, much faster, slightly lower quality
-            DISTIL_LARGE_V3: Fast multilingual
-            LARGE_V3_TURBO_Q4: LARGE_V3_TURBO, quantized to Q4
+    """MLX Whisper model selection options for Apple Silicon.
+
+    Provides various model sizes optimized for Apple Silicon hardware,
+    including quantized variants for improved performance.
+
+    Parameters:
+        TINY: Smallest multilingual model for MLX.
+        MEDIUM: Medium-sized multilingual model for MLX.
+        LARGE_V3: Best quality multilingual model for MLX.
+        LARGE_V3_TURBO: Finetuned, pruned Whisper large-v3, much faster with slightly lower quality.
+        DISTIL_LARGE_V3: Fast multilingual distilled model for MLX.
+        LARGE_V3_TURBO_Q4: LARGE_V3_TURBO quantized to Q4 for reduced memory usage.
     """
 
     # Multilingual models
@@ -256,21 +266,6 @@ class WhisperSTTService(SegmentedSTTService):
 
     This service uses Faster Whisper to perform speech-to-text transcription on audio
     segments. It supports multiple languages and various model sizes.
-
-    Args:
-        model: The Whisper model to use for transcription. Can be a Model enum or string.
-        device: The device to run inference on ('cpu', 'cuda', or 'auto').
-        compute_type: The compute type for inference ('default', 'int8', 'int8_float16', etc.).
-        no_speech_prob: Probability threshold for filtering out non-speech segments.
-        language: The default language for transcription.
-        **kwargs: Additional arguments passed to SegmentedSTTService.
-
-    Attributes:
-        _device: The device used for inference.
-        _compute_type: The compute type for inference.
-        _no_speech_prob: Threshold for non-speech filtering.
-        _model: The loaded Whisper model instance.
-        _settings: Dictionary containing service settings.
     """
 
     def __init__(
@@ -283,6 +278,16 @@ class WhisperSTTService(SegmentedSTTService):
         language: Language = Language.EN,
         **kwargs,
     ):
+        """Initialize the Whisper STT service.
+
+        Args:
+            model: The Whisper model to use for transcription. Can be a Model enum or string.
+            device: The device to run inference on ('cpu', 'cuda', or 'auto').
+            compute_type: The compute type for inference ('default', 'int8', 'int8_float16', etc.).
+            no_speech_prob: Probability threshold for filtering out non-speech segments.
+            language: The default language for transcription.
+            **kwargs: Additional arguments passed to SegmentedSTTService.
+        """
         super().__init__(**kwargs)
         self._device: str = device
         self._compute_type = compute_type
@@ -355,7 +360,7 @@ class WhisperSTTService(SegmentedSTTService):
         pass
 
     async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
-        """Transcribes given audio using Whisper.
+        """Transcribe audio data using Whisper.
 
         Args:
             audio: Raw audio bytes in 16-bit PCM format.
@@ -394,7 +399,12 @@ class WhisperSTTService(SegmentedSTTService):
         if text:
             await self._handle_transcription(text, True, self._settings["language"])
             logger.debug(f"Transcription: [{text}]")
-            yield TranscriptionFrame(text, "", time_now_iso8601(), self._settings["language"])
+            yield TranscriptionFrame(
+                text,
+                self._user_id,
+                time_now_iso8601(),
+                self._settings["language"],
+            )
 
 
 class WhisperSTTServiceMLX(WhisperSTTService):
@@ -402,18 +412,6 @@ class WhisperSTTServiceMLX(WhisperSTTService):
 
     This service uses MLX Whisper to perform speech-to-text transcription on audio
     segments. It's optimized for Apple Silicon and supports multiple languages and quantizations.
-
-    Args:
-        model: The MLX Whisper model to use for transcription. Can be an MLXModel enum or string.
-        no_speech_prob: Probability threshold for filtering out non-speech segments.
-        language: The default language for transcription.
-        temperature: Temperature for sampling. Can be a float or tuple of floats.
-        **kwargs: Additional arguments passed to SegmentedSTTService.
-
-    Attributes:
-        _no_speech_threshold: Threshold for non-speech filtering.
-        _temperature: Temperature for sampling.
-        _settings: Dictionary containing service settings.
     """
 
     def __init__(
@@ -425,6 +423,15 @@ class WhisperSTTServiceMLX(WhisperSTTService):
         temperature: float = 0.0,
         **kwargs,
     ):
+        """Initialize the MLX Whisper STT service.
+
+        Args:
+            model: The MLX Whisper model to use for transcription. Can be an MLXModel enum or string.
+            no_speech_prob: Probability threshold for filtering out non-speech segments.
+            language: The default language for transcription.
+            temperature: Temperature for sampling. Can be a float or tuple of floats.
+            **kwargs: Additional arguments passed to SegmentedSTTService.
+        """
         # Skip WhisperSTTService.__init__ and call its parent directly
         SegmentedSTTService.__init__(self, **kwargs)
 
@@ -455,7 +462,10 @@ class WhisperSTTServiceMLX(WhisperSTTService):
 
     @override
     async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
-        """Transcribes given audio using MLX Whisper.
+        """Transcribe audio data using MLX Whisper.
+
+        The audio is expected to be 16-bit signed PCM data.
+        MLX Whisper will handle the conversion internally.
 
         Args:
             audio: Raw audio bytes in 16-bit PCM format.
@@ -463,10 +473,6 @@ class WhisperSTTServiceMLX(WhisperSTTService):
         Yields:
             Frame: Either a TranscriptionFrame containing the transcribed text
                   or an ErrorFrame if transcription fails.
-
-        Note:
-            The audio is expected to be 16-bit signed PCM data.
-            MLX Whisper will handle the conversion internally.
         """
         try:
             import mlx_whisper
@@ -503,7 +509,12 @@ class WhisperSTTServiceMLX(WhisperSTTService):
             if text:
                 await self._handle_transcription(text, True, self._settings["language"])
                 logger.debug(f"Transcription: [{text}]")
-                yield TranscriptionFrame(text, "", time_now_iso8601(), self._settings["language"])
+                yield TranscriptionFrame(
+                    text,
+                    self._user_id,
+                    time_now_iso8601(),
+                    self._settings["language"],
+                )
 
         except Exception as e:
             logger.exception(f"MLX Whisper transcription error: {e}")