Allowing to define SmartTurnParams

src/pipecat/audio/turn/base_smart_turn.py

@@ -12,6 +12,7 @@ from typing import Dict, Optional
 
 import numpy as np
 from loguru import logger
+from pydantic import BaseModel
 
 
 class EndOfTurnState(Enum):
@@ -19,18 +20,25 @@ class EndOfTurnState(Enum):
     INCOMPLETE = 2
 
 
-# TODO: we should convert all this to params
+STOP_SECS = 1
-STOP_MS = 1000
 PRE_SPEECH_MS = 200
 MAX_DURATION_SECONDS = 8  # Maximum duration for the smart turn model
 
 
-class BaseEndOfTurnAnalyzer(ABC):
+class SmartTurnParams(BaseModel):
-    def __init__(self, *, sample_rate: Optional[int] = None):
+    stop_secs: float = STOP_SECS
+    pre_speech_ms: float = PRE_SPEECH_MS
+    max_duration_secs: float = MAX_DURATION_SECONDS
+
+
+class BaseSmartTurn(ABC):
+    def __init__(self, *, sample_rate: Optional[int] = None, params: SmartTurnParams = SmartTurnParams()):
         self._init_sample_rate = sample_rate
+        self._params = params
         # settings variables
         self._sample_rate = 0
         self._chunk_size_ms = 0
+        self._stop_ms = self._params.stop_secs * 1000
         # inference variables
         self._audio_buffer = []
         self._speech_triggered = False
@@ -67,8 +75,8 @@ class BaseEndOfTurnAnalyzer(ABC):
             if self._speech_triggered:
                 self._audio_buffer.append((time.time(), audio_float32))
                 self._silence_frames += 1
-                if self._silence_frames * self._chunk_size_ms >= STOP_MS:
+                if self._silence_frames * self._chunk_size_ms >= self._stop_ms:
-                    logger.debug("End of Turn complete due to STOP_MS.")
+                    logger.debug("End of Turn complete due to stop_secs.")
                     state = EndOfTurnState.COMPLETE
                     self._clear()
             else:
@@ -76,8 +84,8 @@ class BaseEndOfTurnAnalyzer(ABC):
                 self._audio_buffer.append((time.time(), audio_float32))
                 # Keep the buffer size reasonable, assuming CHUNK is small
                 max_buffer_time = (
-                    PRE_SPEECH_MS + STOP_MS
+                    self._params.pre_speech_ms + self._stop_ms
-                ) / 1000 + MAX_DURATION_SECONDS  # Some extra buffer
+                ) / 1000 + self._params.max_duration_secs  # Some extra buffer
                 while (
                     self._audio_buffer and self._audio_buffer[0][0] < time.time() - max_buffer_time
                 ):
@@ -107,7 +115,7 @@ class BaseEndOfTurnAnalyzer(ABC):
             return state
 
         # Find start and end indices for the segment
-        start_time = self._speech_start_time - (PRE_SPEECH_MS / 1000)
+        start_time = self._speech_start_time - (self._params.pre_speech_ms / 1000)
         start_index = 0
         for i, (t, _) in enumerate(audio_buffer):
             if t >= start_time:
@@ -120,13 +128,13 @@ class BaseEndOfTurnAnalyzer(ABC):
         segment_audio_chunks = [chunk for _, chunk in audio_buffer[start_index : end_index + 1]]
         segment_audio = np.concatenate(segment_audio_chunks)
 
-        # Remove (STOP_MS - 200)ms from the end of the segment
+        # Remove (self._stop_ms - 200)ms from the end of the segment
-        samples_to_remove = int((STOP_MS - 200) / 1000 * self.sample_rate)
+        samples_to_remove = int((self._stop_ms - 200) / 1000 * self.sample_rate)
         segment_audio = segment_audio[:-samples_to_remove]
 
         # Limit maximum duration
-        if len(segment_audio) / self.sample_rate > MAX_DURATION_SECONDS:
+        if len(segment_audio) / self.sample_rate > self._params.max_duration_secs:
-            segment_audio = segment_audio[: int(MAX_DURATION_SECONDS * self.sample_rate)]
+            segment_audio = segment_audio[: int(self._params.max_duration_secs * self.sample_rate)]
 
         # No resampling needed as both recording and prediction use 16000 Hz
         segment_audio_resampled = segment_audio

src/pipecat/audio/turn/local_smart_turn.py

@@ -12,9 +12,7 @@ import numpy as np
 import torch
 from loguru import logger
 
-from pipecat.audio.turn.base_turn_analyzer import (
+from pipecat.audio.turn.base_smart_turn import BaseSmartTurn
-    BaseEndOfTurnAnalyzer,
-)
 
 try:
     import coremltools as ct
@@ -27,9 +25,9 @@ except ModuleNotFoundError as e:
     raise Exception(f"Missing module: {e}")
 
 
-class LocalSmartTurnAnalyzer(BaseEndOfTurnAnalyzer):
+class LocalSmartTurnAnalyzer(BaseSmartTurn):
-    def __init__(self):
+    def __init__(self, **kwargs):
-        super().__init__()
+        super().__init__(**kwargs)
         # To use this locally, set the environment variable LOCAL_SMART_TURN_MODEL_PATH
         # to the path where the smart-turn repo is cloned.
         #

src/pipecat/audio/turn/remote_smart_turn.py

@@ -13,12 +13,12 @@ import numpy as np
 import requests
 from loguru import logger
 
-from pipecat.audio.turn.base_turn_analyzer import BaseEndOfTurnAnalyzer
+from pipecat.audio.turn.base_smart_turn import BaseSmartTurn
 
 
-class RemoteSmartTurnAnalyzer(BaseEndOfTurnAnalyzer):
+class RemoteSmartTurnAnalyzer(BaseSmartTurn):
-    def __init__(self):
+    def __init__(self, **kwargs):
-        super().__init__()
+        super().__init__(**kwargs)
         self.remote_smart_turn_url = os.getenv("REMOTE_SMART_TURN_URL")
 
         if not self.remote_smart_turn_url:

src/pipecat/transports/base_input.py

@@ -10,7 +10,7 @@ from typing import Optional
 
 from loguru import logger
 
-from pipecat.audio.turn.base_turn_analyzer import BaseEndOfTurnAnalyzer, EndOfTurnState
+from pipecat.audio.turn.base_smart_turn import BaseSmartTurn, EndOfTurnState
 from pipecat.audio.vad.vad_analyzer import VADAnalyzer, VADState
 from pipecat.frames.frames import (
     BotInterruptionFrame,
@@ -67,7 +67,7 @@ class BaseInputTransport(FrameProcessor):
         return self._params.vad_analyzer
 
     @property
-    def end_of_turn_analyzer(self) -> Optional[BaseEndOfTurnAnalyzer]:
+    def end_of_turn_analyzer(self) -> Optional[BaseSmartTurn]:
         return self._params.end_of_turn_analyzer
 
     async def start(self, frame: StartFrame):

src/pipecat/transports/base_transport.py

@@ -11,7 +11,7 @@ from pydantic import BaseModel, ConfigDict
 
 from pipecat.audio.filters.base_audio_filter import BaseAudioFilter
 from pipecat.audio.mixers.base_audio_mixer import BaseAudioMixer
-from pipecat.audio.turn.base_turn_analyzer import BaseEndOfTurnAnalyzer
+from pipecat.audio.turn.base_smart_turn import BaseSmartTurn
 from pipecat.audio.vad.vad_analyzer import VADAnalyzer
 from pipecat.processors.frame_processor import FrameProcessor
 from pipecat.utils.base_object import BaseObject
@@ -42,7 +42,7 @@ class TransportParams(BaseModel):
     vad_enabled: bool = False
     vad_audio_passthrough: bool = False
     vad_analyzer: Optional[VADAnalyzer] = None
-    end_of_turn_analyzer: Optional[BaseEndOfTurnAnalyzer] = None
+    end_of_turn_analyzer: Optional[BaseSmartTurn] = None
 
 
 class BaseTransport(BaseObject):