feat(xai): configurable TTS output sample_rate

livekit-plugins/livekit-plugins-xai/livekit/plugins/xai/tts.py

@@ -41,6 +41,7 @@
 
 SAMPLE_RATE = 24000
 NUM_CHANNELS = 1
+_SUPPORTED_SAMPLE_RATES = (8000, 16000, 22050, 24000, 44100, 48000)
 
 XAI_WEBSOCKET_URL = "wss://api.x.ai/v1/tts"
 DEFAULT_VOICE = "ara"
@@ -54,6 +55,7 @@ class _TTSOptions:
     optimize_streaming_latency: NotGivenOr[int]
     speed: NotGivenOr[float]
     text_normalization: NotGivenOr[bool]
+    sample_rate: int
 
 
 class TTS(tts.TTS):
@@ -66,6 +68,7 @@ def __init__(
         optimize_streaming_latency: NotGivenOr[int] = NOT_GIVEN,
         speed: NotGivenOr[float] = NOT_GIVEN,
         text_normalization: NotGivenOr[bool] = NOT_GIVEN,
+        sample_rate: int = SAMPLE_RATE,
         tokenizer: tokenize.WordTokenizer | None = None,
         http_session: aiohttp.ClientSession | None = None,
     ) -> None:
@@ -80,12 +83,17 @@ def __init__(
             optimize_streaming_latency (int, optional): Latency optimization level for the xAI TTS websocket.
             speed (float, optional): Speaking-rate multiplier for the generated audio.
             text_normalization (bool, optional): Whether to normalize text before synthesis.
+            sample_rate (int, optional): Output audio sample rate in Hz. One of 8000, 16000, 22050, 24000, 44100, 48000. Defaults to 24000. Match this to your pipeline (e.g. 16000) to avoid an extra resampling stage.
             api_key (str | None, optional): The xAI API key. If not provided, it will be read from the xAI environment variable.
             http_session (aiohttp.ClientSession | None, optional): An existing aiohttp ClientSession to use. If not provided, a new session will be created.
         """  # noqa: E501
+        if sample_rate not in _SUPPORTED_SAMPLE_RATES:
+            raise ValueError(
+                f"sample_rate must be one of {_SUPPORTED_SAMPLE_RATES}, got {sample_rate}"
+            )
         super().__init__(
             capabilities=tts.TTSCapabilities(streaming=True),
-            sample_rate=SAMPLE_RATE,
+            sample_rate=sample_rate,
             num_channels=NUM_CHANNELS,
         )
 
@@ -105,6 +113,7 @@ def __init__(
             optimize_streaming_latency=optimize_streaming_latency,
             speed=speed,
             text_normalization=text_normalization,
+            sample_rate=sample_rate,
         )
 
         self._session = http_session
@@ -138,7 +147,7 @@ async def _connect_ws(
             "voice": opts.voice,
             "language": opts.language,
             "codec": "pcm",
-            "sample_rate": SAMPLE_RATE,
+            "sample_rate": opts.sample_rate,
         }
         if is_given(opts.optimize_streaming_latency):
             params["optimize_streaming_latency"] = opts.optimize_streaming_latency
@@ -266,7 +275,7 @@ async def _run(self, output_emitter: tts.AudioEmitter) -> None:
         request_id = utils.shortuuid()
         output_emitter.initialize(
             request_id=request_id,
-            sample_rate=SAMPLE_RATE,
+            sample_rate=self._opts.sample_rate,
             num_channels=NUM_CHANNELS,
             stream=True,
             mime_type="audio/pcm",

tests/test_plugin_xai_tts.py

@@ -118,3 +118,34 @@ async def close_ws(ws: _FakeWebSocket) -> None:
     await tts.aclose()
 
     assert websockets[0].closed is True
+
+
+def test_sample_rate_defaults_to_24000() -> None:
+    tts = xai_tts.TTS(api_key="test-key")
+    assert tts.sample_rate == 24000
+    assert tts._opts.sample_rate == 24000  # pyright: ignore[reportPrivateUsage]
+
+
+def test_invalid_sample_rate_raises() -> None:
+    with pytest.raises(ValueError):
+        xai_tts.TTS(api_key="test-key", sample_rate=12345)
+
+
+@pytest.mark.asyncio
+async def test_sample_rate_is_configurable(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured: dict[str, str] = {}
+
+    class _FakeSession:
+        async def ws_connect(self, url: str, **_kwargs: object) -> _FakeWebSocket:
+            captured["url"] = url
+            return _FakeWebSocket()
+
+    tts = xai_tts.TTS(api_key="test-key", sample_rate=16000)
+    assert tts.sample_rate == 16000
+    assert tts._opts.sample_rate == 16000  # pyright: ignore[reportPrivateUsage]
+
+    monkeypatch.setattr(tts, "_ensure_session", lambda: _FakeSession())
+    await tts._connect_ws(1.0, tts._opts)  # pyright: ignore[reportPrivateUsage]
+    assert "sample_rate=16000" in captured["url"]
+
+    await tts.aclose()