stt_streamers/deepgram_streamer.py

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

DEEPGRAM_WEBSOCKET_URL = "wss://api.deepgram.com/v1/listen"


def _connect_deepgram(url_with_params: str, headers: dict[str, str]):
    return connect(
        url_with_params,
        additional_headers=headers,
        proxy=None,
        compression=None,
        ping_interval=None,
        ping_timeout=None,
        open_timeout=20,
        close_timeout=5,
        max_size=None,
    )


class DeepgramStreamer:
    def __init__(self, fs_hz: int, stream_name: str, on_update=None) -> None:
        api_key = os.environ.get("DEEPGRAM_API_KEY")
        if not api_key:
            raise RuntimeError("Missing DEEPGRAM_API_KEY.")
        if not fs_hz:
            raise ValueError("Sample rate (fs_hz) must be specified.")
        self.stream_name = stream_name
        self.api_name = "Deepgram V1 Nova-3"
        self.on_update = on_update
        self.final_tokens: list[dict] = []
        self.lock = threading.Lock()
        self.finished_event = threading.Event()

        # 1. Build the Deepgram URL with query parameters
        config = self.get_config(fs_hz)
        query_string = urllib.parse.urlencode(config)
        url_with_params = f"{DEEPGRAM_WEBSOCKET_URL}?{query_string}"

        # 2. Connect with Authorization header
        # Deepgram requires the API key in the headers
        headers = {"Authorization": f"Token {api_key}"}
        self.ws = _connect_deepgram(url_with_params, headers)
        self._send_lock = threading.Lock()
        self._stop_evt = threading.Event()
        self._last_send_ts = time.monotonic()
        
        # 3. Start the receiving thread
        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"DeepgramStreamer '{self.stream_name}' initialized and connected.")

    def stream_array(self, pcm: np.ndarray) -> str:
        """
        Streams audio chunks to Deepgram and waits for the final result.
        """
        chunk_size = 160  # Keeping the same chunk size as the reference
        num_chunks = int(np.ceil(len(pcm) / chunk_size))
        print(f"Streaming {self.stream_name} audio to Deepgram...")

        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)
        print(f"Finished streaming {self.stream_name} audio.")
        # Signal the end of the stream to Deepgram.
        self.close_stream()

        # Wait for the 'finished' signal (Metadata) from the receive loop
        self.finished_event.wait()

        transcript = self.get_final_transcript()
        self.shutdown()
        return transcript

    def clear_text(self):
        with self.lock:
            self.final_tokens = []
    
    def get_config(self, fs_hz: int) -> dict:
        """
        Returns parameters for the Deepgram V1 URL query string.
        """

        return {
            "model": "nova-3",  # Recommended general model
            "encoding": "linear16",  # Corresponds to pcm_s16le
            "sample_rate": fs_hz,
            "channels": 1,
            "smart_format": "true",  # handling punctuation/formatting
            "interim_results": "true",  # required for non-final updates
            "endpointing": "500",  # ms silence to trigger finalization
            "language": "multi",
            
        }

    def process_chunk(self, chunk: np.ndarray) -> None:
        if self._stop_evt.is_set() or self.ws is None:
            return
        chunk = np.clip(chunk, -1.0, 1.0)
        chunk_int16 = (chunk * 32767).astype(np.int16)
        if len(chunk_int16) == 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"[{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:
        """
        Renders the list of token dicts into a string.
        Matches Soniox logic: treats certain tokens as punctuation triggers.
        """
        text_parts = []
        for token in final_tokens + non_final_tokens:
            text = token["text"]
            text_parts.append(text)
            # Add newline if the text chunk looks like end-of-sentence punctuation
            # Note: Deepgram 'smart_format' usually attaches punctuation to the word.
            if text.strip() in [".", "?", "!"]:
                text_parts.append("\n")
        return "".join(text_parts)

    def _receive_loop(self):
        """
        Background loop to handle incoming JSON messages from Deepgram.
        """
        try:
            while True:
                message = self.ws.recv()
                res = json.loads(message)

                # Check for metadata indicating stream end
                if res.get("type") == "Metadata":
                    self.finished_event.set()
                    break

                # Deepgram error handling
                if "error" in res:
                    print(f"Deepgram Error: {res['error']}")
                    break

                # Process Transcripts
                # Deepgram V1 structure: result -> channel -> alternatives -> [0] -> transcript
                if "channel" in res:
                    is_final = res.get("is_final", False)
                    alternatives = res["channel"].get("alternatives", [])

                    if alternatives:
                        transcript = alternatives[0].get("transcript", "")

                        if transcript:
                            # Wrap the transcript in a dict to match the
                            # 'render_tokens' expectation of a list[dict]
                            token_data = {
                                "text": transcript + " ",  # Add space for readability
                                "is_final": is_final,
                            }

                            non_final_tokens = []
                            with self.lock:
                                if is_final:
                                    self.final_tokens.append(token_data)
                                else:
                                    non_final_tokens.append(token_data)

                                current_finals = list(self.final_tokens)

                            # Trigger the callback
                            text = self.render_tokens(current_finals, non_final_tokens)
                            if self.on_update:
                                self.on_update(text)

        except (ConnectionClosedOK, ConnectionClosedError):
            pass
        except Exception as e:
            print(f"Deepgram receive loop error: {e}")
        finally:
            self.finished_event.set()
            
    def _keepalive_loop(self):
        # Deepgram: KeepAlive als Text-Message senden citeturn4search25
        while not self._stop_evt.is_set():
            time.sleep(0.5)
            if self.ws is None:
                return
            if time.monotonic() - self._last_send_ts >= 3.0:
                try:
                    with self._send_lock:
                        self.ws.send(json.dumps({"type": "KeepAlive"}))
                    self._last_send_ts = time.monotonic()
                except Exception:
                    # bei Fehler: loop endet, send wird später reconnecten können
                    self.ws = None
                    return

    def close_stream(self) -> None:
        # Deepgram CloseStream: {"type":"CloseStream"} citeturn2view3
        if self.ws is None:
            return
        try:
            with self._send_lock:
                self.ws.send(json.dumps({"type": "CloseStream"}))
        except Exception:
            pass

    def close(self) -> None:
        self.close_stream()

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

    def _close_socket(self) -> None:
        if self.ws is None:
            return
        try:
            self.ws.close()
        except Exception:
            pass
        finally:
            self.ws = None
    
    def shutdown(self) -> None:
        self._stop_evt.set()
        self.close_stream()
        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"DeepgramStreamer '{self.stream_name}' shutdown complete.")