stt_streamers/soniox_streamer.py

import json
import os
import threading
import time
import numpy as np
from websockets.sync.client import connect
from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError

SONIOX_WEBSOCKET_URL = "wss://stt-rt.soniox.com/transcribe-websocket"


def _connect_soniox():
    # Bypass proxy auto-detection on hosted environments and rely on Soniox keepalive.
    return connect(
        SONIOX_WEBSOCKET_URL,
        proxy=None,
        compression=None,
        ping_interval=None,
        ping_timeout=None,
        open_timeout=20,
        close_timeout=5,
        max_size=None,
    )

class SonioxStreamer:
    def __init__(self, fs_hz: int, stream_name: str, on_update=None) -> None:
        api_key = os.environ.get("SONIOX_API_KEY")
        if not api_key:
            raise RuntimeError("Missing SONIOX_API_KEY.")
        if not fs_hz:
            raise ValueError("Sample rate (fs_hz) must be specified.")
        self.stream_name = stream_name
        self.api_name = "Soniox RT"
        self.on_update = on_update

        self.final_tokens: list[dict] = []
        self.lock = threading.Lock()
        self.finished_event = threading.Event()

        self._send_lock = threading.Lock()
        self._stop_evt = threading.Event()
        self._last_send_ts = time.monotonic()

        self.ws = _connect_soniox()

        config = self.get_config(api_key, fs_hz)
        with self._send_lock:
            self.ws.send(json.dumps(config))
        self._last_send_ts = time.monotonic()

        self.thread = threading.Thread(target=self._receive_loop, daemon=True)
        self.thread.start()

        self.keepalive_thread = threading.Thread(target=self._keepalive_loop, daemon=True)
        self.keepalive_thread.start()

        print(f"SonioxStreamer '{self.stream_name}' initialized and connected.")

    def stream_array(self, pcm: np.ndarray) -> str:
        chunk_size = 160
        num_chunks = int(np.ceil(len(pcm) / chunk_size))
        print(f"Streaming {self.stream_name} audio to Soniox...")
        for i in range(num_chunks):
            start_idx = i * chunk_size
            end_idx = min((i + 1) * chunk_size, len(pcm))
            chunk = pcm[start_idx:end_idx]
            self.process_chunk(chunk)
        self.close_stream()

        # 3. Wait for the 'finished' message from the receive loop
        self.finished_event.wait()
        print(f"{self.stream_name} streaming complete.")
        transcript = self.get_final_transcript()
        self.shutdown()
        return transcript
        
    
    def get_config(self, api_key: str, fs_hz: int) -> dict:
        return {
            "api_key": api_key,
            "model": "stt-rt-v3",
            "audio_format": "pcm_s16le",
            "sample_rate": fs_hz,
            "num_channels": 1,
            "language_hints": ["en", "de"],
            "language_hints_strict": True,
            "enable_language_identification": True,
            "enable_speaker_diarization": False,
            "enable_endpoint_detection": True,  # liefert <end> am Segmentende citeturn1view0
        }

    def clear_text(self):
        with self.lock:
            self.final_tokens = []

    def process_chunk(self, chunk: np.ndarray) -> None:
        if self.ws is None or self._stop_evt.is_set():
            return

        chunk = np.clip(chunk, -1.0, 1.0)
        chunk_int16 = (chunk * 32767).astype(np.int16)
        if chunk_int16.size == 0:
            return

        try:
            with self._send_lock:
                self.ws.send(chunk_int16.tobytes())
            self._last_send_ts = time.monotonic()
        except Exception as e:
            print(f"[Soniox:{self.stream_name}] send failed: {e}")
            self._stop_evt.set()
            self.ws = None

    def render_tokens(self, final_tokens: list[dict], non_final_tokens: list[dict]) -> str:
        parts = []
        for tok in final_tokens + non_final_tokens:
            text = tok.get("text", "")
            if text == "<end>":
                parts.append("\n")
                continue
            if text == "<fin>":
                continue
            parts.append(text)

            if text.strip() in [".", "?", "!"]:
                parts.append("\n")
        return "".join(parts)

    def get_final_transcript(self) -> str:
        with self.lock:
            return self.render_tokens(self.final_tokens, [])

    def _receive_loop(self):
        try:
            while True:
                message = self.ws.recv()
                if isinstance(message, (bytes, bytearray)):
                    message = message.decode("utf-8", errors="replace")
                res = json.loads(message)

                # Soniox: Error response => server schließt sofort citeturn6view0
                if res.get("error_code") is not None:
                    print(f"[Soniox:{self.stream_name}] ERROR {res.get('error_code')}: {res.get('error_message')}")
                    self.ws = None
                    break

                non_final_tokens: list[dict] = []

                with self.lock:
                    for token in res.get("tokens", []):
                        text = token.get("text")
                        if not text:
                            continue

                        # Final tokens werden nur einmal gesendet (nicht wiederholt) citeturn5view0
                        if token.get("is_final"):
                            self.final_tokens.append(token)
                        else:
                            non_final_tokens.append(token)

                    current_finals = list(self.final_tokens)

                text = self.render_tokens(current_finals, non_final_tokens)
                if self.on_update:
                    self.on_update(text)

                if res.get("finished"):
                    self.finished_event.set()
                    break

        except (ConnectionClosedOK, ConnectionClosedError):
            pass
        except Exception as e:
            print(f"[Soniox:{self.stream_name}] receive loop error: {e}")
        finally:
            self.finished_event.set()

    def _keepalive_loop(self):
        KEEPALIVE_INTERVAL_S = 10.0  # Soniox: mind. 1×/20s, 5–10s üblich citeturn2view0
        while not self._stop_evt.is_set():
            time.sleep(0.5)
            if self.ws is None:
                return

            if time.monotonic() - self._last_send_ts >= KEEPALIVE_INTERVAL_S:
                try:
                    with self._send_lock:
                        self.ws.send(json.dumps({"type": "keepalive"}))  # Soniox-format citeturn2view0
                    self._last_send_ts = time.monotonic()
                except Exception as e:
                    # Normal closure (1000 OK) when stream ends; don't log as failure
                    if "1000" not in str(e):
                        print(f"[Soniox:{self.stream_name}] keepalive failed: {e}")
                    self.ws = None
                    return

    def close_stream(self):
        # Soniox: End-of-stream = leeres WebSocket frame (binary oder text) citeturn6view0
        if self.ws is None:
            return
        try:
            with self._send_lock:
                self.ws.send("")
            self._last_send_ts = time.monotonic()
        except Exception:
            pass

    def shutdown(self) -> None:
        self._stop_evt.set()

        # Graceful end-of-stream
        self.close_stream()

        # optional: kurz warten auf finished (Server sendet finished response) citeturn6view0
        self.finished_event.wait(timeout=2.0)

        self._close_socket()

        if hasattr(self, "thread") and self.thread.is_alive() and self.thread != threading.current_thread():
            self.thread.join(timeout=1.0)
        if hasattr(self, "keepalive_thread") and self.keepalive_thread.is_alive():
            self.keepalive_thread.join(timeout=1.0)
        print(f"SonioxStreamer '{self.stream_name}' shutdown complete.")

    def _close_socket(self) -> None:
        if self.ws is None:
            return
        try:
            self.ws.close()
        except Exception:
            pass
        finally:
            self.ws = None