fix: apply suggestions from preliminary review

Author: rachwalk

src/rai/rai/agents/voice_agent.py

@@ -13,15 +13,16 @@
 # limitations under the License.
 
 
+import time
 from threading import Lock, Thread
-from typing import Any, List, Tuple
+from typing import Any, List, cast
 
 import numpy as np
 from numpy.typing import NDArray
 
 from rai.agents.base import BaseAgent
 from rai.communication import AudioInputDeviceConfig, StreamingAudioInputDevice
-from rai_asr.models.base import BaseVoiceDetectionModel
+from rai_asr.models.base import BaseTranscriptionModel, BaseVoiceDetectionModel
 
 
 class VoiceRecognitionAgent(BaseAgent):
@@ -38,34 +39,50 @@ def __call__(self):
         self.run()
 
     def setup(
-        self, microphone_device_id: int, microphone_config: AudioInputDeviceConfig
+        self,
+        microphone_device_id: int,  # TODO: Change to name based instead of id based identification
+        microphone_config: AudioInputDeviceConfig,
+        transcription_model: BaseTranscriptionModel,
     ):
-        assert isinstance(self.connectors["microphone"], StreamingAudioInputDevice)
+        self.connectors["microphone"] = cast(
+            StreamingAudioInputDevice, self.connectors["microphone"]
+        )
         self.microphone_device_id = str(microphone_device_id)
         self.connectors["microphone"].configure_device(
             target=self.microphone_device_id, config=microphone_config
         )
+        self.transcription_model = transcription_model
         self.ran_setup = True
+        self.running = False
+
+    def add_detection_model(
+        self, model: BaseVoiceDetectionModel, pipeline: str = "record"
+    ):
+        if pipeline == "record":
+            self.should_record_pipeline.append(model)
+        elif pipeline == "stop":
+            self.should_stop_pipeline.append(model)
+        else:
+            raise ValueError("Pipeline should be either 'record' or 'stop'")
 
     def run(self):
+        self.running = True
         self.listener_handle = self.connectors["microphone"].start_action(
             self.microphone_device_id, self.on_new_sample
         )
         self.transcription_thread = Thread(target=self._transcription_function)
         self.transcription_thread.start()
 
     def stop(self):
+        self.running = False
         self.connectors["microphone"].terminate_action(self.listener_handle)
         self.transcription_thread.join()
 
     def on_new_sample(self, indata: np.ndarray, status_flags: dict[str, Any]):
-        should_stop, should_cancel = self.should_stop_recording(indata)
-        print(indata)
-        if should_cancel:
-            self.cancel_task()
-        if (self.recording_started and not should_stop) or (
-            self.should_start_recording(indata)
-        ):
+        should_stop = self.should_stop_recording(indata)
+        if self.should_start_recording(indata):
+            self.recording_started = True
+        if self.recording_started and not should_stop:
             with self.transcription_lock:
                 self.shared_samples.extend(indata)
 
@@ -75,23 +92,27 @@ def should_start_recording(self, audio_data: NDArray[np.int16]) -> bool:
             should_listen, output_parameters = model.detected(
                 audio_data, output_parameters
             )
+            print(should_listen, output_parameters)
             if not should_listen:
                 return False
         return True
 
-    def should_stop_recording(self, audio_data: NDArray[np.int16]) -> Tuple[bool, bool]:
+    def should_stop_recording(self, audio_data: NDArray[np.int16]) -> bool:
         output_parameters = {}
         for model in self.should_stop_pipeline:
             should_listen, output_parameters = model.detected(
                 audio_data, output_parameters
             )
-            # TODO: Add handling output parametrs for checking if should cancel
             if should_listen:
-                return False, False
-        return True, False
+                return True
+        return False
 
     def _transcription_function(self):
-        with self.transcription_lock:
-            samples = np.array(self.shared_samples)
-            print(samples)
-            self.shared_samples = []
+        while self.running:
+            time.sleep(0.1)
+            # critical section for samples
+            with self.transcription_lock:
+                samples = np.array(self.shared_samples)
+                self.shared_samples = []
+            # end critical section for samples
+            self.transcription_model.add_samples(samples)

src/rai_asr/rai_asr/models/base.py

@@ -36,8 +36,9 @@ def __init__(self, model_name: str, sample_rate: int, language: str = "en"):
         self.language = language
 
     @abstractmethod
-    def transcribe(self, data: NDArray[np.int16]) -> str:
+    def add_samples(self, data: NDArray[np.int16]):
         pass
 
-    def __call__(self, data: NDArray[np.int16]) -> str:
-        return self.transcribe(data)
+    @abstractmethod
+    def transcribe(self) -> str:
+        pass

src/rai_asr/rai_asr/models/local_whisper.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import cast
+
 import numpy as np
 import whisper
 from numpy._typing import NDArray
@@ -24,9 +26,22 @@ def __init__(self, model_name: str, sample_rate: int, language: str = "en"):
         super().__init__(model_name, sample_rate, language)
         self.whisper = whisper.load_model(self.model_name)
 
-    def transcribe(self, data: NDArray[np.int16]) -> str:
-        result = whisper.transcribe(self.whisper, data.astype(np.float32) / 32768.0)
+        self.samples = None
+
+    def add_samples(self, data: NDArray[np.int16]):
+        normalized_data = data.astype(np.float32) / 32768.0
+        self.samples = (
+            np.concatenate([self.samples, normalized_data])
+            if self.samples is not None
+            else data
+        )
+
+    def transcribe(self) -> str:
+        if self.samples is None:
+            raise ValueError("No samples to transcribe")
+        result = whisper.transcribe(
+            self.whisper, self.samples
+        )  # TODO: handling of additional transcribe arguments (perhaps in model init)
         transcription = result["text"]
-        # NOTE: this is only for type enforcement, doesn't need to work on runtime
-        assert isinstance(transcription, str)
+        transcription = cast(str, transcription)
         return transcription

src/rai_asr/rai_asr/models/open_ai_whisper.py

@@ -36,10 +36,19 @@ def __init__(self, model_name: str, sample_rate: int, language: str = "en"):
             self.openai_client.audio.transcriptions.create,
             model=self.model_name,
         )
+        self.samples = []
 
-    def transcribe(self, data: NDArray[np.int16]) -> str:
+    def add_samples(self, data: NDArray[np.int16]):
+        normalized_data = data.astype(np.float32) / 32768.0
+        self.samples = (
+            np.concatenate([self.samples, normalized_data])
+            if self.samples is not None
+            else data
+        )
+
+    def transcribe(self) -> str:
         with io.BytesIO() as temp_wav_buffer:
-            wavfile.write(temp_wav_buffer, self.sample_rate, data)
+            wavfile.write(temp_wav_buffer, self.sample_rate, self.samples)
             temp_wav_buffer.seek(0)
             temp_wav_buffer.name = "temp.wav"
             response = self.model(file=temp_wav_buffer, language=self.language)

src/rai_asr/rai_asr/models/silero_vad.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Tuple
+from typing import Any, Literal, Tuple
 
 import numpy as np
 import torch
@@ -22,7 +22,7 @@
 
 
 class SileroVAD(BaseVoiceDetectionModel):
-    def __init__(self, sampling_rate=16000, threshold=0.5):
+    def __init__(self, sampling_rate: Literal[8000, 16000] = 16000, threshold=0.5):
         super(SileroVAD, self).__init__()
         self.model_name = "silero_vad"
         self.model, _ = torch.hub.load(