239 lines
9.9 KiB
Python
239 lines
9.9 KiB
Python
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}")
|