first

speech_to_speech/connections/socket_receiver.py

@@ -0,0 +1,85 @@
+import logging
+import socket
+import time
+from queue import Queue
+from threading import Event
+
+from rich.console import Console
+
+from speech_to_speech.pipeline.messages import PIPELINE_END
+from speech_to_speech.pipeline.queue_types import AudioInItem
+
+logger = logging.getLogger(__name__)
+
+console = Console()
+
+# If should_listen stays cleared for longer than this, something in the
+# pipeline probably failed (LLM exception, TTS crash, etc.).  Re-enable
+# listening so the user isn't permanently locked out.
+SHOULD_LISTEN_TIMEOUT_S = 30.0
+
+
+class SocketReceiver:
+    """
+    Handles reception of the audio packets from the client.
+    """
+
+    def __init__(
+        self,
+        stop_event: Event,
+        queue_out: Queue[AudioInItem],
+        should_listen: Event,
+        host: str = "0.0.0.0",
+        port: int = 12345,
+        chunk_size: int = 1024,
+    ) -> None:
+        self.stop_event = stop_event
+        self.queue_out = queue_out
+        self.should_listen = should_listen
+        self.chunk_size = chunk_size
+        self.host = host
+        self.port = port
+
+    def receive_full_chunk(self, conn: socket.socket, chunk_size: int) -> bytes | None:
+        data = b""
+        while len(data) < chunk_size:
+            packet = conn.recv(chunk_size - len(data))
+            if not packet:
+                # connection closed
+                return None
+            data += packet
+        return data
+
+    def run(self) -> None:
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self.socket.bind((self.host, self.port))
+        self.socket.listen(1)
+        logger.info("Receiver waiting to be connected...")
+        self.conn, _ = self.socket.accept()
+        logger.info("receiver connected")
+
+        self.should_listen.set()
+        listen_cleared_at = None
+        while not self.stop_event.is_set():
+            audio_chunk = self.receive_full_chunk(self.conn, self.chunk_size)
+            if audio_chunk is None:
+                # connection closed
+                self.queue_out.put(PIPELINE_END)
+                break
+            if self.should_listen.is_set():
+                self.queue_out.put(audio_chunk)
+                listen_cleared_at = None
+            else:
+                # Track how long should_listen has been cleared
+                if listen_cleared_at is None:
+                    listen_cleared_at = time.monotonic()
+                elif time.monotonic() - listen_cleared_at > SHOULD_LISTEN_TIMEOUT_S:
+                    logger.warning(
+                        "should_listen has been cleared for %.0fs — pipeline may be stuck, re-enabling listening",
+                        SHOULD_LISTEN_TIMEOUT_S,
+                    )
+                    self.should_listen.set()
+                    listen_cleared_at = None
+        self.conn.close()
+        logger.info("Receiver closed")