This commit is contained in:
2026-08-26 17:42:00 -04:00
commit ad95d530ad
134 changed files with 31750 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
@@ -0,0 +1,71 @@
import logging
import threading
import time
from queue import Queue
import numpy as np
import sounddevice as sd
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE
from speech_to_speech.pipeline.queue_types import AudioInItem, AudioOutItem
logger = logging.getLogger(__name__)
class LocalAudioStreamer:
def __init__(
self,
input_queue: Queue[AudioInItem],
output_queue: Queue[AudioOutItem],
should_listen: threading.Event,
list_play_chunk_size: int = 512,
) -> None:
self.list_play_chunk_size = list_play_chunk_size
self.stop_event = threading.Event()
self.input_queue = input_queue
self.output_queue = output_queue
self.should_listen = should_listen
def run(self) -> None:
# Pre-generate a static dither buffer (±1 LSB, -96 dB) to keep the
# audio sink active without calling numpy inside the real-time callback.
dither = np.random.randint(-1, 2, size=(self.list_play_chunk_size, 1), dtype=np.int16)
def callback(indata: np.ndarray, outdata: np.ndarray, frames: int, time: float, status: str) -> None:
# During shutdown, just output silence
if self.stop_event.is_set():
outdata[:] = 0 * outdata
return
if self.output_queue.empty():
pcm = np.ascontiguousarray(indata, dtype=np.int16)
self.input_queue.put(pcm.tobytes())
outdata[:] = dither
else:
try:
audio_chunk = self.output_queue.get_nowait()
if isinstance(audio_chunk, np.ndarray):
outdata[:] = audio_chunk[:, np.newaxis]
elif audio_chunk == AUDIO_RESPONSE_DONE:
self.should_listen.set()
logger.debug("Response complete, listening re-enabled")
outdata[:] = 0 * outdata
else:
outdata[:] = 0 * outdata
except Exception:
outdata[:] = 0 * outdata
logger.debug("Available devices:")
logger.debug(sd.query_devices())
with sd.Stream(
samplerate=16000,
dtype="int16",
channels=1,
callback=callback,
blocksize=self.list_play_chunk_size,
):
logger.info("Starting local audio stream")
while not self.stop_event.is_set():
time.sleep(0.001)
print("Stopping recording")
@@ -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")
@@ -0,0 +1,69 @@
import logging
import socket
from queue import Empty, Queue
from threading import Event
import numpy as np
from rich.console import Console
from speech_to_speech.pipeline.control import PipelineControlMessage
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END
from speech_to_speech.pipeline.queue_types import AudioOutItem
logger = logging.getLogger(__name__)
console = Console()
class SocketSender:
"""
Handles sending generated audio packets to the clients.
"""
def __init__(
self,
stop_event: Event,
queue_in: Queue[AudioOutItem],
should_listen: Event,
host: str = "0.0.0.0",
port: int = 12346,
) -> None:
self.stop_event = stop_event
self.queue_in = queue_in
self.should_listen = should_listen
self.host = host
self.port = port
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("Sender waiting to be connected...")
self.conn, _ = self.socket.accept()
logger.info("sender connected")
while not self.stop_event.is_set():
try:
audio_chunk = self.queue_in.get(timeout=0.5)
except Empty:
continue
if isinstance(audio_chunk, PipelineControlMessage):
continue
if isinstance(audio_chunk, bytes) and audio_chunk == AUDIO_RESPONSE_DONE:
self.should_listen.set()
continue
payload: bytes
if isinstance(audio_chunk, bytes):
payload = audio_chunk
elif isinstance(audio_chunk, np.ndarray):
payload = audio_chunk.tobytes()
elif hasattr(audio_chunk, "tobytes"):
payload = audio_chunk.tobytes()
else:
continue
self.conn.sendall(payload)
if isinstance(audio_chunk, bytes) and audio_chunk == PIPELINE_END:
break
self.conn.close()
logger.info("Sender closed")
@@ -0,0 +1,238 @@
import asyncio
import json
import logging
from queue import Empty, Queue
from threading import Event
from typing import Any
import numpy as np
from websockets.asyncio.server import ServerConnection
from speech_to_speech.pipeline.control import SESSION_END, PipelineControlMessage, is_control_message
from speech_to_speech.pipeline.events import PipelineEvent
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END
from speech_to_speech.pipeline.queue_types import AudioInItem, AudioOutItem, TextEventItem
logger = logging.getLogger(__name__)
class WebSocketStreamer:
"""
Handles bidirectional audio streaming over WebSocket.
Receives audio from clients and puts it in the input_queue.
Sends audio from the output_queue to clients.
Sends text messages (transcripts/tools) from text_output_queue to clients.
"""
def __init__(
self,
stop_event: Event,
input_queue: Queue[AudioInItem],
output_queue: Queue[AudioOutItem],
should_listen: Event,
text_output_queue: Queue[TextEventItem] | None = None,
host: str = "0.0.0.0",
port: int = 8765,
) -> None:
self.stop_event = stop_event
self.input_queue = input_queue # clients -> VAD
self.output_queue = output_queue # TTS -> clients
self.text_output_queue = text_output_queue # Text messages -> clients
self.should_listen = should_listen
self.host = host
self.port = port
self.clients: set[ServerConnection] = set()
self.loop: asyncio.AbstractEventLoop | None = None
self.server: Any = None
def run(self) -> None:
"""Run the WebSocket server (called from a thread)."""
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop)
try:
self.loop.run_until_complete(self._run_server())
except Exception as e:
logger.error(f"WebSocket server error: {e}")
finally:
self.loop.close()
async def _run_server(self) -> None:
"""Main async server loop."""
import websockets
logger.info(f"WebSocket server starting on ws://{self.host}:{self.port}")
self.server = await websockets.serve(
self._handle_client,
self.host,
self.port,
)
logger.info("WebSocket server ready, waiting for connections...")
# Start the sender task
sender_task = asyncio.create_task(self._send_loop())
# Wait until stop_event is set
while not self.stop_event.is_set():
await asyncio.sleep(0.1)
# Cleanup
sender_task.cancel()
try:
await sender_task
except asyncio.CancelledError:
pass
# Close all clients
for client in list(self.clients):
try:
await client.close()
except Exception:
pass
self.server.close()
await self.server.wait_closed()
logger.info("WebSocket server closed")
async def _handle_client(self, websocket: ServerConnection) -> None:
"""Handle a single WebSocket client connection."""
client_id = id(websocket)
logger.info(f"Client {client_id} connected")
self.clients.add(websocket)
recv_buffer = bytearray()
# Enable listening when first client connects
if len(self.clients) == 1:
# Drain edge queues so stale data from a previous session doesn't
# leak into the new one (SESSION_END may not have flushed everything).
for q in (self.output_queue, self.text_output_queue):
if q is not None:
while not q.empty():
try:
q.get_nowait()
except Empty:
break
self.should_listen.set()
logger.debug("Listening enabled, edge queues drained (should_listen.set())")
try:
logger.debug(f"Client {client_id}: Starting message receive loop")
async for message in websocket:
if isinstance(message, bytes):
logger.debug(f"Client {client_id}: Received {len(message)} bytes of audio")
if self.should_listen.is_set():
# Split into 512-sample (1024 bytes) chunks for VAD.
# Keep a per-client remainder buffer so no samples are dropped
# when WebSocket frame boundaries are not aligned.
chunk_size_bytes = 512 * 2 # 512 samples * 2 bytes per int16
recv_buffer.extend(message)
num_chunks = 0
while len(recv_buffer) >= chunk_size_bytes:
chunk = bytes(recv_buffer[:chunk_size_bytes])
del recv_buffer[:chunk_size_bytes]
self.input_queue.put(chunk)
num_chunks += 1
logger.debug(f"Client {client_id}: Queued {num_chunks} chunks for processing")
else:
logger.debug(f"Client {client_id}: Skipping audio (should_listen not set)")
except Exception as e:
logger.error(f"Client {client_id} error: {type(e).__name__}: {e}", exc_info=True)
finally:
self.clients.discard(websocket)
logger.info(f"Client {client_id} disconnected (finally block)")
if len(self.clients) == 0:
logger.debug("Last WebSocket client disconnected, ending session")
self.input_queue.put(SESSION_END)
async def _send_loop(self) -> None:
"""Send audio and text from queues to all connected clients."""
# Buffer audio until we have at least 100ms worth (3200 bytes = 1600 samples at 16kHz int16)
MIN_AUDIO_BYTES = 3200
audio_buffer = bytearray()
while not self.stop_event.is_set():
try:
# Check for audio
try:
audio_chunk = self.output_queue.get_nowait()
if isinstance(audio_chunk, bytes) and audio_chunk == PIPELINE_END:
if audio_buffer and self.clients:
data = bytes(audio_buffer)
audio_buffer.clear()
await asyncio.gather(
*[client.send(data) for client in self.clients],
return_exceptions=True,
)
break
if isinstance(audio_chunk, bytes) and audio_chunk == AUDIO_RESPONSE_DONE:
if audio_buffer and self.clients:
data = bytes(audio_buffer)
audio_buffer.clear()
await asyncio.gather(
*[client.send(data) for client in self.clients],
return_exceptions=True,
)
self.should_listen.set()
logger.debug("Response complete, listening re-enabled")
continue
if is_control_message(audio_chunk, SESSION_END.kind):
audio_buffer.clear()
continue
if isinstance(audio_chunk, PipelineControlMessage):
continue
if self.clients:
chunk_bytes: bytes
if isinstance(audio_chunk, bytes):
chunk_bytes = audio_chunk
elif isinstance(audio_chunk, np.ndarray):
chunk_bytes = audio_chunk.tobytes()
elif hasattr(audio_chunk, "tobytes"):
chunk_bytes = audio_chunk.tobytes()
else:
continue
audio_buffer.extend(chunk_bytes)
if len(audio_buffer) >= MIN_AUDIO_BYTES:
data = bytes(audio_buffer)
audio_buffer.clear()
logger.debug(f"Sending {len(data)} bytes of audio to {len(self.clients)} client(s)")
await asyncio.gather(
*[client.send(data) for client in self.clients], return_exceptions=True
)
except Empty:
# Flush any buffered audio when queue is empty
if audio_buffer and self.clients:
data = bytes(audio_buffer)
audio_buffer.clear()
logger.debug(f"Flushing {len(data)} bytes of audio to {len(self.clients)} client(s)")
await asyncio.gather(*[client.send(data) for client in self.clients], return_exceptions=True)
# Check for text/tool messages
if self.text_output_queue:
try:
text_message = self.text_output_queue.get_nowait()
if self.clients:
if isinstance(text_message, PipelineEvent):
payload = text_message.model_dump()
await asyncio.gather(
*[client.send(json.dumps(payload)) for client in self.clients],
return_exceptions=True,
)
elif isinstance(text_message, (PipelineControlMessage, bytes)):
continue
except Empty:
pass
await asyncio.sleep(0.01)
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"Send loop error: {e}")