Merge pull request #3236 from pipecat-ai/mb/websocket-stt-services
Update websocket STT services to use the WebsocketSTTService base class
This commit is contained in:
Mark Backman
1
changelog/3236.added.md
Normal file
1
changelog/3236.added.md
Normal file
@@ -0,0 +1 @@
|
||||
- Add Gladia session id to logs for `GladiaSTTService`.
|
||||
6
changelog/3236.changed.md
Normal file
6
changelog/3236.changed.md
Normal file
@@ -0,0 +1,6 @@
|
||||
- Updated websocket STT services to use the `WebsocketSTTService` base class. This base class manages the websocket connection and handles reconnects. Updated services:
|
||||
|
||||
- `AssemblyAISTTService`
|
||||
- `AWSTranscribeSTTService`
|
||||
- `GladiaSTTService`
|
||||
- `SonioxSTTService`
|
||||
1
changelog/3236.fixed.md
Normal file
1
changelog/3236.fixed.md
Normal file
@@ -0,0 +1 @@
|
||||
- In `GladiaSTTService`, reset the `_bytes_sent` counter on connecting the websocket. This avoids unnecessary audio buffer trimming.
|
||||
@@ -21,7 +21,6 @@ from pipecat import __version__ as pipecat_version
|
||||
from pipecat.frames.frames import (
|
||||
CancelFrame,
|
||||
EndFrame,
|
||||
ErrorFrame,
|
||||
Frame,
|
||||
InterimTranscriptionFrame,
|
||||
StartFrame,
|
||||
@@ -30,7 +29,7 @@ from pipecat.frames.frames import (
|
||||
UserStoppedSpeakingFrame,
|
||||
)
|
||||
from pipecat.processors.frame_processor import FrameDirection
|
||||
from pipecat.services.stt_service import STTService
|
||||
from pipecat.services.stt_service import WebsocketSTTService
|
||||
from pipecat.transcriptions.language import Language
|
||||
from pipecat.utils.time import time_now_iso8601
|
||||
from pipecat.utils.tracing.service_decorators import traced_stt
|
||||
@@ -44,15 +43,15 @@ from .models import (
|
||||
)
|
||||
|
||||
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 AssemblyAI, you need to `pip install "pipecat-ai[assemblyai]"`.')
|
||||
raise Exception(f"Missing module: {e}")
|
||||
|
||||
|
||||
class AssemblyAISTTService(STTService):
|
||||
class AssemblyAISTTService(WebsocketSTTService):
|
||||
"""AssemblyAI real-time speech-to-text service.
|
||||
|
||||
Provides real-time speech transcription using AssemblyAI's WebSocket API.
|
||||
@@ -80,15 +79,14 @@ class AssemblyAISTTService(STTService):
|
||||
vad_force_turn_endpoint: Whether to force turn endpoint on VAD stop. Defaults to True.
|
||||
**kwargs: Additional arguments passed to parent STTService class.
|
||||
"""
|
||||
super().__init__(sample_rate=connection_params.sample_rate, **kwargs)
|
||||
|
||||
self._api_key = api_key
|
||||
self._language = language
|
||||
self._api_endpoint_base_url = api_endpoint_base_url
|
||||
self._connection_params = connection_params
|
||||
self._vad_force_turn_endpoint = vad_force_turn_endpoint
|
||||
|
||||
super().__init__(sample_rate=self._connection_params.sample_rate, **kwargs)
|
||||
|
||||
self._websocket = None
|
||||
self._termination_event = asyncio.Event()
|
||||
self._received_termination = False
|
||||
self._connected = False
|
||||
@@ -114,7 +112,7 @@ class AssemblyAISTTService(STTService):
|
||||
frame: Start frame to begin processing.
|
||||
"""
|
||||
await super().start(frame)
|
||||
self._chunk_size_bytes = int(self._chunk_size_ms * self._sample_rate * 2 / 1000)
|
||||
self._chunk_size_bytes = int(self._chunk_size_ms * self.sample_rate * 2 / 1000)
|
||||
await self._connect()
|
||||
|
||||
async def stop(self, frame: EndFrame):
|
||||
@@ -146,10 +144,11 @@ class AssemblyAISTTService(STTService):
|
||||
"""
|
||||
self._audio_buffer.extend(audio)
|
||||
|
||||
while len(self._audio_buffer) >= self._chunk_size_bytes:
|
||||
chunk = bytes(self._audio_buffer[: self._chunk_size_bytes])
|
||||
self._audio_buffer = self._audio_buffer[self._chunk_size_bytes :]
|
||||
await self._websocket.send(chunk)
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
while len(self._audio_buffer) >= self._chunk_size_bytes:
|
||||
chunk = bytes(self._audio_buffer[: self._chunk_size_bytes])
|
||||
self._audio_buffer = self._audio_buffer[self._chunk_size_bytes :]
|
||||
await self._websocket.send(chunk)
|
||||
|
||||
yield None
|
||||
|
||||
@@ -164,7 +163,11 @@ class AssemblyAISTTService(STTService):
|
||||
if isinstance(frame, UserStartedSpeakingFrame):
|
||||
await self.start_ttfb_metrics()
|
||||
elif isinstance(frame, UserStoppedSpeakingFrame):
|
||||
if self._vad_force_turn_endpoint:
|
||||
if (
|
||||
self._vad_force_turn_endpoint
|
||||
and self._websocket
|
||||
and self._websocket.state is State.OPEN
|
||||
):
|
||||
await self._websocket.send(json.dumps({"type": "ForceEndpoint"}))
|
||||
await self.start_processing_metrics()
|
||||
|
||||
@@ -191,7 +194,63 @@ class AssemblyAISTTService(STTService):
|
||||
return self._api_endpoint_base_url
|
||||
|
||||
async def _connect(self):
|
||||
"""Connect to the AssemblyAI service.
|
||||
|
||||
Establishes websocket connection and starts receive task.
|
||||
"""
|
||||
await self._connect_websocket()
|
||||
|
||||
if self._websocket and not self._receive_task:
|
||||
self._receive_task = self.create_task(self._receive_task_handler(self._report_error))
|
||||
|
||||
async def _disconnect(self):
|
||||
"""Disconnect from the AssemblyAI service.
|
||||
|
||||
Sends termination message, waits for acknowledgment, and cleans up.
|
||||
"""
|
||||
if not self._connected or not self._websocket:
|
||||
return
|
||||
|
||||
try:
|
||||
self._termination_event.clear()
|
||||
self._received_termination = False
|
||||
|
||||
if self._websocket.state is State.OPEN:
|
||||
# Send any remaining audio
|
||||
if len(self._audio_buffer) > 0:
|
||||
await self._websocket.send(bytes(self._audio_buffer))
|
||||
self._audio_buffer.clear()
|
||||
|
||||
# Send termination message and wait for acknowledgment
|
||||
try:
|
||||
await self._websocket.send(json.dumps({"type": "Terminate"}))
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(self._termination_event.wait(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("Timed out waiting for termination message from server")
|
||||
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
finally:
|
||||
# Clean up tasks and connection
|
||||
if self._receive_task:
|
||||
await self.cancel_task(self._receive_task)
|
||||
self._receive_task = None
|
||||
|
||||
await self._disconnect_websocket()
|
||||
|
||||
async def _connect_websocket(self):
|
||||
"""Establish the websocket connection to AssemblyAI."""
|
||||
try:
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
return
|
||||
|
||||
logger.debug("Connecting to AssemblyAI WebSocket")
|
||||
|
||||
ws_url = self._build_ws_url()
|
||||
headers = {
|
||||
"Authorization": self._api_key,
|
||||
@@ -202,68 +261,50 @@ class AssemblyAISTTService(STTService):
|
||||
additional_headers=headers,
|
||||
)
|
||||
self._connected = True
|
||||
self._receive_task = self.create_task(self._receive_task_handler())
|
||||
|
||||
await self._call_event_handler("on_connected")
|
||||
logger.debug(f"{self} Connected to AssemblyAI WebSocket")
|
||||
except Exception as e:
|
||||
self._connected = False
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
await self.push_error(error_msg=f"Unable to connect to AssemblyAI: {e}", exception=e)
|
||||
raise
|
||||
|
||||
async def _disconnect(self):
|
||||
"""Disconnect from AssemblyAI WebSocket and wait for termination message."""
|
||||
if not self._connected or not self._websocket:
|
||||
return
|
||||
|
||||
async def _disconnect_websocket(self):
|
||||
"""Close the websocket connection to AssemblyAI."""
|
||||
try:
|
||||
self._termination_event.clear()
|
||||
self._received_termination = False
|
||||
|
||||
if len(self._audio_buffer) > 0:
|
||||
await self._websocket.send(bytes(self._audio_buffer))
|
||||
self._audio_buffer.clear()
|
||||
|
||||
try:
|
||||
await self._websocket.send(json.dumps({"type": "Terminate"}))
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(self._termination_event.wait(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("Timed out waiting for termination message from server")
|
||||
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
|
||||
if self._receive_task:
|
||||
await self.cancel_task(self._receive_task)
|
||||
|
||||
await self._websocket.close()
|
||||
|
||||
if self._websocket:
|
||||
logger.debug("Disconnecting from AssemblyAI WebSocket")
|
||||
await self._websocket.close()
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
|
||||
await self.push_error(error_msg=f"Error closing websocket: {e}", exception=e)
|
||||
finally:
|
||||
self._websocket = None
|
||||
self._connected = False
|
||||
self._receive_task = None
|
||||
await self._call_event_handler("on_disconnected")
|
||||
|
||||
async def _receive_task_handler(self):
|
||||
"""Handle incoming WebSocket messages."""
|
||||
try:
|
||||
while self._connected:
|
||||
try:
|
||||
message = await self._websocket.recv()
|
||||
data = json.loads(message)
|
||||
await self._handle_message(data)
|
||||
except websockets.exceptions.ConnectionClosedOK:
|
||||
break
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
break
|
||||
def _get_websocket(self):
|
||||
"""Get the current WebSocket connection.
|
||||
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
Returns:
|
||||
The WebSocket connection.
|
||||
|
||||
Raises:
|
||||
Exception: If WebSocket is not connected.
|
||||
"""
|
||||
if self._websocket:
|
||||
return self._websocket
|
||||
raise Exception("Websocket not connected")
|
||||
|
||||
async def _receive_messages(self):
|
||||
"""Receive and process websocket messages.
|
||||
|
||||
Continuously processes messages from the websocket connection.
|
||||
"""
|
||||
async for message in self._get_websocket():
|
||||
try:
|
||||
data = json.loads(message)
|
||||
await self._handle_message(data)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Received non-JSON message: {message}")
|
||||
|
||||
def _parse_message(self, message: Dict[str, Any]) -> BaseMessage:
|
||||
"""Parse a raw message into the appropriate message type."""
|
||||
|
||||
@@ -29,13 +29,12 @@ from pipecat.frames.frames import (
|
||||
TranscriptionFrame,
|
||||
)
|
||||
from pipecat.services.aws.utils import build_event_message, decode_event, get_presigned_url
|
||||
from pipecat.services.stt_service import STTService
|
||||
from pipecat.services.stt_service import WebsocketSTTService
|
||||
from pipecat.transcriptions.language import Language, resolve_language
|
||||
from pipecat.utils.time import time_now_iso8601
|
||||
from pipecat.utils.tracing.service_decorators import traced_stt
|
||||
|
||||
try:
|
||||
import websockets
|
||||
from websockets.asyncio.client import connect as websocket_connect
|
||||
from websockets.protocol import State
|
||||
except ModuleNotFoundError as e:
|
||||
@@ -44,7 +43,7 @@ except ModuleNotFoundError as e:
|
||||
raise Exception(f"Missing module: {e}")
|
||||
|
||||
|
||||
class AWSTranscribeSTTService(STTService):
|
||||
class AWSTranscribeSTTService(WebsocketSTTService):
|
||||
"""AWS Transcribe Speech-to-Text service using WebSocket streaming.
|
||||
|
||||
Provides real-time speech transcription using AWS Transcribe's streaming API.
|
||||
@@ -99,9 +98,6 @@ class AWSTranscribeSTTService(STTService):
|
||||
"region": region or os.getenv("AWS_REGION", "us-east-1"),
|
||||
}
|
||||
|
||||
self._ws_client = None
|
||||
self._connection_lock = asyncio.Lock()
|
||||
self._connecting = False
|
||||
self._receive_task = None
|
||||
|
||||
def get_service_encoding(self, encoding: str) -> str:
|
||||
@@ -123,29 +119,9 @@ class AWSTranscribeSTTService(STTService):
|
||||
|
||||
Args:
|
||||
frame: Start frame signaling service initialization.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If WebSocket connection cannot be established after retries.
|
||||
"""
|
||||
await super().start(frame)
|
||||
logger.info("Starting AWS Transcribe service...")
|
||||
retry_count = 0
|
||||
max_retries = 3
|
||||
|
||||
while retry_count < max_retries:
|
||||
try:
|
||||
await self._connect()
|
||||
if self._ws_client and self._ws_client.state is State.OPEN:
|
||||
logger.info("Successfully established WebSocket connection")
|
||||
return
|
||||
logger.warning("WebSocket connection not established after connect")
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
retry_count += 1
|
||||
if retry_count < max_retries:
|
||||
await asyncio.sleep(1) # Wait before retrying
|
||||
|
||||
raise RuntimeError("Failed to establish WebSocket connection after multiple attempts")
|
||||
await self._connect()
|
||||
|
||||
async def stop(self, frame: EndFrame):
|
||||
"""Stop the service and disconnect from AWS Transcribe.
|
||||
@@ -174,140 +150,127 @@ class AWSTranscribeSTTService(STTService):
|
||||
Yields:
|
||||
ErrorFrame: If processing fails or connection issues occur.
|
||||
"""
|
||||
try:
|
||||
# Ensure WebSocket is connected
|
||||
if not self._ws_client or self._ws_client.state is State.CLOSED:
|
||||
logger.debug("WebSocket not connected, attempting to reconnect...")
|
||||
try:
|
||||
await self._connect()
|
||||
except Exception as e:
|
||||
yield ErrorFrame(error=f"Unknown error occurred: {e}")
|
||||
return
|
||||
|
||||
# Format the audio data according to AWS event stream format
|
||||
event_message = build_event_message(audio)
|
||||
|
||||
# Send the formatted event message
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
try:
|
||||
await self._ws_client.send(event_message)
|
||||
# Format the audio data according to AWS event stream format
|
||||
event_message = build_event_message(audio)
|
||||
|
||||
# Send the formatted event message
|
||||
await self._websocket.send(event_message)
|
||||
# Start metrics after first chunk sent
|
||||
await self.start_processing_metrics()
|
||||
await self.start_ttfb_metrics()
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
logger.warning(f"Connection closed while sending: {e}")
|
||||
await self._disconnect()
|
||||
# Don't yield error here - we'll retry on next frame
|
||||
except Exception as e:
|
||||
yield ErrorFrame(error=f"Unknown error occurred: {e}")
|
||||
await self._disconnect()
|
||||
yield ErrorFrame(error=f"Error sending audio: {e}")
|
||||
|
||||
except Exception as e:
|
||||
yield ErrorFrame(error=f"Unknown error occurred: {e}")
|
||||
await self._disconnect()
|
||||
yield None
|
||||
|
||||
async def _connect(self):
|
||||
"""Connect to AWS Transcribe with connection state management."""
|
||||
if self._ws_client and self._ws_client.state is State.OPEN and self._receive_task:
|
||||
logger.debug(f"{self} Already connected")
|
||||
return
|
||||
"""Connect to the AWS Transcribe service.
|
||||
|
||||
async with self._connection_lock:
|
||||
if self._connecting:
|
||||
logger.debug(f"{self} Connection already in progress")
|
||||
return
|
||||
Establishes websocket connection and starts receive task.
|
||||
"""
|
||||
await self._connect_websocket()
|
||||
|
||||
try:
|
||||
self._connecting = True
|
||||
logger.debug(f"{self} Starting connection process...")
|
||||
|
||||
if self._ws_client:
|
||||
await self._disconnect()
|
||||
|
||||
language_code = self.language_to_service_language(
|
||||
Language(self._settings["language"])
|
||||
)
|
||||
if not language_code:
|
||||
raise ValueError(f"Unsupported language: {self._settings['language']}")
|
||||
|
||||
# Generate random websocket key
|
||||
websocket_key = "".join(
|
||||
random.choices(
|
||||
string.ascii_uppercase + string.ascii_lowercase + string.digits, k=20
|
||||
)
|
||||
)
|
||||
|
||||
# Add required headers
|
||||
additional_headers = {
|
||||
"Origin": "https://localhost",
|
||||
"Sec-WebSocket-Key": websocket_key,
|
||||
"Sec-WebSocket-Version": "13",
|
||||
"Connection": "keep-alive",
|
||||
}
|
||||
|
||||
# Get presigned URL
|
||||
presigned_url = get_presigned_url(
|
||||
region=self._credentials["region"],
|
||||
credentials={
|
||||
"access_key": self._credentials["aws_access_key_id"],
|
||||
"secret_key": self._credentials["aws_secret_access_key"],
|
||||
"session_token": self._credentials["aws_session_token"],
|
||||
},
|
||||
language_code=language_code,
|
||||
media_encoding=self.get_service_encoding(
|
||||
self._settings["media_encoding"]
|
||||
), # Convert to AWS format
|
||||
sample_rate=self._settings["sample_rate"],
|
||||
number_of_channels=self._settings["number_of_channels"],
|
||||
enable_partial_results_stabilization=True,
|
||||
partial_results_stability="high",
|
||||
show_speaker_label=self._settings["show_speaker_label"],
|
||||
enable_channel_identification=self._settings["enable_channel_identification"],
|
||||
)
|
||||
|
||||
logger.debug(f"{self} Connecting to WebSocket with URL: {presigned_url[:100]}...")
|
||||
|
||||
# Connect with the required headers and settings
|
||||
self._ws_client = await websocket_connect(
|
||||
presigned_url,
|
||||
additional_headers=additional_headers,
|
||||
subprotocols=["mqtt"],
|
||||
ping_interval=None,
|
||||
ping_timeout=None,
|
||||
compression=None,
|
||||
)
|
||||
|
||||
logger.debug(f"{self} WebSocket connected, starting receive task...")
|
||||
|
||||
# Start receive task
|
||||
self._receive_task = self.create_task(self._receive_loop())
|
||||
|
||||
logger.info(f"{self} Successfully connected to AWS Transcribe")
|
||||
|
||||
await self._call_event_handler("on_connected")
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
await self._disconnect()
|
||||
raise
|
||||
|
||||
finally:
|
||||
self._connecting = False
|
||||
if self._websocket and not self._receive_task:
|
||||
self._receive_task = self.create_task(self._receive_task_handler(self._report_error))
|
||||
|
||||
async def _disconnect(self):
|
||||
"""Disconnect from AWS Transcribe."""
|
||||
"""Disconnect from the AWS Transcribe service.
|
||||
|
||||
Sends end-stream message and cleans up.
|
||||
"""
|
||||
if self._receive_task:
|
||||
await self.cancel_task(self._receive_task)
|
||||
self._receive_task = None
|
||||
|
||||
try:
|
||||
if self._ws_client and self._ws_client.state is State.OPEN:
|
||||
# Send end-stream message
|
||||
# Send end-stream message before closing
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
try:
|
||||
end_stream = {"message-type": "event", "event": "end"}
|
||||
await self._ws_client.send(json.dumps(end_stream))
|
||||
await self._ws_client.close()
|
||||
await self._websocket.send(json.dumps(end_stream))
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Error sending end-stream: {e}", exception=e)
|
||||
|
||||
await self._disconnect_websocket()
|
||||
|
||||
async def _connect_websocket(self):
|
||||
"""Establish the websocket connection to AWS Transcribe."""
|
||||
try:
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
return
|
||||
|
||||
logger.debug("Connecting to AWS Transcribe WebSocket")
|
||||
|
||||
language_code = self.language_to_service_language(Language(self._settings["language"]))
|
||||
if not language_code:
|
||||
raise ValueError(f"Unsupported language: {self._settings['language']}")
|
||||
|
||||
# Generate random websocket key
|
||||
websocket_key = "".join(
|
||||
random.choices(
|
||||
string.ascii_uppercase + string.ascii_lowercase + string.digits, k=20
|
||||
)
|
||||
)
|
||||
|
||||
# Add required headers
|
||||
additional_headers = {
|
||||
"Origin": "https://localhost",
|
||||
"Sec-WebSocket-Key": websocket_key,
|
||||
"Sec-WebSocket-Version": "13",
|
||||
"Connection": "keep-alive",
|
||||
}
|
||||
|
||||
# Get presigned URL
|
||||
presigned_url = get_presigned_url(
|
||||
region=self._credentials["region"],
|
||||
credentials={
|
||||
"access_key": self._credentials["aws_access_key_id"],
|
||||
"secret_key": self._credentials["aws_secret_access_key"],
|
||||
"session_token": self._credentials["aws_session_token"],
|
||||
},
|
||||
language_code=language_code,
|
||||
media_encoding=self.get_service_encoding(
|
||||
self._settings["media_encoding"]
|
||||
), # Convert to AWS format
|
||||
sample_rate=self._settings["sample_rate"],
|
||||
number_of_channels=self._settings["number_of_channels"],
|
||||
enable_partial_results_stabilization=True,
|
||||
partial_results_stability="high",
|
||||
show_speaker_label=self._settings["show_speaker_label"],
|
||||
enable_channel_identification=self._settings["enable_channel_identification"],
|
||||
)
|
||||
|
||||
logger.debug(f"{self} Connecting to WebSocket with URL: {presigned_url[:100]}...")
|
||||
|
||||
# Connect with the required headers and settings
|
||||
self._websocket = await websocket_connect(
|
||||
presigned_url,
|
||||
additional_headers=additional_headers,
|
||||
subprotocols=["mqtt"],
|
||||
ping_interval=None,
|
||||
ping_timeout=None,
|
||||
compression=None,
|
||||
)
|
||||
|
||||
await self._call_event_handler("on_connected")
|
||||
logger.info(f"{self} Successfully connected to AWS Transcribe")
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
await self.push_error(
|
||||
error_msg=f"Unable to connect to AWS Transcribe: {e}", exception=e
|
||||
)
|
||||
raise
|
||||
|
||||
async def _disconnect_websocket(self):
|
||||
"""Close the websocket connection to AWS Transcribe."""
|
||||
try:
|
||||
if self._websocket:
|
||||
logger.debug("Disconnecting from AWS Transcribe WebSocket")
|
||||
await self._websocket.close()
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Error closing websocket: {e}", exception=e)
|
||||
finally:
|
||||
self._ws_client = None
|
||||
self._websocket = None
|
||||
await self._call_event_handler("on_disconnected")
|
||||
|
||||
def language_to_service_language(self, language: Language) -> str | None:
|
||||
@@ -471,16 +434,26 @@ class AWSTranscribeSTTService(STTService):
|
||||
):
|
||||
pass
|
||||
|
||||
async def _receive_loop(self):
|
||||
"""Background task to receive and process messages from AWS Transcribe."""
|
||||
while True:
|
||||
if not self._ws_client or self._ws_client.state is State.CLOSED:
|
||||
logger.warning(f"{self} WebSocket closed in receive loop")
|
||||
break
|
||||
def _get_websocket(self):
|
||||
"""Get the current WebSocket connection.
|
||||
|
||||
Returns:
|
||||
The WebSocket connection.
|
||||
|
||||
Raises:
|
||||
Exception: If WebSocket is not connected.
|
||||
"""
|
||||
if self._websocket:
|
||||
return self._websocket
|
||||
raise Exception("Websocket not connected")
|
||||
|
||||
async def _receive_messages(self):
|
||||
"""Receive and process websocket messages.
|
||||
|
||||
Continuously processes messages from the websocket connection.
|
||||
"""
|
||||
async for response in self._get_websocket():
|
||||
try:
|
||||
response = await self._ws_client.recv()
|
||||
|
||||
headers, payload = decode_event(response)
|
||||
|
||||
if headers.get(":message-type") == "event":
|
||||
@@ -527,11 +500,5 @@ class AWSTranscribeSTTService(STTService):
|
||||
else:
|
||||
logger.debug(f"{self} Other message type received: {headers}")
|
||||
logger.debug(f"{self} Payload: {payload}")
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
await self.push_error(
|
||||
error_msg=f"WebSocket connection closed in receive loop", exception=e
|
||||
)
|
||||
break
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
break
|
||||
logger.warning(f"Error processing message: {e}")
|
||||
|
||||
@@ -23,7 +23,6 @@ from pipecat import __version__ as pipecat_version
|
||||
from pipecat.frames.frames import (
|
||||
CancelFrame,
|
||||
EndFrame,
|
||||
ErrorFrame,
|
||||
Frame,
|
||||
InterimTranscriptionFrame,
|
||||
StartFrame,
|
||||
@@ -31,7 +30,7 @@ from pipecat.frames.frames import (
|
||||
TranslationFrame,
|
||||
)
|
||||
from pipecat.services.gladia.config import GladiaInputParams
|
||||
from pipecat.services.stt_service import STTService
|
||||
from pipecat.services.stt_service import WebsocketSTTService
|
||||
from pipecat.transcriptions.language import Language, resolve_language
|
||||
from pipecat.utils.time import time_now_iso8601
|
||||
from pipecat.utils.tracing.service_decorators import traced_stt
|
||||
@@ -176,7 +175,7 @@ class _InputParamsDescriptor:
|
||||
return GladiaInputParams
|
||||
|
||||
|
||||
class GladiaSTTService(STTService):
|
||||
class GladiaSTTService(WebsocketSTTService):
|
||||
"""Speech-to-Text service using Gladia's API.
|
||||
|
||||
This service connects to Gladia's WebSocket API for real-time transcription
|
||||
@@ -202,8 +201,6 @@ class GladiaSTTService(STTService):
|
||||
sample_rate: Optional[int] = None,
|
||||
model: str = "solaria-1",
|
||||
params: Optional[GladiaInputParams] = None,
|
||||
max_reconnection_attempts: int = 5,
|
||||
reconnection_delay: float = 1.0,
|
||||
max_buffer_size: int = 1024 * 1024 * 20, # 20MB default buffer
|
||||
**kwargs,
|
||||
):
|
||||
@@ -222,8 +219,6 @@ class GladiaSTTService(STTService):
|
||||
sample_rate: Audio sample rate in Hz. If None, uses service default.
|
||||
model: Model to use for transcription. Defaults to "solaria-1".
|
||||
params: Additional configuration parameters for Gladia service.
|
||||
max_reconnection_attempts: Maximum number of reconnection attempts. Defaults to 5.
|
||||
reconnection_delay: Initial delay between reconnection attempts in seconds.
|
||||
max_buffer_size: Maximum size of audio buffer in bytes. Defaults to 20MB.
|
||||
**kwargs: Additional arguments passed to the STTService parent class.
|
||||
"""
|
||||
@@ -256,16 +251,13 @@ class GladiaSTTService(STTService):
|
||||
self._url = url
|
||||
self.set_model_name(model)
|
||||
self._params = params
|
||||
self._websocket = None
|
||||
self._receive_task = None
|
||||
self._keepalive_task = None
|
||||
self._settings = {}
|
||||
|
||||
# Reconnection settings
|
||||
self._max_reconnection_attempts = max_reconnection_attempts
|
||||
self._reconnection_delay = reconnection_delay
|
||||
self._reconnection_attempts = 0
|
||||
# Session management
|
||||
self._session_url = None
|
||||
self._session_id = None
|
||||
self._connection_active = False
|
||||
|
||||
# Audio buffer management
|
||||
@@ -274,9 +266,8 @@ class GladiaSTTService(STTService):
|
||||
self._max_buffer_size = max_buffer_size
|
||||
self._buffer_lock = asyncio.Lock()
|
||||
|
||||
# Connection management
|
||||
self._connection_task = None
|
||||
self._should_reconnect = True
|
||||
def __str__(self):
|
||||
return f"{self.name} [{self._session_id}]"
|
||||
|
||||
def can_generate_metrics(self) -> bool:
|
||||
"""Check if the service can generate performance metrics.
|
||||
@@ -355,11 +346,7 @@ class GladiaSTTService(STTService):
|
||||
frame: The start frame triggering service startup.
|
||||
"""
|
||||
await super().start(frame)
|
||||
if self._connection_task:
|
||||
return
|
||||
|
||||
self._should_reconnect = True
|
||||
self._connection_task = self.create_task(self._connection_handler())
|
||||
await self._connect()
|
||||
|
||||
async def stop(self, frame: EndFrame):
|
||||
"""Stop the Gladia STT websocket connection.
|
||||
@@ -368,14 +355,8 @@ class GladiaSTTService(STTService):
|
||||
frame: The end frame triggering service shutdown.
|
||||
"""
|
||||
await super().stop(frame)
|
||||
self._should_reconnect = False
|
||||
await self._send_stop_recording()
|
||||
|
||||
if self._connection_task:
|
||||
await self.cancel_task(self._connection_task)
|
||||
self._connection_task = None
|
||||
|
||||
await self._cleanup_connection()
|
||||
await self._disconnect()
|
||||
|
||||
async def cancel(self, frame: CancelFrame):
|
||||
"""Cancel the Gladia STT websocket connection.
|
||||
@@ -384,13 +365,7 @@ class GladiaSTTService(STTService):
|
||||
frame: The cancel frame triggering service cancellation.
|
||||
"""
|
||||
await super().cancel(frame)
|
||||
self._should_reconnect = False
|
||||
|
||||
if self._connection_task:
|
||||
await self.cancel_task(self._connection_task)
|
||||
self._connection_task = None
|
||||
|
||||
await self._cleanup_connection()
|
||||
await self._disconnect()
|
||||
|
||||
async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
|
||||
"""Run speech-to-text on audio data.
|
||||
@@ -412,88 +387,93 @@ class GladiaSTTService(STTService):
|
||||
trim_size = len(self._audio_buffer) - self._max_buffer_size
|
||||
self._audio_buffer = self._audio_buffer[trim_size:]
|
||||
self._bytes_sent = max(0, self._bytes_sent - trim_size)
|
||||
logger.warning(f"Audio buffer exceeded max size, trimmed {trim_size} bytes")
|
||||
logger.warning(f"{self} Audio buffer exceeded max size, trimmed {trim_size} bytes")
|
||||
|
||||
# Send audio if connected
|
||||
if self._connection_active and self._websocket and self._websocket.state is State.OPEN:
|
||||
try:
|
||||
await self._send_audio(audio)
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
logger.warning(f"Websocket closed while sending audio chunk: {e}")
|
||||
logger.warning(f"{self} Websocket closed while sending audio chunk: {e}")
|
||||
self._connection_active = False
|
||||
|
||||
yield None
|
||||
|
||||
async def _connection_handler(self):
|
||||
"""Handle WebSocket connection with automatic reconnection."""
|
||||
while self._should_reconnect:
|
||||
try:
|
||||
# Initialize session if needed
|
||||
if not self._session_url:
|
||||
settings = self._prepare_settings()
|
||||
response = await self._setup_gladia(settings)
|
||||
self._session_url = response["url"]
|
||||
self._reconnection_attempts = 0
|
||||
logger.info(f"Session URL : {self._session_url}")
|
||||
async def _connect(self):
|
||||
"""Connect to the Gladia service.
|
||||
|
||||
# Connect with automatic reconnection
|
||||
async with websocket_connect(self._session_url) as websocket:
|
||||
try:
|
||||
self._websocket = websocket
|
||||
self._connection_active = True
|
||||
logger.debug(f"{self} Connected to Gladia WebSocket")
|
||||
Initializes the session if needed and establishes websocket connection.
|
||||
"""
|
||||
# Initialize session if needed
|
||||
if not self._session_url:
|
||||
settings = self._prepare_settings()
|
||||
response = await self._setup_gladia(settings)
|
||||
self._session_url = response["url"]
|
||||
self._session_id = response["id"]
|
||||
logger.info(f"{self} Session URL: {self._session_url}")
|
||||
|
||||
# Send buffered audio if any
|
||||
await self._send_buffered_audio()
|
||||
await self._connect_websocket()
|
||||
|
||||
# Start tasks
|
||||
self._receive_task = self.create_task(self._receive_task_handler())
|
||||
self._keepalive_task = self.create_task(self._keepalive_task_handler())
|
||||
if self._websocket and not self._receive_task:
|
||||
self._receive_task = self.create_task(self._receive_task_handler(self._report_error))
|
||||
|
||||
# Wait for tasks to complete
|
||||
await asyncio.gather(self._receive_task, self._keepalive_task)
|
||||
if self._websocket and not self._keepalive_task:
|
||||
self._keepalive_task = self.create_task(self._keepalive_task_handler())
|
||||
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
logger.warning(f"WebSocket connection closed: {e}")
|
||||
self._connection_active = False
|
||||
async def _disconnect(self):
|
||||
"""Disconnect from the Gladia service.
|
||||
|
||||
# Clean up tasks
|
||||
if self._receive_task:
|
||||
await self.cancel_task(self._receive_task)
|
||||
if self._keepalive_task:
|
||||
await self.cancel_task(self._keepalive_task)
|
||||
|
||||
# Attempt reconnect using helper
|
||||
if not await self._maybe_reconnect():
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
self._connection_active = False
|
||||
|
||||
if not self._should_reconnect:
|
||||
break
|
||||
|
||||
# Reset session URL to get a new one
|
||||
self._session_url = None
|
||||
await asyncio.sleep(self._reconnection_delay)
|
||||
|
||||
async def _cleanup_connection(self):
|
||||
"""Clean up connection resources."""
|
||||
Cleans up tasks and closes websocket connection.
|
||||
"""
|
||||
self._connection_active = False
|
||||
|
||||
if self._keepalive_task:
|
||||
await self.cancel_task(self._keepalive_task)
|
||||
self._keepalive_task = None
|
||||
|
||||
if self._websocket:
|
||||
await self._websocket.close()
|
||||
self._websocket = None
|
||||
|
||||
if self._receive_task:
|
||||
await self.cancel_task(self._receive_task)
|
||||
self._receive_task = None
|
||||
|
||||
await self._disconnect_websocket()
|
||||
|
||||
async def _connect_websocket(self):
|
||||
"""Establish the websocket connection to Gladia."""
|
||||
try:
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
return
|
||||
|
||||
logger.debug(f"{self}Connecting to Gladia WebSocket")
|
||||
|
||||
self._websocket = await websocket_connect(self._session_url)
|
||||
self._connection_active = True
|
||||
|
||||
# Reset byte tracking for new connection
|
||||
async with self._buffer_lock:
|
||||
self._bytes_sent = 0
|
||||
|
||||
await self._call_event_handler("on_connected")
|
||||
|
||||
# Send buffered audio if any
|
||||
await self._send_buffered_audio()
|
||||
|
||||
logger.debug(f"{self} Connected to Gladia WebSocket")
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unable to connect to Gladia: {e}", exception=e)
|
||||
raise
|
||||
|
||||
async def _disconnect_websocket(self):
|
||||
"""Close the websocket connection to Gladia."""
|
||||
try:
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
logger.debug(f"{self} Disconnecting from Gladia WebSocket")
|
||||
await self._websocket.close()
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Error closing websocket: {e}", exception=e)
|
||||
finally:
|
||||
self._websocket = None
|
||||
await self._call_event_handler("on_disconnected")
|
||||
|
||||
async def _setup_gladia(self, settings: Dict[str, Any]):
|
||||
async with aiohttp.ClientSession() as session:
|
||||
params = {}
|
||||
@@ -510,10 +490,10 @@ class GladiaSTTService(STTService):
|
||||
else:
|
||||
error_text = await response.text()
|
||||
logger.error(
|
||||
f"Gladia error: {response.status}: {error_text or response.reason}"
|
||||
f"{self} Gladia error: {response.status}: {error_text or response.reason}"
|
||||
)
|
||||
raise Exception(
|
||||
f"Failed to initialize Gladia session: {response.status} - {error_text}"
|
||||
f"{self} Failed to initialize Gladia session: {response.status} - {error_text}"
|
||||
)
|
||||
|
||||
@traced_stt
|
||||
@@ -541,28 +521,26 @@ class GladiaSTTService(STTService):
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
await self._websocket.send(json.dumps({"type": "stop_recording"}))
|
||||
|
||||
async def _keepalive_task_handler(self):
|
||||
"""Send periodic empty audio chunks to keep the connection alive."""
|
||||
try:
|
||||
KEEPALIVE_SLEEP = 20
|
||||
while self._connection_active:
|
||||
# Send keepalive (Gladia times out after 30 seconds)
|
||||
await asyncio.sleep(KEEPALIVE_SLEEP)
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
# Send an empty audio chunk as keepalive
|
||||
empty_audio = b""
|
||||
await self._send_audio(empty_audio)
|
||||
else:
|
||||
logger.debug("Websocket closed, stopping keepalive")
|
||||
break
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
logger.debug("Connection closed during keepalive")
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
def _get_websocket(self):
|
||||
"""Get the current WebSocket connection.
|
||||
|
||||
async def _receive_task_handler(self):
|
||||
try:
|
||||
async for message in self._websocket:
|
||||
Returns:
|
||||
The WebSocket connection.
|
||||
|
||||
Raises:
|
||||
Exception: If WebSocket is not connected.
|
||||
"""
|
||||
if self._websocket:
|
||||
return self._websocket
|
||||
raise Exception("Websocket not connected")
|
||||
|
||||
async def _receive_messages(self):
|
||||
"""Receive and process websocket messages.
|
||||
|
||||
Continuously processes messages from the websocket connection.
|
||||
"""
|
||||
async for message in self._get_websocket():
|
||||
try:
|
||||
content = json.loads(message)
|
||||
|
||||
# Handle audio chunk acknowledgments
|
||||
@@ -617,26 +595,24 @@ class GladiaSTTService(STTService):
|
||||
translation, "", time_now_iso8601(), translated_language
|
||||
)
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"{self} Received non-JSON message: {message}")
|
||||
|
||||
async def _keepalive_task_handler(self):
|
||||
"""Send periodic empty audio chunks to keep the connection alive."""
|
||||
try:
|
||||
KEEPALIVE_SLEEP = 20
|
||||
while self._connection_active:
|
||||
# Send keepalive (Gladia times out after 30 seconds)
|
||||
await asyncio.sleep(KEEPALIVE_SLEEP)
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
# Send an empty audio chunk as keepalive
|
||||
empty_audio = b""
|
||||
await self._send_audio(empty_audio)
|
||||
else:
|
||||
logger.debug(f"{self} Websocket closed, stopping keepalive")
|
||||
break
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
# Expected when closing the connection
|
||||
pass
|
||||
logger.debug(f"{self} Connection closed during keepalive")
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
|
||||
async def _maybe_reconnect(self) -> bool:
|
||||
"""Handle exponential backoff reconnection logic."""
|
||||
if not self._should_reconnect:
|
||||
return False
|
||||
self._reconnection_attempts += 1
|
||||
if self._reconnection_attempts > self._max_reconnection_attempts:
|
||||
await self.push_error(
|
||||
error_msg=f"Max reconnection attempts ({self._max_reconnection_attempts}) reached",
|
||||
)
|
||||
self._should_reconnect = False
|
||||
return False
|
||||
delay = self._reconnection_delay * (2 ** (self._reconnection_attempts - 1))
|
||||
logger.debug(
|
||||
f"{self} Reconnecting in {delay} seconds (attempt {self._reconnection_attempts}/{self._max_reconnection_attempts})"
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
return True
|
||||
|
||||
@@ -17,7 +17,6 @@ from pydantic import BaseModel
|
||||
from pipecat.frames.frames import (
|
||||
CancelFrame,
|
||||
EndFrame,
|
||||
ErrorFrame,
|
||||
Frame,
|
||||
InterimTranscriptionFrame,
|
||||
StartFrame,
|
||||
@@ -25,13 +24,12 @@ from pipecat.frames.frames import (
|
||||
UserStoppedSpeakingFrame,
|
||||
)
|
||||
from pipecat.processors.frame_processor import FrameDirection
|
||||
from pipecat.services.stt_service import STTService
|
||||
from pipecat.services.stt_service import WebsocketSTTService
|
||||
from pipecat.transcriptions.language import Language
|
||||
from pipecat.utils.time import time_now_iso8601
|
||||
from pipecat.utils.tracing.service_decorators import traced_stt
|
||||
|
||||
try:
|
||||
import websockets
|
||||
from websockets.asyncio.client import connect as websocket_connect
|
||||
from websockets.protocol import State
|
||||
except ModuleNotFoundError as e:
|
||||
@@ -134,7 +132,7 @@ def _prepare_language_hints(
|
||||
return list(set(prepared_languages))
|
||||
|
||||
|
||||
class SonioxSTTService(STTService):
|
||||
class SonioxSTTService(WebsocketSTTService):
|
||||
"""Speech-to-Text service using Soniox's WebSocket API.
|
||||
|
||||
This service connects to Soniox's WebSocket API for real-time transcription
|
||||
@@ -173,7 +171,6 @@ class SonioxSTTService(STTService):
|
||||
self.set_model_name(params.model)
|
||||
self._params = params
|
||||
self._vad_force_turn_endpoint = vad_force_turn_endpoint
|
||||
self._websocket = None
|
||||
|
||||
self._final_transcription_buffer = []
|
||||
self._last_tokens_received: Optional[float] = None
|
||||
@@ -188,59 +185,7 @@ class SonioxSTTService(STTService):
|
||||
frame: The start frame containing initialization parameters.
|
||||
"""
|
||||
await super().start(frame)
|
||||
if self._websocket:
|
||||
return
|
||||
|
||||
self._websocket = await websocket_connect(self._url)
|
||||
|
||||
if not self._websocket:
|
||||
await self.push_error(error_msg=f"Unable to connect to Soniox API at {self._url}")
|
||||
|
||||
# If vad_force_turn_endpoint is not enabled, we need to enable endpoint detection.
|
||||
# Either one or the other is required.
|
||||
enable_endpoint_detection = not self._vad_force_turn_endpoint
|
||||
|
||||
context = self._params.context
|
||||
if isinstance(context, SonioxContextObject):
|
||||
context = context.model_dump()
|
||||
|
||||
# Send the initial configuration message.
|
||||
config = {
|
||||
"api_key": self._api_key,
|
||||
"model": self._model_name,
|
||||
"audio_format": self._params.audio_format,
|
||||
"num_channels": self._params.num_channels or 1,
|
||||
"enable_endpoint_detection": enable_endpoint_detection,
|
||||
"sample_rate": self.sample_rate,
|
||||
"language_hints": _prepare_language_hints(self._params.language_hints),
|
||||
"context": context,
|
||||
"enable_speaker_diarization": self._params.enable_speaker_diarization,
|
||||
"enable_language_identification": self._params.enable_language_identification,
|
||||
"client_reference_id": self._params.client_reference_id,
|
||||
}
|
||||
|
||||
# Send the configuration message.
|
||||
await self._websocket.send(json.dumps(config))
|
||||
|
||||
if self._websocket and not self._receive_task:
|
||||
self._receive_task = self.create_task(self._receive_task_handler())
|
||||
if self._websocket and not self._keepalive_task:
|
||||
self._keepalive_task = self.create_task(self._keepalive_task_handler())
|
||||
|
||||
async def _cleanup(self):
|
||||
if self._keepalive_task:
|
||||
await self.cancel_task(self._keepalive_task)
|
||||
self._keepalive_task = None
|
||||
|
||||
if self._websocket:
|
||||
await self._websocket.close()
|
||||
self._websocket = None
|
||||
|
||||
if self._receive_task:
|
||||
# Task cannot cancel itself. If task called _cleanup() we expect it to cancel itself.
|
||||
if self._receive_task != asyncio.current_task():
|
||||
await self._receive_task
|
||||
self._receive_task = None
|
||||
await self._connect()
|
||||
|
||||
async def stop(self, frame: EndFrame):
|
||||
"""Stop the Soniox STT websocket connection.
|
||||
@@ -253,6 +198,7 @@ class SonioxSTTService(STTService):
|
||||
"""
|
||||
await super().stop(frame)
|
||||
await self._send_stop_recording()
|
||||
await self._disconnect()
|
||||
|
||||
async def cancel(self, frame: CancelFrame):
|
||||
"""Cancel the Soniox STT websocket connection.
|
||||
@@ -265,7 +211,7 @@ class SonioxSTTService(STTService):
|
||||
frame: The cancel frame.
|
||||
"""
|
||||
await super().cancel(frame)
|
||||
await self._cleanup()
|
||||
await self._disconnect()
|
||||
|
||||
async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
|
||||
"""Send audio data to Soniox STT Service.
|
||||
@@ -311,28 +257,110 @@ class SonioxSTTService(STTService):
|
||||
# Send stop recording message
|
||||
await self._websocket.send("")
|
||||
|
||||
async def _keepalive_task_handler(self):
|
||||
"""Connection has to be open all the time."""
|
||||
async def _connect(self):
|
||||
"""Connect to the Soniox service.
|
||||
|
||||
Establishes websocket connection and starts receive and keepalive tasks.
|
||||
"""
|
||||
await self._connect_websocket()
|
||||
|
||||
if self._websocket and not self._receive_task:
|
||||
self._receive_task = self.create_task(self._receive_task_handler(self._report_error))
|
||||
|
||||
if self._websocket and not self._keepalive_task:
|
||||
self._keepalive_task = self.create_task(self._keepalive_task_handler())
|
||||
|
||||
async def _disconnect(self):
|
||||
"""Disconnect from the Soniox service.
|
||||
|
||||
Cleans up tasks and closes websocket connection.
|
||||
"""
|
||||
if self._keepalive_task:
|
||||
await self.cancel_task(self._keepalive_task)
|
||||
self._keepalive_task = None
|
||||
|
||||
if self._receive_task:
|
||||
await self.cancel_task(self._receive_task)
|
||||
self._receive_task = None
|
||||
|
||||
await self._disconnect_websocket()
|
||||
|
||||
async def _connect_websocket(self):
|
||||
"""Establish the websocket connection to Soniox."""
|
||||
try:
|
||||
while True:
|
||||
logger.trace("Sending keepalive message")
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
await self._websocket.send(KEEPALIVE_MESSAGE)
|
||||
else:
|
||||
logger.debug("WebSocket connection closed.")
|
||||
break
|
||||
await asyncio.sleep(5)
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
return
|
||||
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
# Expected when closing the connection
|
||||
logger.debug("WebSocket connection closed, keepalive task stopped.")
|
||||
logger.debug("Connecting to Soniox STT")
|
||||
|
||||
self._websocket = await websocket_connect(self._url)
|
||||
|
||||
if not self._websocket:
|
||||
await self.push_error(error_msg=f"Unable to connect to Soniox API at {self._url}")
|
||||
raise Exception(f"Unable to connect to Soniox API at {self._url}")
|
||||
|
||||
# If vad_force_turn_endpoint is not enabled, we need to enable endpoint detection.
|
||||
# Either one or the other is required.
|
||||
enable_endpoint_detection = not self._vad_force_turn_endpoint
|
||||
|
||||
context = self._params.context
|
||||
if isinstance(context, SonioxContextObject):
|
||||
context = context.model_dump()
|
||||
|
||||
# Send the initial configuration message.
|
||||
config = {
|
||||
"api_key": self._api_key,
|
||||
"model": self._model_name,
|
||||
"audio_format": self._params.audio_format,
|
||||
"num_channels": self._params.num_channels or 1,
|
||||
"enable_endpoint_detection": enable_endpoint_detection,
|
||||
"sample_rate": self.sample_rate,
|
||||
"language_hints": _prepare_language_hints(self._params.language_hints),
|
||||
"context": context,
|
||||
"enable_speaker_diarization": self._params.enable_speaker_diarization,
|
||||
"enable_language_identification": self._params.enable_language_identification,
|
||||
"client_reference_id": self._params.client_reference_id,
|
||||
}
|
||||
|
||||
# Send the configuration message.
|
||||
await self._websocket.send(json.dumps(config))
|
||||
|
||||
await self._call_event_handler("on_connected")
|
||||
logger.debug("Connected to Soniox STT")
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
await self.push_error(error_msg=f"Unable to connect to Soniox: {e}", exception=e)
|
||||
raise
|
||||
|
||||
async def _receive_task_handler(self):
|
||||
if not self._websocket:
|
||||
return
|
||||
async def _disconnect_websocket(self):
|
||||
"""Close the websocket connection to Soniox."""
|
||||
try:
|
||||
if self._websocket:
|
||||
logger.debug("Disconnecting from Soniox STT")
|
||||
await self._websocket.close()
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Error closing websocket: {e}", exception=e)
|
||||
finally:
|
||||
self._websocket = None
|
||||
await self._call_event_handler("on_disconnected")
|
||||
|
||||
def _get_websocket(self):
|
||||
"""Get the current WebSocket connection.
|
||||
|
||||
Returns:
|
||||
The WebSocket connection.
|
||||
|
||||
Raises:
|
||||
Exception: If WebSocket is not connected.
|
||||
"""
|
||||
if self._websocket:
|
||||
return self._websocket
|
||||
raise Exception("Websocket not connected")
|
||||
|
||||
async def _receive_messages(self):
|
||||
"""Receive and process websocket messages.
|
||||
|
||||
Continuously processes messages from the websocket connection.
|
||||
"""
|
||||
# Transcription frame will be only sent after we get the "endpoint" event.
|
||||
self._final_transcription_buffer = []
|
||||
|
||||
@@ -351,8 +379,8 @@ class SonioxSTTService(STTService):
|
||||
await self.stop_processing_metrics()
|
||||
self._final_transcription_buffer = []
|
||||
|
||||
try:
|
||||
async for message in self._websocket:
|
||||
async for message in self._get_websocket():
|
||||
try:
|
||||
content = json.loads(message)
|
||||
|
||||
tokens = content["tokens"]
|
||||
@@ -404,7 +432,7 @@ class SonioxSTTService(STTService):
|
||||
# In case of error, still send the final transcript (if any remaining in the buffer).
|
||||
await send_endpoint_transcript()
|
||||
await self.push_error(
|
||||
error_msg=f"Error: {error_code} (_receive_task_handler) - {error_message}"
|
||||
error_msg=f"Error: {error_code} (_receive_messages) - {error_message}"
|
||||
)
|
||||
|
||||
finished = content.get("finished")
|
||||
@@ -412,11 +440,24 @@ class SonioxSTTService(STTService):
|
||||
# When finished, still send the final transcript (if any remaining in the buffer).
|
||||
await send_endpoint_transcript()
|
||||
logger.debug("Transcription finished.")
|
||||
await self._cleanup()
|
||||
return
|
||||
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
# Expected when closing the connection.
|
||||
pass
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Received non-JSON message: {message}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error processing message: {e}")
|
||||
|
||||
async def _keepalive_task_handler(self):
|
||||
"""Connection has to be open all the time."""
|
||||
try:
|
||||
while True:
|
||||
logger.trace("Sending keepalive message")
|
||||
if self._websocket and self._websocket.state is State.OPEN:
|
||||
await self._websocket.send(KEEPALIVE_MESSAGE)
|
||||
else:
|
||||
logger.debug("WebSocket connection closed.")
|
||||
break
|
||||
await asyncio.sleep(5)
|
||||
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Error receiving message: {e}", exception=e)
|
||||
logger.debug(f"Keepalive task stopped: {e}")
|
||||
|
||||
Reference in New Issue
Block a user