529 lines
19 KiB
Python
529 lines
19 KiB
Python
|
"""Deepgram-based transcription engine using WebSocket streaming.
|
||
|
|
|
||
|
|
Supports two modes:
|
||
|
|
- Managed mode: connects to a proxy server that handles Deepgram credentials
|
||
|
|
- BYOK mode: connects directly to the Deepgram API with a user-provided key
|
||
|
|
|
||
|
|
Implements the same duck-type interface as RealtimeTranscriptionEngine so
|
||
|
|
MainWindow can use it as a drop-in replacement.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
import numpy as np
|
||
|
|
import threading
|
||
|
|
from datetime import datetime
|
||
|
|
from queue import Queue, Empty
|
||
|
|
from typing import Optional, Callable
|
||
|
|
|
||
|
|
from client.transcription_engine_realtime import TranscriptionResult
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class DeepgramTranscriptionEngine:
|
||
|
|
"""
|
||
|
|
Transcription engine that streams audio to Deepgram via WebSocket.
|
||
|
|
|
||
|
|
In managed mode the connection goes through a proxy at
|
||
|
|
``wss://<server>/ws/transcribe`` which handles authentication and
|
||
|
|
Deepgram credentials. In BYOK (bring-your-own-key) mode the
|
||
|
|
connection goes directly to the Deepgram API.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
# Construction / configuration
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
|
||
|
|
def __init__(self, config, user_name: str = "User", input_device_index: Optional[int] = None):
|
||
|
|
"""
|
||
|
|
Initialise the engine from a :class:`client.config.Config` object.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
config: Application ``Config`` instance.
|
||
|
|
user_name: Display name attached to transcriptions.
|
||
|
|
input_device_index: Index of the audio input device to use
|
||
|
|
(``None`` for the system default).
|
||
|
|
"""
|
||
|
|
self.config = config
|
||
|
|
self.user_name = user_name
|
||
|
|
self.input_device_index = input_device_index
|
||
|
|
|
||
|
|
# Mode: 'managed' (proxy) or 'byok' (direct Deepgram)
|
||
|
|
self.mode: str = config.get("remote.mode", "managed")
|
||
|
|
|
||
|
|
# Managed-mode settings
|
||
|
|
self.server_url: str = config.get("remote.server_url", "")
|
||
|
|
self.auth_token: str = config.get("remote.auth_token", "")
|
||
|
|
|
||
|
|
# BYOK-mode settings
|
||
|
|
self.byok_api_key: str = config.get("remote.byok_api_key", "")
|
||
|
|
|
||
|
|
# Deepgram model / language (used in both modes)
|
||
|
|
self.deepgram_model: str = config.get("remote.deepgram_model", "nova-2")
|
||
|
|
self.language: str = config.get("remote.language", "en-US")
|
||
|
|
|
||
|
|
# Audio parameters
|
||
|
|
self.sample_rate: int = 16000
|
||
|
|
self.channels: int = 1
|
||
|
|
self.blocksize: int = 4096
|
||
|
|
|
||
|
|
# Callbacks
|
||
|
|
self.realtime_callback: Optional[Callable[[TranscriptionResult], None]] = None
|
||
|
|
self.final_callback: Optional[Callable[[TranscriptionResult], None]] = None
|
||
|
|
self._on_error: Optional[Callable[[str], None]] = None
|
||
|
|
self._on_credits_low: Optional[Callable[[int], None]] = None
|
||
|
|
|
||
|
|
# Internal state
|
||
|
|
self._is_initialized: bool = False
|
||
|
|
self._is_recording: bool = False
|
||
|
|
self._stop_event: threading.Event = threading.Event()
|
||
|
|
self._audio_queue: Queue = Queue()
|
||
|
|
|
||
|
|
# Asyncio event loop running in a daemon thread
|
||
|
|
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
||
|
|
self._thread: Optional[threading.Thread] = None
|
||
|
|
|
||
|
|
# WebSocket handle (set inside the async context)
|
||
|
|
self._ws = None
|
||
|
|
|
||
|
|
# sounddevice InputStream
|
||
|
|
self._stream = None
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
# Callback setters
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
|
||
|
|
def set_callbacks(
|
||
|
|
self,
|
||
|
|
realtime_callback: Optional[Callable[[TranscriptionResult], None]] = None,
|
||
|
|
final_callback: Optional[Callable[[TranscriptionResult], None]] = None,
|
||
|
|
):
|
||
|
|
"""Set transcription result callbacks (matches RealtimeTranscriptionEngine API)."""
|
||
|
|
self.realtime_callback = realtime_callback
|
||
|
|
self.final_callback = final_callback
|
||
|
|
|
||
|
|
def set_error_callback(self, fn: Optional[Callable[[str], None]]):
|
||
|
|
"""Set a callback invoked on errors. ``fn`` receives a string message."""
|
||
|
|
self._on_error = fn
|
||
|
|
|
||
|
|
def set_credits_low_callback(self, fn: Optional[Callable[[int], None]]):
|
||
|
|
"""Set a callback for low-credit warnings. ``fn`` receives seconds remaining."""
|
||
|
|
self._on_credits_low = fn
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
# Public interface (duck-typed with RealtimeTranscriptionEngine)
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
|
||
|
|
def initialize(self) -> bool:
|
||
|
|
"""Validate configuration and mark the engine as ready.
|
||
|
|
|
||
|
|
Returns ``True`` when the engine is ready to start recording.
|
||
|
|
"""
|
||
|
|
if self._is_initialized:
|
||
|
|
return True
|
||
|
|
|
||
|
|
if self.mode == "managed":
|
||
|
|
if not self.server_url:
|
||
|
|
logger.error("Managed mode requires a server URL (remote.server_url)")
|
||
|
|
return False
|
||
|
|
if not self.auth_token:
|
||
|
|
logger.error("Managed mode requires an auth token (remote.auth_token)")
|
||
|
|
return False
|
||
|
|
elif self.mode == "byok":
|
||
|
|
if not self.byok_api_key:
|
||
|
|
logger.error("BYOK mode requires an API key (remote.byok_api_key)")
|
||
|
|
return False
|
||
|
|
else:
|
||
|
|
logger.error("Unknown remote mode: %s (expected 'managed' or 'byok')", self.mode)
|
||
|
|
return False
|
||
|
|
|
||
|
|
self._is_initialized = True
|
||
|
|
logger.info("DeepgramTranscriptionEngine initialised in %s mode", self.mode)
|
||
|
|
return True
|
||
|
|
|
||
|
|
def start_recording(self) -> bool:
|
||
|
|
"""Open the audio stream and connect the WebSocket.
|
||
|
|
|
||
|
|
Returns ``True`` on success.
|
||
|
|
"""
|
||
|
|
if not self._is_initialized:
|
||
|
|
logger.error("Engine not initialised -- call initialize() first")
|
||
|
|
return False
|
||
|
|
|
||
|
|
if self._is_recording:
|
||
|
|
return True
|
||
|
|
|
||
|
|
self._stop_event.clear()
|
||
|
|
self._is_recording = True
|
||
|
|
|
||
|
|
# Start the asyncio event-loop thread (handles WS send/receive)
|
||
|
|
self._thread = threading.Thread(target=self._run_event_loop, daemon=True)
|
||
|
|
self._thread.start()
|
||
|
|
|
||
|
|
# Start the audio capture stream
|
||
|
|
try:
|
||
|
|
self._start_audio_stream()
|
||
|
|
except Exception as exc:
|
||
|
|
logger.error("Failed to open audio stream: %s", exc)
|
||
|
|
self._is_recording = False
|
||
|
|
self._stop_event.set()
|
||
|
|
return False
|
||
|
|
|
||
|
|
logger.info("Recording started")
|
||
|
|
return True
|
||
|
|
|
||
|
|
def stop_recording(self):
|
||
|
|
"""Stop audio capture and close the WebSocket."""
|
||
|
|
if not self._is_recording:
|
||
|
|
return
|
||
|
|
|
||
|
|
self._is_recording = False
|
||
|
|
self._stop_event.set()
|
||
|
|
|
||
|
|
# Stop audio stream
|
||
|
|
self._stop_audio_stream()
|
||
|
|
|
||
|
|
# Close WebSocket from outside the event-loop thread
|
||
|
|
if self._ws is not None and self._loop is not None and not self._loop.is_closed():
|
||
|
|
asyncio.run_coroutine_threadsafe(self._close_ws(), self._loop)
|
||
|
|
|
||
|
|
# Wait for the thread to finish
|
||
|
|
if self._thread is not None:
|
||
|
|
self._thread.join(timeout=5)
|
||
|
|
self._thread = None
|
||
|
|
|
||
|
|
logger.info("Recording stopped")
|
||
|
|
|
||
|
|
def stop(self):
|
||
|
|
"""Full shutdown -- stop recording and release all resources."""
|
||
|
|
self.stop_recording()
|
||
|
|
self._is_initialized = False
|
||
|
|
logger.info("DeepgramTranscriptionEngine shut down")
|
||
|
|
|
||
|
|
def is_ready(self) -> bool:
|
||
|
|
"""Return ``True`` if the engine has been successfully initialised."""
|
||
|
|
return self._is_initialized
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
# Audio capture (sounddevice)
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
|
||
|
|
def _start_audio_stream(self):
|
||
|
|
"""Open a ``sounddevice.InputStream`` that feeds the audio queue."""
|
||
|
|
import sounddevice as sd
|
||
|
|
|
||
|
|
def _audio_callback(indata, frames, time_info, status): # noqa: ARG001
|
||
|
|
if status:
|
||
|
|
logger.warning("Audio stream status: %s", status)
|
||
|
|
if self._is_recording:
|
||
|
|
# float32 -> int16 PCM bytes
|
||
|
|
pcm = (indata * 32767).astype(np.int16).tobytes()
|
||
|
|
self._audio_queue.put(pcm)
|
||
|
|
|
||
|
|
self._stream = sd.InputStream(
|
||
|
|
samplerate=self.sample_rate,
|
||
|
|
blocksize=self.blocksize,
|
||
|
|
channels=self.channels,
|
||
|
|
dtype="float32",
|
||
|
|
device=self.input_device_index,
|
||
|
|
callback=_audio_callback,
|
||
|
|
)
|
||
|
|
self._stream.start()
|
||
|
|
|
||
|
|
def _stop_audio_stream(self):
|
||
|
|
"""Close the audio input stream."""
|
||
|
|
if self._stream is not None:
|
||
|
|
try:
|
||
|
|
self._stream.stop()
|
||
|
|
self._stream.close()
|
||
|
|
except Exception as exc:
|
||
|
|
logger.debug("Error closing audio stream: %s", exc)
|
||
|
|
finally:
|
||
|
|
self._stream = None
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
# Asyncio event-loop (runs in daemon thread)
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
|
||
|
|
def _run_event_loop(self):
|
||
|
|
"""Entry point for the daemon thread -- runs the async event loop."""
|
||
|
|
self._loop = asyncio.new_event_loop()
|
||
|
|
asyncio.set_event_loop(self._loop)
|
||
|
|
try:
|
||
|
|
self._loop.run_until_complete(self._ws_lifecycle())
|
||
|
|
except Exception as exc:
|
||
|
|
logger.error("Event-loop error: %s", exc)
|
||
|
|
finally:
|
||
|
|
try:
|
||
|
|
self._loop.run_until_complete(self._loop.shutdown_asyncgens())
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
self._loop.close()
|
||
|
|
self._loop = None
|
||
|
|
|
||
|
|
async def _ws_lifecycle(self):
|
||
|
|
"""Connect, authenticate (if managed), then run send/receive loops."""
|
||
|
|
import websockets
|
||
|
|
|
||
|
|
try:
|
||
|
|
ws_url, extra_headers = self._build_ws_url_and_headers()
|
||
|
|
|
||
|
|
logger.info("Connecting to %s", ws_url)
|
||
|
|
self._ws = await websockets.connect(
|
||
|
|
ws_url,
|
||
|
|
additional_headers=extra_headers,
|
||
|
|
ping_interval=20,
|
||
|
|
ping_timeout=10,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Managed mode: send auth message and wait for ready
|
||
|
|
if self.mode == "managed":
|
||
|
|
if not await self._managed_handshake():
|
||
|
|
return
|
||
|
|
|
||
|
|
# Run send and receive concurrently
|
||
|
|
await asyncio.gather(
|
||
|
|
self._send_loop(),
|
||
|
|
self._receive_loop(),
|
||
|
|
)
|
||
|
|
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
pass
|
||
|
|
except Exception as exc:
|
||
|
|
msg = f"WebSocket error: {exc}"
|
||
|
|
logger.error(msg)
|
||
|
|
if self._on_error:
|
||
|
|
self._on_error(msg)
|
||
|
|
finally:
|
||
|
|
await self._close_ws()
|
||
|
|
|
||
|
|
def _build_ws_url_and_headers(self):
|
||
|
|
"""Return ``(url, headers)`` depending on the current mode."""
|
||
|
|
if self.mode == "managed":
|
||
|
|
# Ensure the server URL uses wss:// and append the path
|
||
|
|
url = self.server_url.rstrip("/")
|
||
|
|
if not url.startswith("ws://") and not url.startswith("wss://"):
|
||
|
|
url = f"wss://{url}"
|
||
|
|
url = f"{url}/ws/transcribe"
|
||
|
|
return url, {}
|
||
|
|
|
||
|
|
# BYOK -- connect directly to Deepgram
|
||
|
|
params = (
|
||
|
|
f"model={self.deepgram_model}"
|
||
|
|
f"&language={self.language}"
|
||
|
|
"&interim_results=true"
|
||
|
|
"&encoding=linear16"
|
||
|
|
f"&sample_rate={self.sample_rate}"
|
||
|
|
f"&channels={self.channels}"
|
||
|
|
)
|
||
|
|
url = f"wss://api.deepgram.com/v1/listen?{params}"
|
||
|
|
headers = {"Authorization": f"Token {self.byok_api_key}"}
|
||
|
|
return url, headers
|
||
|
|
|
||
|
|
# -- managed-mode handshake ---------------------------------------- #
|
||
|
|
|
||
|
|
async def _managed_handshake(self) -> bool:
|
||
|
|
"""Send auth message and wait for ``ready`` (managed mode).
|
||
|
|
|
||
|
|
Returns ``True`` on success.
|
||
|
|
"""
|
||
|
|
auth_msg = {
|
||
|
|
"type": "auth",
|
||
|
|
"token": self.auth_token,
|
||
|
|
"config": {
|
||
|
|
"model": self.deepgram_model,
|
||
|
|
"language": self.language,
|
||
|
|
"sample_rate": self.sample_rate,
|
||
|
|
"channels": self.channels,
|
||
|
|
"encoding": "linear16",
|
||
|
|
"interim_results": True,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
await self._ws.send(json.dumps(auth_msg))
|
||
|
|
|
||
|
|
try:
|
||
|
|
raw = await asyncio.wait_for(self._ws.recv(), timeout=15)
|
||
|
|
data = json.loads(raw)
|
||
|
|
if data.get("type") == "ready":
|
||
|
|
logger.info("Managed proxy is ready")
|
||
|
|
return True
|
||
|
|
|
||
|
|
if data.get("type") == "error":
|
||
|
|
err = data.get("message", "unknown error")
|
||
|
|
logger.error("Auth error from proxy: %s", err)
|
||
|
|
if self._on_error:
|
||
|
|
self._on_error(f"Proxy auth error: {err}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
logger.warning("Unexpected handshake message: %s", data)
|
||
|
|
return False
|
||
|
|
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
logger.error("Timed out waiting for proxy ready message")
|
||
|
|
if self._on_error:
|
||
|
|
self._on_error("Timed out waiting for proxy ready message")
|
||
|
|
return False
|
||
|
|
|
||
|
|
# -- send loop ----------------------------------------------------- #
|
||
|
|
|
||
|
|
async def _send_loop(self):
|
||
|
|
"""Drain the audio queue and push raw PCM bytes over the WebSocket."""
|
||
|
|
while not self._stop_event.is_set():
|
||
|
|
try:
|
||
|
|
pcm_bytes = self._audio_queue.get(timeout=0.1)
|
||
|
|
except Empty:
|
||
|
|
continue
|
||
|
|
|
||
|
|
try:
|
||
|
|
await self._ws.send(pcm_bytes)
|
||
|
|
except Exception as exc:
|
||
|
|
if not self._stop_event.is_set():
|
||
|
|
logger.error("Send error: %s", exc)
|
||
|
|
break
|
||
|
|
|
||
|
|
# -- receive loop -------------------------------------------------- #
|
||
|
|
|
||
|
|
async def _receive_loop(self):
|
||
|
|
"""Listen for messages from the WebSocket and dispatch them."""
|
||
|
|
while not self._stop_event.is_set():
|
||
|
|
try:
|
||
|
|
raw = await asyncio.wait_for(self._ws.recv(), timeout=1.0)
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
continue
|
||
|
|
except Exception as exc:
|
||
|
|
if not self._stop_event.is_set():
|
||
|
|
logger.error("Receive error: %s", exc)
|
||
|
|
break
|
||
|
|
|
||
|
|
try:
|
||
|
|
data = json.loads(raw)
|
||
|
|
except (json.JSONDecodeError, TypeError):
|
||
|
|
logger.debug("Non-JSON message received, ignoring")
|
||
|
|
continue
|
||
|
|
|
||
|
|
if self.mode == "managed":
|
||
|
|
self._handle_managed_message(data)
|
||
|
|
else:
|
||
|
|
self._handle_byok_message(data)
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
# Message handlers
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
|
||
|
|
def _handle_managed_message(self, data: dict):
|
||
|
|
"""Process a message from the managed proxy."""
|
||
|
|
msg_type = data.get("type", "")
|
||
|
|
|
||
|
|
if msg_type == "transcript":
|
||
|
|
text = data.get("text", "")
|
||
|
|
is_final = data.get("is_final", False)
|
||
|
|
if text.strip():
|
||
|
|
result = TranscriptionResult(
|
||
|
|
text=text,
|
||
|
|
is_final=is_final,
|
||
|
|
timestamp=datetime.now(),
|
||
|
|
user_name=self.user_name,
|
||
|
|
)
|
||
|
|
if is_final:
|
||
|
|
if self.final_callback:
|
||
|
|
self.final_callback(result)
|
||
|
|
else:
|
||
|
|
if self.realtime_callback:
|
||
|
|
self.realtime_callback(result)
|
||
|
|
|
||
|
|
elif msg_type == "credits_low":
|
||
|
|
seconds_remaining = data.get("seconds_remaining", 0)
|
||
|
|
logger.warning("Credits low -- %d seconds remaining", seconds_remaining)
|
||
|
|
if self._on_credits_low:
|
||
|
|
self._on_credits_low(int(seconds_remaining))
|
||
|
|
|
||
|
|
elif msg_type == "error":
|
||
|
|
code = data.get("code", "")
|
||
|
|
message = data.get("message", "Unknown error")
|
||
|
|
logger.error("Proxy error [%s]: %s", code, message)
|
||
|
|
if self._on_error:
|
||
|
|
self._on_error(f"[{code}] {message}" if code else message)
|
||
|
|
|
||
|
|
elif msg_type == "session_end":
|
||
|
|
seconds_used = data.get("seconds_used", 0)
|
||
|
|
logger.info("Session ended -- %d seconds used", seconds_used)
|
||
|
|
|
||
|
|
elif msg_type == "ready":
|
||
|
|
# May arrive again after reconnects; safe to ignore.
|
||
|
|
logger.debug("Received ready message (already connected)")
|
||
|
|
|
||
|
|
else:
|
||
|
|
logger.debug("Unhandled managed message type: %s", msg_type)
|
||
|
|
|
||
|
|
def _handle_byok_message(self, data: dict):
|
||
|
|
"""Process a message received directly from the Deepgram API."""
|
||
|
|
msg_type = data.get("type", "")
|
||
|
|
|
||
|
|
if msg_type == "Results":
|
||
|
|
channel = data.get("channel", {})
|
||
|
|
alternatives = channel.get("alternatives", [])
|
||
|
|
if not alternatives:
|
||
|
|
return
|
||
|
|
|
||
|
|
transcript = alternatives[0].get("transcript", "")
|
||
|
|
is_final = data.get("is_final", False)
|
||
|
|
|
||
|
|
if transcript.strip():
|
||
|
|
result = TranscriptionResult(
|
||
|
|
text=transcript,
|
||
|
|
is_final=is_final,
|
||
|
|
timestamp=datetime.now(),
|
||
|
|
user_name=self.user_name,
|
||
|
|
)
|
||
|
|
if is_final:
|
||
|
|
if self.final_callback:
|
||
|
|
self.final_callback(result)
|
||
|
|
else:
|
||
|
|
if self.realtime_callback:
|
||
|
|
self.realtime_callback(result)
|
||
|
|
|
||
|
|
elif msg_type == "Metadata":
|
||
|
|
logger.debug("Deepgram metadata: %s", data)
|
||
|
|
|
||
|
|
elif msg_type == "UtteranceEnd":
|
||
|
|
logger.debug("Deepgram utterance end")
|
||
|
|
|
||
|
|
else:
|
||
|
|
logger.debug("Unhandled Deepgram message type: %s", msg_type)
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
# Helpers
|
||
|
|
# ------------------------------------------------------------------ #
|
||
|
|
|
||
|
|
async def _close_ws(self):
|
||
|
|
"""Close the WebSocket connection if open."""
|
||
|
|
if self._ws is not None:
|
||
|
|
try:
|
||
|
|
await self._ws.close()
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
self._ws = None
|
||
|
|
|
||
|
|
def set_user_name(self, user_name: str):
|
||
|
|
"""Update the user name attached to future transcriptions."""
|
||
|
|
self.user_name = user_name
|
||
|
|
|
||
|
|
def is_recording_active(self) -> bool:
|
||
|
|
"""Return ``True`` if audio is currently being captured."""
|
||
|
|
return self._is_recording
|
||
|
|
|
||
|
|
def __repr__(self) -> str:
|
||
|
|
return (
|
||
|
|
f"DeepgramTranscriptionEngine(mode={self.mode}, "
|
||
|
|
f"recording={self._is_recording})"
|
||
|
|
)
|
||
|
|
|
||
|
|
def __del__(self):
|
||
|
|
"""Best-effort cleanup."""
|
||
|
|
try:
|
||
|
|
self.stop()
|
||
|
|
except Exception:
|
||
|
|
pass
|