feat: sglang backend support

1 files changed, 364 insertions, 0 deletions

diff --git a/app/tests/test_tts_sglomni.py b/app/tests/test_tts_sglomni.py
new file mode 100644
index 0000000..2dee364
--- /dev/null
+++ b/app/tests/test_tts_sglomni.py
@@ -0,0 +1,364 @@
+"""Tests for the SGLang-Omni TTS client (converter/clients/sglomni.py)."""
+
+import base64
+import io
+import json
+import tempfile
+import unittest
+import urllib.error
+import wave
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from converter import config
+from converter.clients import SgOmniTTSClient
+from converter.clients.base import NonRetryableTTSError
+from converter.clients.sglomni import _data_url, _is_loopback
+from converter.clients.speakers import QWEN3_TTS_SPEAKERS
+
+
+class CatalogConsistencyTests(unittest.TestCase):
+    """The backend catalog's vendored facts match the converter's."""
+
+    def test_customvoice_speakers_match_the_qwen_table(self):
+        from backends.sglomni.catalog import QWEN_CUSTOMVOICE_SPEAKERS
+        self.assertEqual(QWEN_CUSTOMVOICE_SPEAKERS, QWEN3_TTS_SPEAKERS)
+
+_DUMMY_CHUNKS = Path(tempfile.gettempdir()) / "audiobook_sglomni_test_chunks"
+
+
+def _make_wav() -> bytes:
+    """A real minimal RIFF/WAVE file (what a server response looks like)."""
+    buffer = io.BytesIO()
+    with wave.open(buffer, "wb") as wav_file:
+        wav_file.setnchannels(1)
+        wav_file.setsampwidth(2)
+        wav_file.setframerate(24000)
+        wav_file.writeframes(b"\x01\x00" * 16)
+    return buffer.getvalue()
+
+
+_WAV_BYTES = _make_wav()
+_WAV_FRAMES = b"\x01\x00" * 16
+
+
+class LoopbackTests(unittest.TestCase):
+    def test_loopback_hosts(self):
+        self.assertTrue(_is_loopback("http://127.0.0.1:8100"))
+        self.assertTrue(_is_loopback("http://localhost:8100"))
+        self.assertFalse(_is_loopback("http://10.20.30.40:8100"))
+
+    def test_data_url_carries_mime_and_bytes(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "ref.wav"
+            path.write_bytes(b"abc")
+            url = _data_url(path)
+        self.assertTrue(url.startswith("data:audio/wav;base64,"))
+        self.assertEqual(
+            base64.b64decode(url.partition(";base64,")[2]), b"abc")
+
+
+class ConnectInputTests(unittest.TestCase):
+    """Capability-driven validation before any HTTP is attempted."""
+
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.ref = Path(self._tmp.name) / "narrator.wav"
+        self.ref.write_bytes(b"abc")
+        self.addCleanup(self._tmp.cleanup)
+
+    def _client(self, model="higgs_audio_v3_tts", **kwargs):
+        # Bypass _connect (HTTP) — these tests cover the input checks.
+        with patch.object(SgOmniTTSClient, "_connect"):
+            return SgOmniTTSClient(_DUMMY_CHUNKS, model=model, **kwargs)
+
+    def test_unknown_model_raises(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            self._client(model="nope")
+        self.assertIn("Unknown SGLang-Omni model", str(ctx.exception))
+
+    def test_design_model_requires_instructions(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            self._client(model="qwen3_tts_1_7b_voicedesign")
+        self.assertIn("--instructions", str(ctx.exception))
+
+    def test_reference_required_model_refuses_to_connect_without_one(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            self._client(model="qwen3_tts_1_7b_base")
+        self.assertIn("requires reference audio", str(ctx.exception))
+
+    def test_missing_reference_file_raises(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            self._client(model="higgs_audio_v3_tts",
+                         ref_audio=str(Path(self._tmp.name) / "gone.wav"))
+        self.assertIn("Reference audio not found", str(ctx.exception))
+
+    def test_clone_capable_model_allows_text_only(self):
+        client = self._client(model="higgs_audio_v3_tts")
+        self.assertIsNone(client.ref_audio)
+
+    def test_speaker_model_ignores_the_clone_reference(self):
+        client = self._client(model="qwen3_tts_0_6b_customvoice",
+                              ref_audio=str(self.ref))
+        self.assertIsNone(client.ref_audio)
+
+    def test_seed_only_sent_for_models_that_accept_it(self):
+        with patch("converter.clients.sglomni.resolve_request_seed",
+                   return_value=42):
+            client = self._client(model="qwen3_tts_1_7b_base",
+                                  ref_audio=str(self.ref))
+            self.assertEqual(client._seed, 42)
+            client = self._client(model="higgs_audio_v3_tts")
+            self.assertIsNone(client._seed)
+
+    def test_negative_seed_is_not_sent(self):
+        with patch("converter.clients.sglomni.resolve_request_seed",
+                   return_value=-1):
+            client = self._client(model="qwen3_tts_1_7b_base",
+                                  ref_audio=str(self.ref))
+            self.assertIsNone(client._seed)
+
+
+class ConnectHealthTests(unittest.TestCase):
+    """_connect gates on /health and the hosted model."""
+
+    def _response(self, payload):
+        response = MagicMock()
+        response.__enter__.return_value = response
+        response.read.return_value = json.dumps(payload).encode("utf-8")
+        return response
+
+    def _connect(self, payloads, **kwargs):
+        # urlopen is called once per _get_json call, in order.
+        responses = [self._response(payload) for payload in payloads]
+        with patch("converter.clients.sglomni.urllib.request.urlopen",
+                   side_effect=responses):
+            with patch.object(SgOmniTTSClient, "_resolve_reference_text"):
+                return SgOmniTTSClient(_DUMMY_CHUNKS,
+                                       model="higgs_audio_v3_tts", **kwargs)
+
+    def test_unreachable_server_raises_with_guidance(self):
+        import urllib.error
+        with patch("converter.clients.sglomni.urllib.request.urlopen",
+                   side_effect=urllib.error.URLError("refused")):
+            with self.assertRaises(RuntimeError) as ctx:
+                SgOmniTTSClient(_DUMMY_CHUNKS, model="higgs_audio_v3_tts")
+        self.assertIn("not reachable", str(ctx.exception))
+        self.assertIn("sgl-omni", str(ctx.exception))
+
+    def test_booting_server_raises(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            self._connect([{"status": "unhealthy"}])
+        self.assertIn("not healthy", str(ctx.exception))
+
+    def test_foreign_hosted_model_raises_with_both_names(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            self._connect([
+                {"status": "healthy", "stages": []},
+                {"data": [{"id": "Zyphra/zonos2"}]},
+            ])
+        message = str(ctx.exception)
+        self.assertIn("Zyphra/zonos2", message)
+        self.assertIn("bosonai/higgs-audio-v3-tts-4b", message)
+
+    def test_matching_model_connects(self):
+        client = self._connect([
+            {"status": "healthy", "stages": []},
+            {"data": [{"id": "bosonai/higgs-audio-v3-tts-4b"}]},
+        ])
+        self.assertEqual(client.entry.repo, "bosonai/higgs-audio-v3-tts-4b")
+
+
+class PayloadTests(unittest.TestCase):
+    """The /v1/audio/speech request shape per voice capability."""
+
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.ref = Path(self._tmp.name) / "narrator.wav"
+        self.ref.write_bytes(b"abc")
+        self.addCleanup(self._tmp.cleanup)
+
+    def _make_client(self, model, **kwargs):
+        client = SgOmniTTSClient.__new__(SgOmniTTSClient)
+        from backends.sglomni.catalog import entry_by_key
+        client.entry = entry_by_key(model)
+        client.api_url = "http://127.0.0.1:8100"
+        client.voice = kwargs.get("voice")
+        if "ref_audio" in kwargs:
+            kwargs["ref_audio"] = str(self.ref)
+        client.ref_audio = kwargs.get("ref_audio")
+        client.ref_text = kwargs.get("ref_text", "")
+        client.instructions = kwargs.get("instructions", "")
+        client.language = "English"
+        client._seed = None
+        return client
+
+    def test_speaker_payload_sends_the_preset_name(self):
+        client = self._make_client("qwen3_tts_0_6b_customvoice",
+                                   voice="Vivian")
+        payload = client._request_payload("Hello.")
+        self.assertEqual(payload["voice"], "Vivian")
+        self.assertEqual(payload["model"],
+                         "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
+        self.assertEqual(payload["response_format"], "wav")
+        self.assertNotIn("ref_audio", payload)
+        self.assertNotIn("task_type", payload)
+
+    def test_speaker_without_voice_uses_the_server_default(self):
+        client = self._make_client("voxtral_tts")
+        self.assertEqual(client._request_payload("Hello.")["voice"],
+                         "default")
+
+    def test_design_payload_sends_task_type_and_instructions(self):
+        client = self._make_client("qwen3_tts_1_7b_voicedesign",
+                                   instructions="A warm narrator.")
+        payload = client._request_payload("Hello.")
+        self.assertEqual(payload["task_type"], "VoiceDesign")
+        self.assertEqual(payload["instructions"], "A warm narrator.")
+
+    def test_clone_payload_sends_reference_path_on_loopback(self):
+        client = self._make_client("higgs_audio_v3_tts",
+                                   ref_audio=str(self.ref),
+                                   ref_text="A transcript.")
+        payload = client._request_payload("Hello.")
+        self.assertEqual(payload["ref_audio"], str(self.ref.resolve()))
+        self.assertEqual(payload["ref_text"], "A transcript.")
+
+    def test_clone_payload_inlines_audio_for_remote_servers(self):
+        client = self._make_client("higgs_audio_v3_tts",
+                                   ref_audio=str(self.ref))
+        client.api_url = "http://10.20.30.40:8100"
+        payload = client._request_payload("Hello.")
+        self.assertTrue(payload["ref_audio"].startswith(
+            "data:audio/wav;base64,"))
+        self.assertEqual(
+            base64.b64decode(payload["ref_audio"].partition(";base64,")[2]),
+            b"abc")
+        self.assertNotIn("ref_text", payload)
+
+    def test_clone_without_reference_sends_no_reference_fields(self):
+        client = self._make_client("higgs_audio_v3_tts")
+        payload = client._request_payload("Hello.")
+        self.assertNotIn("ref_audio", payload)
+        self.assertEqual(payload["voice"], "default")
+
+    def test_seed_included_when_resolved(self):
+        client = self._make_client("qwen3_tts_1_7b_base",
+                                   ref_audio="x.wav")
+        client._seed = 7
+        self.assertEqual(client._request_payload("Hello.")["seed"], 7)
+
+
+class RequestErrorTests(unittest.TestCase):
+    """OpenAI-style error envelopes decide retryability."""
+
+    def _make_client(self):
+        return SgOmniTTSClient.__new__(SgOmniTTSClient)
+
+    def test_bad_request_envelope_is_not_retryable(self):
+        client = self._make_client()
+        detail = json.dumps({"error": {
+            "message": "voice 'nope' not found",
+            "type": "BadRequestError", "code": 400}})
+        error = client._request_error(400, detail)
+        self.assertIsInstance(error, NonRetryableTTSError)
+        self.assertIn("voice 'nope' not found", str(error))
+
+    def test_server_error_is_retryable(self):
+        client = self._make_client()
+        error = client._request_error(503, "overloaded")
+        self.assertNotIsInstance(error, NonRetryableTTSError)
+
+    def test_non_json_4xx_is_not_retryable(self):
+        client = self._make_client()
+        error = client._request_error(422, "plain text rejection")
+        self.assertIsInstance(error, NonRetryableTTSError)
+
+
+class GenerateChunkTests(unittest.TestCase):
+    """Chunk generation: WAV output, sub-chunking, bookkeeping."""
+
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self._sleep = patch("converter.clients.base.time.sleep")
+        self._sleep.start()
+        self.addCleanup(self._sleep.stop)
+        self.addCleanup(self._tmp.cleanup)
+
+    def _make_client(self):
+        client = SgOmniTTSClient.__new__(SgOmniTTSClient)
+        from backends.sglomni.catalog import entry_by_key
+        client.entry = entry_by_key("higgs_audio_v3_tts")
+        client.chunks_dir = Path(self._tmp.name)
+        client.api_url = "http://127.0.0.1:8100"
+        client.voice = None
+        client.ref_audio = None
+        client.ref_text = ""
+        client.instructions = ""
+        client.language = "English"
+        client._seed = None
+        return client
+
+    def _read_wav(self, path):
+        with wave.open(str(path), "rb") as wav_file:
+            return wav_file.readframes(wav_file.getnframes())
+
+    def test_generate_chunk_writes_the_wav_response(self):
+        client = self._make_client()
+        with patch.object(client, "_request_wav", return_value=_WAV_BYTES):
+            result = client.generate_chunk("Hello world.", 1)
+        self.assertIsNotNone(result)
+        path = Path(result)
+        self.assertEqual(path.name, "chunk_0001.wav")
+        self.assertEqual(self._read_wav(path), _WAV_FRAMES)
+
+    def test_long_text_is_subchunked_and_concatenated(self):
+        client = self._make_client()
+        text = " ".join(f"word{i}" for i in range(24))
+        responses = [_WAV_BYTES, _WAV_BYTES, _WAV_BYTES]
+        with patch.object(config, "CHUNK_SIZE", 10), \
+                patch.object(client, "_request_wav",
+                             side_effect=responses) as mock_wav, \
+                patch("converter.clients.sglomni.concat_audio_files") as mock_concat:
+            result = client.generate_chunk(text, 1)
+        # 24 words at CHUNK_SIZE 10 -> three sub-requests (10/10/4).
+        self.assertEqual(mock_wav.call_count, 3)
+        self.assertIsNotNone(result)
+        mock_concat.assert_called_once()
+        args = mock_concat.call_args[0]
+        self.assertEqual(len(args[0]), 3)
+        self.assertEqual(args[1], Path(result))
+
+    def test_single_subchunk_skips_concatenation(self):
+        client = self._make_client()
+        with patch.object(client, "_request_wav", return_value=_WAV_BYTES), \
+                patch("converter.clients.sglomni.concat_audio_files") as mock_concat:
+            client.generate_chunk("Hello.", 1)
+        mock_concat.assert_not_called()
+
+    def test_empty_text_fails_the_chunk(self):
+        client = self._make_client()
+        with patch.object(client, "_request_wav") as mock_wav:
+            self.assertIsNone(client.generate_chunk("   ", 1))
+        mock_wav.assert_not_called()
+
+    def test_request_failure_fails_the_chunk_attempt(self):
+        client = self._make_client()
+        with patch.object(client, "_request_wav",
+                          side_effect=RuntimeError("down")) as mock_wav:
+            self.assertIsNone(client.generate_chunk("Hello.", 1))
+        self.assertEqual(mock_wav.call_count, 1)
+
+    def test_stale_chunk_files_are_removed(self):
+        stale = Path(self._tmp.name) / "chunk_0001.mp3"
+        stale.write_bytes(b"old")
+        client = self._make_client()
+        with patch.object(client, "_request_wav", return_value=_WAV_BYTES):
+            client.generate_chunk("Hello.", 1)
+        remaining = sorted(path.name for path in
+                           Path(self._tmp.name).glob("chunk_0001.*"))
+        self.assertEqual(remaining, ["chunk_0001.wav"])
+
+
+if __name__ == "__main__":
+    unittest.main()