first
This commit is contained in:
@@ -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