From a71c5785704b2053abf0c7390463d76c1160d79d Mon Sep 17 00:00:00 2001 From: Quang Tran Date: Wed, 8 Jul 2026 22:45:06 +0700 Subject: [PATCH] fix: address PR review comments on StreamingTranscriptionClient --- README.md | 2 +- examples/manual_audio_chunking.py | 4 +- tests/test_streaming_client.py | 164 ++++++++++++++++++++++++++++++ whisper_live/client.py | 34 ++++--- 4 files changed, 186 insertions(+), 18 deletions(-) create mode 100644 tests/test_streaming_client.py diff --git a/README.md b/README.md index 4f4646f..ba1231a 100644 --- a/README.md +++ b/README.md @@ -265,7 +265,7 @@ Audio is automatically normalized to float32 range [-1.0, 1.0]. Clients can also `StreamingTranscriptionClient` lets you push raw PCM audio bytes from any source — a live microphone capture loop, a network stream, an audio pipeline — and receive transcripts via callbacks as speech is detected. Unlike `TranscriptionClient`, it does not manage audio capture internally; you control when and how audio is fed. -A runnable example that reads from an audio file and stream the chunks is at [`examples/manual_audio_chunking.py`](examples/manual_audio_chunking.py): +A runnable example that reads from an audio file and streams the chunks is at [`examples/manual_audio_chunking.py`](examples/manual_audio_chunking.py): ```bash python examples/manual_audio_chunking.py --file assets/jfk.flac diff --git a/examples/manual_audio_chunking.py b/examples/manual_audio_chunking.py index 2022ca7..790853b 100644 --- a/examples/manual_audio_chunking.py +++ b/examples/manual_audio_chunking.py @@ -69,8 +69,8 @@ def main(): print("\n[INFO] Final transcript:") for seg in client.transcript: print(f" [{float(seg['start']):.2f}s → {float(seg['end']):.2f}s] {seg['text'].strip()}") - if client.last_partial: - seg = client.last_segment + seg = client.last_partial + if seg: print(f" [{float(seg['start']):.2f}s → {float(seg['end']):.2f}s] {seg['text'].strip()} (partial)") diff --git a/tests/test_streaming_client.py b/tests/test_streaming_client.py new file mode 100644 index 0000000..4c8933a --- /dev/null +++ b/tests/test_streaming_client.py @@ -0,0 +1,164 @@ +import json +import unittest +from unittest.mock import patch, MagicMock + +import numpy as np + +from whisper_live.client import Client, StreamingTranscriptionClient + + +class StreamingClientTestCase(unittest.TestCase): + @patch('whisper_live.client.websocket.WebSocketApp') + def setUp(self, mock_websocket): + self.mock_websocket = mock_websocket + self.mock_ws_app = mock_websocket.return_value + self.mock_ws_app.send = MagicMock() + + self.committed = [] + self.partials = [] + self.session_started = [] + + self.client = StreamingTranscriptionClient( + host='localhost', + port=9090, + lang="en", + on_session_started=lambda: self.session_started.append(True), + on_committed_transcript=lambda text, segs: self.committed.append((text, segs)), + on_partial_transcript=lambda text, segs: self.partials.append((text, segs)), + ) + self._inner = self.client._client + + def tearDown(self): + self._inner.close_websocket() + self.mock_websocket.stop() + + def _server_ready(self, backend="faster_whisper"): + self._inner.on_message(self.mock_ws_app, json.dumps({ + "uid": self._inner.uid, + "message": "SERVER_READY", + "backend": backend, + })) + + def _send_segments(self, segments): + self._inner.on_message(self.mock_ws_app, json.dumps({ + "uid": self._inner.uid, + "segments": segments, + })) + + +class TestPcmFormatConversion(StreamingClientTestCase): + def test_int16_is_normalized_to_float32(self): + self._server_ready() + raw = np.array([0, 16384, -32768], dtype=np.int16).tobytes() + with patch.object(self._inner, 'send_packet_to_server') as mock_send: + self.client.send(raw, pcm_format="int16") + sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32) + np.testing.assert_allclose(sent, [0.0, 0.5, -1.0], atol=1e-4) + + def test_float32_passes_through(self): + self._server_ready() + raw = np.array([0.1, -0.2], dtype=np.float32).tobytes() + with patch.object(self._inner, 'send_packet_to_server') as mock_send: + self.client.send(raw, pcm_format="float32") + self.assertEqual(mock_send.call_args[0][0], raw) + + def test_default_format_is_int16(self): + self._server_ready() + raw = np.array([32767], dtype=np.int16).tobytes() + with patch.object(self._inner, 'send_packet_to_server') as mock_send: + self.client.send(raw) + sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32) + self.assertAlmostEqual(float(sent[0]), 32767 / 32768.0, places=4) + + def test_unsupported_format_raises(self): + self._server_ready() + with self.assertRaises(ValueError): + self.client.send(b"\x00\x00", pcm_format="int8") + + def test_send_after_close_raises(self): + self.client._closed = True + with self.assertRaises(RuntimeError): + self.client.send(b"\x00\x00", pcm_format="int16") + + def test_send_array_normalizes_integers(self): + with patch.object(self._inner, 'send_packet_to_server') as mock_send: + self.client.send_array(np.array([0, 16384, -32768], dtype=np.int16)) + sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32) + np.testing.assert_allclose(sent, [0.0, 0.5, -1.0], atol=1e-4) + + +class TestTranscriptDispatch(StreamingClientTestCase): + def test_partial_then_committed(self): + self._server_ready() + self._send_segments([{"start": 0, "end": 1, "text": "hello", "completed": False}]) + self.assertEqual(len(self.partials), 1) + self.assertEqual(self.partials[0][0], "hello") + self.assertEqual(len(self.committed), 0) + + self._send_segments([{"start": 0, "end": 1, "text": "hello world", "completed": True}]) + self.assertEqual(len(self.committed), 1) + self.assertEqual(self.committed[0][0], "hello world") + self.assertEqual(len(self.client.transcript), 1) + + def test_committed_deduplicated(self): + self._server_ready() + seg = {"start": 0, "end": 1, "text": "hi", "completed": True} + self._send_segments([seg]) + self._send_segments([seg]) + self.assertEqual(len(self.committed), 1) + self.assertEqual(len(self.client.transcript), 1) + + def test_committed_backend_agnostic(self): + """Committed dispatch must work for non-faster_whisper backends.""" + self._server_ready(backend="tensorrt") + self._send_segments([{"start": 0, "end": 1, "text": "trt seg", "completed": True}]) + self.assertEqual(len(self.committed), 1) + self.assertEqual(len(self.client.transcript), 1) + + def test_last_partial_alias(self): + self._server_ready() + self._send_segments([{"start": 0, "end": 1, "text": "pending", "completed": False}]) + self.assertIsNotNone(self.client.last_partial) + self.assertIs(self.client.last_partial, self.client.last_segment) + + +class TestConnectLifecycle(StreamingClientTestCase): + def test_connect_returns_after_ready(self): + self._server_ready() + self.assertIs(self.client.connect(), self.client) + self.assertEqual(len(self.session_started), 1) + + def test_connect_times_out(self): + self.client._ready_timeout = 0.1 + with self.assertRaises(TimeoutError): + self.client.connect() + + def test_connect_raises_on_server_error(self): + self._inner.on_message(self.mock_ws_app, json.dumps({ + "uid": self._inner.uid, + "status": "ERROR", + "message": "boom", + })) + with self.assertRaises(RuntimeError): + self.client.connect() + + def test_connect_raises_when_server_full(self): + self._inner.on_message(self.mock_ws_app, json.dumps({ + "uid": self._inner.uid, + "status": "WAIT", + "message": 5, + })) + with self.assertRaises(RuntimeError): + self.client.connect() + + def test_close_sends_end_of_audio(self): + self._server_ready() + with patch.object(self._inner, 'send_packet_to_server') as mock_send, \ + patch.object(self._inner, 'close_websocket') as mock_close: + self.client.close(drain_seconds=0) + mock_send.assert_called_once_with(Client.END_OF_AUDIO.encode("utf-8")) + mock_close.assert_called_once() + + +if __name__ == '__main__': + unittest.main() diff --git a/whisper_live/client.py b/whisper_live/client.py index fb6555c..88a9677 100644 --- a/whisper_live/client.py +++ b/whisper_live/client.py @@ -927,10 +927,10 @@ class _HookedClient(Client): super().__init__(*args, **kwargs) def on_message(self, ws, message): - data = json.loads(message) - if data.get("message") == "SERVER_READY" and self._on_session_started: - self._on_session_started() + was_recording = self.recording super().on_message(ws, message) + if not was_recording and self.recording and self._on_session_started: + self._on_session_started() def on_error(self, ws, error): if self._on_error_hook: @@ -1007,7 +1007,8 @@ class StreamingTranscriptionClient: self._on_committed_transcript = on_committed_transcript self._ready_timeout = ready_timeout self._closed = False - self._last_committed_count = 0 + self._transcript = [] + self._committed_keys = set() self._client = _HookedClient( host=host, @@ -1031,11 +1032,16 @@ class StreamingTranscriptionClient: ) def _dispatch_transcript(self, text: str, segments: list) -> None: - new_committed = self._client.transcript[self._last_committed_count:] - for seg in new_committed: + for seg in segments: + if not seg.get("completed", False): + continue + key = (seg.get("start"), seg.get("end"), seg.get("text")) + if key in self._committed_keys: + continue + self._committed_keys.add(key) + self._transcript.append(seg) if self._on_committed_transcript: self._on_committed_transcript(seg["text"].strip(), [seg]) - self._last_committed_count = len(self._client.transcript) last = segments[-1] if segments else None if last and not last.get("completed", False) and self._on_partial_transcript: @@ -1054,12 +1060,12 @@ class StreamingTranscriptionClient: time.sleep(0.05) return self - def send(self, audio_bytes: bytes, pcm_format: PcmFormat = "float32") -> None: + def send(self, audio_bytes: bytes, pcm_format: PcmFormat = "int16") -> None: """Send one PCM chunk. Any chunk size is fine; must be mono 16 kHz. Args: audio_bytes: Raw PCM payload. - pcm_format: ``"float32"`` passes through; ``"int16"`` is normalized to float32. + pcm_format: ``"int16"`` is normalized to float32; ``"float32"`` passes through. """ if self._closed: raise RuntimeError("Client is already closed.") @@ -1092,17 +1098,15 @@ class StreamingTranscriptionClient: @property def transcript(self) -> list: """All committed segments received so far.""" - return self._client.transcript + return self._transcript @property def last_partial(self) -> Optional[dict]: """The most recent in-progress segment, or ``None`` if none pending.""" return self._client.last_segment - - @property - def last_segment(self) -> Optional[dict]: - """The most recent in-progress segment, or ``None`` if none pending.""" - return self._client.last_segment + + # Alias for ``last_partial``; kept for readability at call sites. + last_segment = last_partial def close(self, drain_seconds: float = 2.0) -> None: """Signal end-of-stream, wait briefly for final transcripts, then close.