first
This commit is contained in:
@@ -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}")
|
||||
Reference in New Issue
Block a user