diff --git a/tests/test_base_backend.py b/tests/test_base_backend.py index 8f90b77..624ec7f 100644 --- a/tests/test_base_backend.py +++ b/tests/test_base_backend.py @@ -266,6 +266,33 @@ class TestCleanup(unittest.TestCase): self.assertTrue(client.exit) +class TestTrimTranscript(unittest.TestCase): + def setUp(self): + self.ws = MagicMock() + self.client = ConcreteServeClient(client_uid="test", websocket=self.ws) + + def test_transcript_trimmed_when_over_max(self): + self.client.transcript = [ + {"start": f"{i}.000", "end": f"{i+1}.000", "text": f"seg{i}", "completed": True} + for i in range(self.client.MAX_TRANSCRIPT_LENGTH + 100) + ] + self.client._trim_transcript() + self.assertEqual(len(self.client.transcript), self.client.MAX_TRANSCRIPT_LENGTH) + self.assertEqual(self.client.transcript[0]["text"], "seg100") + + def test_transcript_not_trimmed_when_under_max(self): + self.client.transcript = [ + {"start": "0.000", "end": "1.000", "text": "a", "completed": True} + ] + self.client._trim_transcript() + self.assertEqual(len(self.client.transcript), 1) + + def test_text_list_trimmed(self): + self.client.text = ["word"] * (self.client.MAX_TRANSCRIPT_LENGTH + 50) + self.client._trim_transcript() + self.assertEqual(len(self.client.text), self.client.MAX_TRANSCRIPT_LENGTH) + + class TestUpdateSegments(unittest.TestCase): """Tests for the core update_segments() logic.""" diff --git a/whisper_live/backend/base.py b/whisper_live/backend/base.py index e45e503..855e310 100644 --- a/whisper_live/backend/base.py +++ b/whisper_live/backend/base.py @@ -24,6 +24,9 @@ class ServeClientBase(object): same_output_threshold: int """Number of repeated outputs before considering it as a valid segment.""" + MAX_TRANSCRIPT_LENGTH = 500 + MAX_TRANSLATION_QUEUE_SIZE = 100 + def __init__( self, client_uid, @@ -376,4 +379,12 @@ class ServeClientBase(object): with self.lock: self.timestamp_offset += offset + self._trim_transcript() return last_segment + + def _trim_transcript(self): + """Trims transcript and text lists to prevent unbounded memory growth.""" + if len(self.transcript) > self.MAX_TRANSCRIPT_LENGTH: + self.transcript = self.transcript[-self.MAX_TRANSCRIPT_LENGTH:] + if len(self.text) > self.MAX_TRANSCRIPT_LENGTH: + self.text = self.text[-self.MAX_TRANSCRIPT_LENGTH:] diff --git a/whisper_live/server.py b/whisper_live/server.py index e3aaf6b..e342119 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -191,7 +191,7 @@ class TranscriptionServer: if enable_translation: target_language = options.get("target_language", "fr") - translation_queue = queue.Queue() + translation_queue = queue.Queue(maxsize=ServeClientBase.MAX_TRANSLATION_QUEUE_SIZE) from whisper_live.backend.translation_backend import ServeClientTranslation translation_client = ServeClientTranslation( client_uid=options["uid"],