feat(deepgram): make Flux audio chunk duration configurable

livekit-plugins/livekit-plugins-deepgram/livekit/plugins/deepgram/stt_v2.py

@@ -54,6 +54,7 @@ class STTOptions:
     sample_rate: int
     keyterm: str | Sequence[str]
     endpoint_url: str
+    chunk_size_ms: int = 50
     language: str = "en"
     eager_eot_threshold: NotGivenOr[float] = NOT_GIVEN
     eot_threshold: NotGivenOr[float] = NOT_GIVEN
@@ -72,6 +73,7 @@ def __init__(
         *,
         model: V2Models | str = "flux-general-en",
         sample_rate: int = 16000,
+        chunk_size_ms: int = 50,
         eager_eot_threshold: NotGivenOr[float] = NOT_GIVEN,
         eot_threshold: NotGivenOr[float] = NOT_GIVEN,
         eot_timeout_ms: NotGivenOr[int] = NOT_GIVEN,
@@ -93,6 +95,10 @@ def __init__(
         Args:
             model: The Deepgram model to use for speech recognition. Defaults to "flux-general-en".
             sample_rate: The sample rate of the audio in Hz. Defaults to 16000.
+            chunk_size_ms: Duration of outgoing PCM audio chunks in milliseconds. Defaults to 50.
+                Deepgram recommends 80 ms for Flux. This controls WebSocket audio batching,
+                independently of incoming RTC frame size and endpointing delay. Set at creation time.
+                See https://developers.deepgram.com/docs/flux/quickstart.
             eager_eot_threshold: The threshold for eager end of turn to enable preemptive generation. Disabled by default. Set to 0.3-0.9 to enable preemptive generation.
             eot_threshold: The threshold for end of speech detection, ranges 0.5-0.9. Defaults to 0.7. If using eager_eot_threshold, set this higher to allow a higher eager value.
             eot_timeout_ms: The timeout for end of speech detection. Defaults to 3000.
@@ -108,7 +114,8 @@ def __init__(
             redact: Redact numbers from the transcription, "numbers" or "aggressive_numbers". Flux does not support entity redaction (pci, pii, ...). Applied at connection time. Defaults to NOT_GIVEN.
 
         Raises:
-            ValueError: If no API key is provided or found in environment variables.
+            ValueError: If no API key is provided or found in environment variables, or
+                chunk_size_ms is not a positive integer that yields at least one audio sample.
 
         Note:
             The api_key must be set either through the constructor argument or by setting
@@ -125,6 +132,16 @@ def __init__(
             )
         )
 
+        if (
+            isinstance(chunk_size_ms, bool)
+            or not isinstance(chunk_size_ms, int)
+            or chunk_size_ms <= 0
+            or sample_rate * chunk_size_ms // 1000 < 1
+        ):
+            raise ValueError(
+                "chunk_size_ms must be a positive integer yielding at least one sample"
+            )
+
         deepgram_api_key = api_key if is_given(api_key) else os.environ.get("DEEPGRAM_API_KEY")
         if not deepgram_api_key:
             raise ValueError("Deepgram API key is required")
@@ -152,6 +169,7 @@ def __init__(
         self._opts = STTOptions(
             model=model,
             sample_rate=sample_rate,
+            chunk_size_ms=chunk_size_ms,
             keyterm=([keyterm] if isinstance(keyterm, str) else list(keyterm))
             if is_given(keyterm)
             else [],
@@ -462,12 +480,12 @@ async def _run(self) -> None:
         async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
             nonlocal closing_ws
 
-            # forward audio to deepgram in chunks of 50ms
-            samples_50ms = self._opts.sample_rate // 20
+            # Repack incoming frames into the configured outgoing PCM chunk duration.
+            chunk_samples = self._opts.sample_rate * self._opts.chunk_size_ms // 1000
             audio_bstream = utils.audio.AudioByteStream(
                 sample_rate=self._opts.sample_rate,
                 num_channels=1,
-                samples_per_channel=samples_50ms,
+                samples_per_channel=chunk_samples,
             )
 
             has_ended = False

tests/test_plugin_deepgram_stt.py

@@ -291,3 +291,84 @@ async def test_flush_finalizes_after_the_buffered_audio():
         await _wait_until(lambda: ws.sent() == ["audio", "Finalize"])
     finally:
         await stream.aclose()
+
+
+@pytest.mark.parametrize("chunk_size_ms", [0, -1, 1.5, True])
+def test_flux_rejects_invalid_chunk_size(chunk_size_ms):
+    from livekit.plugins.deepgram import STTv2
+
+    with pytest.raises(ValueError, match="chunk_size_ms"):
+        STTv2(api_key="test-key", chunk_size_ms=chunk_size_ms)
+
+
+def test_flux_rejects_chunk_size_smaller_than_one_sample():
+    from livekit.plugins.deepgram import STTv2
+
+    with pytest.raises(ValueError, match="chunk_size_ms"):
+        STTv2(api_key="test-key", sample_rate=100, chunk_size_ms=1)
+
+
+@pytest.mark.parametrize("chunk_size_ms", [None, 20, 80])
+@pytest.mark.parametrize("sample_rate", [8000, 16000, 48000])
+@pytest.mark.parametrize("duration_ms", [400, 410])
+async def test_flux_chunk_size_preserves_audio_and_flushes_tail(
+    chunk_size_ms, sample_rate, duration_ms
+):
+    from typing import Any, cast
+
+    from livekit import rtc
+    from livekit.plugins.deepgram import STTv2
+
+    class AudioWS(_LiveWS):
+        def __init__(self) -> None:
+            super().__init__()
+            self.audio: list[bytes] = []
+
+        async def send_bytes(self, data: bytes) -> None:
+            self.audio.append(data)
+            await super().send_bytes(data)
+
+    ws = AudioWS()
+    # Omitting the argument exercises the public constructor's existing default.
+    options = {} if chunk_size_ms is None else {"chunk_size_ms": chunk_size_ms}
+    instance = STTv2(
+        api_key="test-key",
+        sample_rate=sample_rate,
+        http_session=cast(Any, SimpleNamespace(closed=False)),
+        **options,
+    )
+    stream = instance.stream()
+
+    async def connect():
+        return ws
+
+    # Replace only the network connection; exercise the real stream and send loop.
+    stream._connect_ws = connect
+    try:
+        assert "chunk_size_ms" not in stream._live_config()
+        input_samples = sample_rate // 100  # 10 ms RTC frames
+        input_bytes = input_samples * 2
+        pcm = bytes(i % 256 for i in range(sample_rate * duration_ms // 1000 * 2))
+        for offset in range(0, len(pcm), input_bytes):
+            stream.push_frame(
+                rtc.AudioFrame(
+                    data=pcm[offset : offset + input_bytes],
+                    sample_rate=sample_rate,
+                    num_channels=1,
+                    samples_per_channel=input_samples,
+                )
+            )
+        stream.end_input()
+        await _wait_until(lambda: "CloseStream" in ws.wire)
+
+        chunk_bytes = sample_rate * (chunk_size_ms or 50) // 1000 * 2
+        full_chunks, remainder = divmod(len(pcm), chunk_bytes)
+        expected_sizes = [chunk_bytes] * full_chunks
+        if remainder:
+            expected_sizes.append(remainder)
+        assert [len(data) for data in ws.audio] == expected_sizes
+        assert b"".join(ws.audio) == pcm
+        assert ws.sent() == ["audio"] * len(expected_sizes) + ["CloseStream"]
+    finally:
+        await stream.aclose()
+        await instance.aclose()