234 lines
8.8 KiB
Python
234 lines
8.8 KiB
Python
"""Voice service — STT and TTS over one WebSocket.
|
|
|
|
This process holds the GPU models and nothing else. It has no idea what the
|
|
CRM is: it receives audio and returns text, receives text and returns audio.
|
|
All orchestration, auth and business logic stay in the Node service, so voice
|
|
is just another channel into the same agent rather than a parallel system.
|
|
|
|
Protocol (ws /ws/voice), JSON control + binary audio:
|
|
|
|
client → server
|
|
binary 16 kHz mono PCM16 mic frames
|
|
{"type":"config","lang":"ta"} set the session language
|
|
{"type":"speak","text":"…","id":"…"} synthesise
|
|
{"type":"cancel"} stop speaking now (barge-in)
|
|
{"type":"reset"} clear the endpointer
|
|
|
|
server → client
|
|
{"type":"ready", …}
|
|
{"type":"speech_start"} VAD opened a turn → caller ducks TTS
|
|
{"type":"transcript","text":…} a finished utterance
|
|
{"type":"audio_start","id":…,"sample_rate":24000}
|
|
binary float32 mono TTS chunks
|
|
{"type":"audio_end","id":…}
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
|
|
import numpy as np
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
|
|
from .config import settings
|
|
from .stt import SUPPORTED, Transcriber
|
|
from .tts import Synthesizer, split_sentences
|
|
from .vad import Endpointer
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)-5s %(message)s", datefmt="%H:%M:%S")
|
|
logger = logging.getLogger("voice")
|
|
|
|
app = FastAPI(title="WeLe Voice Service")
|
|
|
|
stt = Transcriber()
|
|
tts = Synthesizer()
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def _startup() -> None:
|
|
t0 = time.perf_counter()
|
|
logger.info("loading models on %s…", settings.device)
|
|
stt.load()
|
|
tts.load()
|
|
if settings.warmup:
|
|
stt.warmup()
|
|
tts.warmup()
|
|
logger.info("voice service ready in %.1fs", time.perf_counter() - t0)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict:
|
|
import torch
|
|
|
|
return {
|
|
"ok": True,
|
|
"device": settings.device,
|
|
"stt_model": settings.stt_model,
|
|
"tts_model": settings.tts_model,
|
|
"tts_sample_rate": tts.sample_rate,
|
|
"languages": SUPPORTED,
|
|
"vram_gb": round(torch.cuda.memory_reserved() / 1e9, 2) if settings.device == "cuda" else None,
|
|
}
|
|
|
|
|
|
@app.get("/languages")
|
|
async def languages() -> dict:
|
|
return {"languages": SUPPORTED}
|
|
|
|
|
|
class Session:
|
|
"""One browser connection. Owns its endpointer and its speaking state."""
|
|
|
|
def __init__(self, ws: WebSocket) -> None:
|
|
self.ws = ws
|
|
# "auto" detects per utterance; `prefer` breaks ties when the detector
|
|
# is unsure, which is common on short code-mixed ("Tanglish") speech.
|
|
self.lang = "auto"
|
|
self.prefer = settings.default_lang
|
|
self.endpointer = Endpointer()
|
|
self.endpointer.load()
|
|
self._speak_task: asyncio.Task | None = None
|
|
self._cancel = asyncio.Event()
|
|
|
|
async def send(self, payload: dict) -> None:
|
|
await self.ws.send_text(json.dumps(payload, ensure_ascii=False))
|
|
|
|
# ── microphone ───────────────────────────────────────────────────────────
|
|
async def on_audio(self, raw: bytes) -> None:
|
|
pcm = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
|
|
loop = asyncio.get_running_loop()
|
|
# VAD is a small torch model but still blocking; keep the event loop free.
|
|
utterances, started = await loop.run_in_executor(None, self.endpointer.push, pcm)
|
|
|
|
if started:
|
|
# Barge-in: the user talking wins immediately.
|
|
await self.stop_speaking()
|
|
await self.send({"type": "speech_start"})
|
|
|
|
for utt in utterances:
|
|
request_lang = f"auto:{self.prefer}" if self.lang == "auto" else self.lang
|
|
result = await loop.run_in_executor(None, stt.transcribe, utt.audio, request_lang)
|
|
if result["text"]:
|
|
await self.send({"type": "transcript", **result, "truncated": utt.truncated})
|
|
else:
|
|
await self.send({"type": "transcript_empty", "reason": result.get("note", "no speech")})
|
|
|
|
# ── speaking ─────────────────────────────────────────────────────────────
|
|
async def speak(self, text: str, msg_id: str, lang: str | None = None, voice: str | None = None) -> None:
|
|
await self.stop_speaking()
|
|
self._cancel.clear()
|
|
self._speak_task = asyncio.create_task(self._speak(text, msg_id, lang or self.lang, voice))
|
|
|
|
async def _speak(self, text: str, msg_id: str, lang: str, voice: str | None) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
try:
|
|
await self.send({"type": "audio_start", "id": msg_id, "sample_rate": tts.sample_rate})
|
|
|
|
# Sentence at a time: shorter prompts reach first audio sooner, and
|
|
# a boundary is a clean place to stop when interrupted.
|
|
for sentence in split_sentences(text):
|
|
if self._cancel.is_set():
|
|
break
|
|
queue: asyncio.Queue = asyncio.Queue(maxsize=32)
|
|
|
|
def produce() -> None:
|
|
try:
|
|
for chunk in tts.stream(sentence, lang, voice):
|
|
if self._cancel.is_set():
|
|
break
|
|
asyncio.run_coroutine_threadsafe(queue.put(chunk), loop).result()
|
|
finally:
|
|
asyncio.run_coroutine_threadsafe(queue.put(None), loop).result()
|
|
|
|
loop.run_in_executor(None, produce)
|
|
while True:
|
|
chunk = await queue.get()
|
|
if chunk is None:
|
|
break
|
|
if self._cancel.is_set():
|
|
continue # drain, don't send
|
|
await self.ws.send_bytes(np.asarray(chunk, dtype=np.float32).tobytes())
|
|
|
|
await self.send({"type": "audio_end", "id": msg_id, "cancelled": self._cancel.is_set()})
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception as e: # noqa: BLE001
|
|
logger.exception("synthesis failed")
|
|
try:
|
|
await self.send({"type": "error", "where": "tts", "message": str(e)[:200]})
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
|
|
async def stop_speaking(self) -> None:
|
|
if self._speak_task and not self._speak_task.done():
|
|
self._cancel.set()
|
|
try:
|
|
await asyncio.wait_for(self._speak_task, timeout=2.0)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
self._speak_task.cancel()
|
|
self._speak_task = None
|
|
|
|
|
|
@app.websocket("/ws/voice")
|
|
async def voice(ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
session = Session(ws)
|
|
await session.send({
|
|
"type": "ready",
|
|
"sample_rate_in": settings.sample_rate_in,
|
|
"sample_rate_out": tts.sample_rate,
|
|
"languages": SUPPORTED,
|
|
})
|
|
logger.info("voice session opened")
|
|
|
|
try:
|
|
while True:
|
|
msg = await ws.receive()
|
|
if msg["type"] == "websocket.disconnect":
|
|
break
|
|
|
|
if (raw := msg.get("bytes")) is not None:
|
|
await session.on_audio(raw)
|
|
continue
|
|
|
|
if (text := msg.get("text")) is None:
|
|
continue
|
|
try:
|
|
data = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
kind = data.get("type")
|
|
if kind == "config":
|
|
session.lang = data.get("lang", session.lang)
|
|
if session.lang != "auto":
|
|
session.prefer = session.lang
|
|
elif data.get("prefer"):
|
|
session.prefer = data["prefer"]
|
|
session.endpointer.reset()
|
|
await session.send({"type": "config_ok", "lang": session.lang, "prefer": session.prefer})
|
|
elif kind == "speak":
|
|
await session.speak(data.get("text", ""), data.get("id", ""), data.get("lang"), data.get("voice"))
|
|
elif kind == "cancel":
|
|
await session.stop_speaking()
|
|
await session.send({"type": "cancelled"})
|
|
elif kind == "reset":
|
|
session.endpointer.reset()
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
await session.stop_speaking()
|
|
logger.info("voice session closed")
|
|
|
|
|
|
def main() -> None:
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host=settings.host, port=settings.port, log_level="info", ws_max_size=16 * 1024 * 1024)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|