From 0d74790c670e95f461db45374b9bc1cd8c8f8e81 Mon Sep 17 00:00:00 2001
From: makaveli10 <vineet.suryan@collabora.com>
Date: Thu, 10 Oct 2024 08:37:45 -0400
Subject: [PATCH] Expose ClientManager arguments to be passed from client

Signed-off-by: makaveli10 <vineet.suryan@collabora.com>
---
 whisper_live/client.py | 19 ++++++++++++++++---
 whisper_live/server.py |  8 +++++++-
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/whisper_live/client.py b/whisper_live/client.py
index c252607..4cfb63c 100644
--- a/whisper_live/client.py
+++ b/whisper_live/client.py
@@ -30,7 +30,9 @@ class Client:
         model="small",
         srt_file_path="output.srt",
         use_vad=True,
-        log_transcription=True
+        log_transcription=True,
+        max_clients=4,
+        max_connection_time=600,
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -59,6 +61,8 @@ class Client:
         self.last_segment = None
         self.last_received_segment = None
         self.log_transcription = log_transcription
+        self.max_clients = max_clients
+        self.max_connection_time = max_connection_time
 
         if translate:
             self.task = "translate"
@@ -199,7 +203,9 @@ class Client:
                     "language": self.language,
                     "task": self.task,
                     "model": self.model,
-                    "use_vad": self.use_vad
+                    "use_vad": self.use_vad,
+                    "max_clients": self.max_clients,
+                    "max_connection_time": self.max_connection_time,
                 }
             )
         )
@@ -681,8 +687,15 @@ class TranscriptionClient(TranscriptionTeeClient):
         output_recording_filename="./output_recording.wav",
         output_transcription_path="./output.srt",
         log_transcription=True,
+        max_clients=4,
+        max_connection_time=600,
     ):
-        self.client = Client(host, port, lang, translate, model, srt_file_path=output_transcription_path, use_vad=use_vad, log_transcription=log_transcription)
+        self.client = Client(
+            host, port, lang, translate, model, srt_file_path=output_transcription_path,
+            use_vad=use_vad, log_transcription=log_transcription, max_clients=max_clients,
+            max_connection_time=max_connection_time
+        )
+
         if save_output_recording and not output_recording_filename.endswith(".wav"):
             raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}")
         if not output_transcription_path.endswith(".srt"):
diff --git a/whisper_live/server.py b/whisper_live/server.py
index e3346d2..575a2cf 100644
--- a/whisper_live/server.py
+++ b/whisper_live/server.py
@@ -147,7 +147,7 @@ class TranscriptionServer:
     RATE = 16000
 
     def __init__(self):
-        self.client_manager = ClientManager()
+        self.client_manager = None
         self.no_voice_activity_chunks = 0
         self.use_vad = True
         self.single_model = False
@@ -224,6 +224,12 @@ class TranscriptionServer:
             logging.info("New client connected")
             options = websocket.recv()
             options = json.loads(options)
+            
+            if self.client_manager is None:
+                max_clients = options.get('max_clients', 4)
+                max_connection_time = options.get('max_connection_time', 600)
+                self.client_manager = ClientManager(max_clients, max_connection_time)
+
             self.use_vad = options.get('use_vad')
             if self.client_manager.is_server_full(websocket, options):
                 websocket.close()