Files
product/main.py
T
2026-03-29 23:01:39 +08:00

637 lines
21 KiB
Python

from __future__ import annotations
import asyncio
import io
import json
import logging
import os
import time
import wave
from dataclasses import dataclass, field
import numpy as np
from aiortc import RTCPeerConnection, RTCSessionDescription
from dotenv import load_dotenv
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from config import settings
from core.pipeline import ChatPipeline
from core.state_machine import Arbitrator, SessionState
from services.asr import ASRService
from services.avatar import AvatarService
from services.llm import LLMService
from services.tts import TTSService
from services.vad import VADService
from webrtc.tracks import AudioBus, AvatarAudioTrack
load_dotenv()
if settings.hf_token:
os.environ["HF_TOKEN"] = settings.hf_token
logging.basicConfig(level=logging.INFO)
app = FastAPI(title="Visual Voice Chat")
app.mount("/web", StaticFiles(directory="web"), name="web")
app.add_middleware(
CORSMiddleware,
allow_origins=[o.strip() for o in settings.cors_origins.split(",")] if settings.cors_origins else ["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
arbitrator = Arbitrator()
asr_service = ASRService()
llm_service = LLMService()
tts_service = TTSService()
vad_service = VADService()
avatar_service = AvatarService()
pipeline = ChatPipeline(arbitrator, asr_service, llm_service, tts_service)
audio_bus = AudioBus(sample_rate=48000)
pcs: set[RTCPeerConnection] = set()
subtitle_clients: set[WebSocket] = set()
animation_clients: set[WebSocket] = set()
audio_clients: set[WebSocket] = set()
MAX_BUFFERED_AUDIO_SAMPLES = 16000 * 45
MIN_VOICE_TURN_SAMPLES = 16000 // 2
MIN_VOICE_RMS = 0.012
@dataclass
class RuntimeState:
last_user_text: str = ""
last_reply_text: str = ""
pipeline_runs: int = 0
last_latency_ms: int = 0
last_asr_latency_ms: int = 0
last_llm_latency_ms: int = 0
last_llm_source: str = "unknown"
last_tts_latency_ms: int = 0
last_tts_first_chunk_ms: int = 0
last_barge_in_ms: int = 0
barge_in_count: int = 0
last_input_mode: str = "none"
vad_start_count: int = 0
vad_end_count: int = 0
last_animation_mode: str = settings.avatar_driver_mode
last_animation_frame_count: int = 0
busy: bool = False
pending_voice_turn: bool = False
buffer_16k: list[np.ndarray] = field(default_factory=list)
buffered_16k_samples: int = 0
runtime = RuntimeState()
class TextChatRequest(BaseModel):
text: str
async def _broadcast_json(clients: set[WebSocket], payload: dict) -> None:
stale: list[WebSocket] = []
message = json.dumps(payload, ensure_ascii=False)
for ws in list(clients):
try:
await ws.send_text(message)
except Exception:
stale.append(ws)
for ws in stale:
clients.discard(ws)
async def _broadcast_subtitle(
role: str,
text: str,
source: str,
*,
partial: bool = False,
final: bool = True,
) -> None:
if not text:
return
await _broadcast_json(
subtitle_clients,
{
"role": role,
"text": text,
"source": source,
"partial": partial,
"final": final,
"ts_ms": int(time.time() * 1000),
},
)
async def _broadcast_animation(payload: dict) -> None:
runtime.last_animation_mode = str(payload.get("driver", settings.avatar_driver_mode))
runtime.last_animation_frame_count = int(payload.get("frame_count", 0))
await _broadcast_json(animation_clients, payload)
def _runtime_payload() -> dict:
return {
"state": str(arbitrator.state),
"peers": len(pcs),
"audio_clients": len(audio_clients),
"subtitle_clients": len(subtitle_clients),
"animation_clients": len(animation_clients),
"pipeline_runs": runtime.pipeline_runs,
"last_latency_ms": runtime.last_latency_ms,
"last_asr_latency_ms": runtime.last_asr_latency_ms,
"last_llm_latency_ms": runtime.last_llm_latency_ms,
"last_llm_source": runtime.last_llm_source,
"last_tts_latency_ms": runtime.last_tts_latency_ms,
"last_tts_first_chunk_ms": runtime.last_tts_first_chunk_ms,
"last_barge_in_ms": runtime.last_barge_in_ms,
"barge_in_count": runtime.barge_in_count,
"last_input_mode": runtime.last_input_mode,
"vad_start_count": runtime.vad_start_count,
"vad_end_count": runtime.vad_end_count,
"last_user_text": runtime.last_user_text,
"last_reply_text": runtime.last_reply_text,
"last_animation_mode": runtime.last_animation_mode,
"last_animation_frame_count": runtime.last_animation_frame_count,
"pipeline_busy": runtime.busy,
"llm": llm_service.health,
"asr": asr_service.health,
"tts": tts_service.health,
"vad": vad_service.health,
"avatar": avatar_service.health,
}
def _clear_buffered_audio() -> None:
runtime.buffer_16k.clear()
runtime.buffered_16k_samples = 0
def _append_buffered_audio(chunk: np.ndarray) -> None:
runtime.buffer_16k.append(chunk)
runtime.buffered_16k_samples += int(chunk.shape[0])
while runtime.buffered_16k_samples > MAX_BUFFERED_AUDIO_SAMPLES and runtime.buffer_16k:
dropped = runtime.buffer_16k.pop(0)
runtime.buffered_16k_samples -= int(dropped.shape[0])
def _resample_mono(audio: np.ndarray, from_sr: int, to_sr: int) -> np.ndarray:
if from_sr == to_sr:
return audio.astype(np.float32, copy=False)
in_len = audio.shape[0]
out_len = max(1, int(in_len * to_sr / from_sr))
x_old = np.linspace(0.0, 1.0, in_len, endpoint=False)
x_new = np.linspace(0.0, 1.0, out_len, endpoint=False)
return np.interp(x_new, x_old, audio).astype(np.float32)
def _resample_to_16k_mono(pcm: np.ndarray, sample_rate: int, channels: int) -> np.ndarray:
if pcm.ndim == 2:
mono = pcm.mean(axis=0)
else:
mono = pcm
mono = mono.astype(np.float32) / 32768.0
if sample_rate == 16000:
return mono
in_len = mono.shape[0]
out_len = max(1, int(in_len * 16000 / sample_rate))
x_old = np.linspace(0.0, 1.0, in_len, endpoint=False)
x_new = np.linspace(0.0, 1.0, out_len, endpoint=False)
return np.interp(x_new, x_old, mono).astype(np.float32)
def _encode_wav_bytes(audio: np.ndarray, sample_rate: int) -> bytes:
pcm = np.clip(audio.astype(np.float32, copy=False) * 32767.0, -32768, 32767).astype(np.int16)
with io.BytesIO() as buffer:
with wave.open(buffer, "wb") as wav_file:
wav_file.setnchannels(1)
wav_file.setsampwidth(2)
wav_file.setframerate(sample_rate)
wav_file.writeframes(pcm.tobytes())
return buffer.getvalue()
async def _broadcast_audio_chunk(audio: np.ndarray, sample_rate: int) -> None:
if not audio_clients:
return
stale: list[WebSocket] = []
payload = _encode_wav_bytes(audio, sample_rate)
for ws in list(audio_clients):
try:
await ws.send_bytes(payload)
except Exception:
stale.append(ws)
for ws in stale:
audio_clients.discard(ws)
async def _broadcast_audio_reset(reason: str) -> None:
await _broadcast_json(
audio_clients,
{
"type": "audio_reset",
"reason": reason,
"ts_ms": int(time.time() * 1000),
},
)
async def _enqueue_audio_and_animation(
seg_audio: np.ndarray,
seg_sr: int,
seg_text: str,
idx: int,
total: int,
) -> None:
await _broadcast_audio_chunk(seg_audio, seg_sr)
animation_payload = await avatar_service.build_controls_from_audio(
seg_audio,
seg_sr,
text=seg_text,
chunk_index=idx,
total_chunks=total,
)
await _broadcast_animation(animation_payload)
async def _run_pipeline_from_buffer() -> None:
if runtime.busy or not runtime.buffer_16k or not runtime.pending_voice_turn:
return
runtime.busy = True
runtime.pending_voice_turn = False
start_t = time.perf_counter()
try:
audio_16k = np.concatenate(runtime.buffer_16k, axis=0)
_clear_buffered_audio()
if audio_16k.shape[0] < MIN_VOICE_TURN_SAMPLES:
return
rms = float(np.sqrt(np.mean(audio_16k.astype(np.float32) ** 2))) if audio_16k.size else 0.0
if rms < MIN_VOICE_RMS:
return
runtime.last_input_mode = "voice"
async def on_user_text_now(user_text: str) -> None:
runtime.last_user_text = user_text
await _broadcast_subtitle("user", user_text, "voice")
async def on_reply_text_now(reply_text: str) -> None:
runtime.last_reply_text = reply_text
async def on_reply_segment_now(seg_text: str, idx: int, total: int) -> None:
await _broadcast_subtitle("ai", seg_text, "voice", partial=True, final=(idx >= total - 1))
result, audio, sr = await pipeline.process_turn(
audio_16k,
on_audio_chunk=_enqueue_audio_and_animation,
on_user_text=on_user_text_now,
on_reply_text=on_reply_text_now,
on_reply_segment=on_reply_segment_now,
)
if not result["user_text"]:
return
runtime.last_user_text = result["user_text"]
runtime.last_reply_text = result["reply_text"]
runtime.pipeline_runs += 1
runtime.last_latency_ms = int((time.perf_counter() - start_t) * 1000)
runtime.last_asr_latency_ms = int(result.get("asr_latency_ms", 0))
runtime.last_llm_latency_ms = int(result.get("llm_latency_ms", 0))
runtime.last_llm_source = str(result.get("llm_source", "unknown"))
runtime.last_tts_latency_ms = int(result.get("tts_latency_ms", 0))
runtime.last_tts_first_chunk_ms = int(result.get("tts_first_chunk_ms", 0))
if result.get("audio_samples", 0) <= 1:
await _broadcast_audio_chunk(audio, sr)
await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
finally:
runtime.busy = False
if runtime.pending_voice_turn and runtime.buffer_16k:
asyncio.create_task(_run_pipeline_from_buffer())
async def _process_user_audio_chunk(mono_16k: np.ndarray, vad_buffer: np.ndarray) -> np.ndarray:
_append_buffered_audio(mono_16k)
vad_buffer = np.concatenate([vad_buffer, mono_16k], axis=0)
while vad_buffer.shape[0] >= 512:
chunk = vad_buffer[:512]
vad_buffer = vad_buffer[512:]
try:
event = vad_service.push_chunk(chunk)
except Exception as exc:
logging.exception("VAD chunk processing failed: %s", exc)
event = None
if event and "start" in event:
barge_t0 = time.perf_counter()
was_avatar_speaking = arbitrator.state == SessionState.AVATAR_SPEAKING
await arbitrator.on_speech_start()
await audio_bus.clear()
await _broadcast_audio_reset("speech-start")
await _broadcast_animation(await avatar_service.build_reset_payload(reason="speech-start"))
if was_avatar_speaking:
runtime.last_barge_in_ms = int((time.perf_counter() - barge_t0) * 1000)
runtime.barge_in_count += 1
runtime.vad_start_count += 1
if event and "end" in event:
await arbitrator.on_speech_end()
runtime.vad_end_count += 1
runtime.pending_voice_turn = True
asyncio.create_task(_run_pipeline_from_buffer())
return vad_buffer
async def _consume_user_audio(track) -> None:
vad_buffer = np.zeros(0, dtype=np.float32)
while True:
frame = await track.recv()
pcm = frame.to_ndarray()
sample_rate = getattr(frame, "sample_rate", 48000) or 48000
layout = getattr(frame, "layout", None)
channels = len(layout.channels) if layout and layout.channels else 1
mono_16k = _resample_to_16k_mono(pcm, sample_rate, channels)
vad_buffer = await _process_user_audio_chunk(mono_16k, vad_buffer)
@app.get("/health")
async def health():
return _runtime_payload()
@app.get("/meta")
async def meta():
return {
"stun_url": settings.stun_url,
"https_enabled": bool(settings.ssl_certfile and settings.ssl_keyfile),
"host": settings.webrtc_host,
"port": settings.webrtc_port,
"avatar_protocol": settings.avatar_control_protocol,
}
@app.get("/avatar/schema")
async def avatar_schema():
return avatar_service.schema
@app.get("/events")
async def events():
async def gen():
while True:
yield f"data: {json.dumps(_runtime_payload(), ensure_ascii=False)}\n\n"
await asyncio.sleep(1.0)
return StreamingResponse(gen(), media_type="text/event-stream")
@app.websocket("/ws/subtitles")
async def ws_subtitles(websocket: WebSocket):
await websocket.accept()
subtitle_clients.add(websocket)
try:
while True:
await websocket.send_text(json.dumps({"type": "ping"}, ensure_ascii=False))
await asyncio.sleep(15)
except WebSocketDisconnect:
subtitle_clients.discard(websocket)
except Exception:
subtitle_clients.discard(websocket)
@app.websocket("/ws/animation")
async def ws_animation(websocket: WebSocket):
await websocket.accept()
animation_clients.add(websocket)
try:
await websocket.send_text(json.dumps({"type": "animation_ready", **avatar_service.schema}, ensure_ascii=False))
await websocket.send_text(json.dumps(await avatar_service.build_idle_payload(reason="client-connected"), ensure_ascii=False))
while True:
await websocket.send_text(json.dumps({"type": "ping"}, ensure_ascii=False))
await asyncio.sleep(15)
except WebSocketDisconnect:
animation_clients.discard(websocket)
except Exception:
animation_clients.discard(websocket)
@app.websocket("/ws/audio")
async def ws_audio(websocket: WebSocket):
await websocket.accept()
replaced_previous = len(audio_clients) > 0
if replaced_previous:
old_clients = list(audio_clients)
await asyncio.gather(*(client.close(code=1012, reason="replaced") for client in old_clients), return_exceptions=True)
for client in old_clients:
audio_clients.discard(client)
audio_clients.add(websocket)
sample_rate = 48000
channels = 1
vad_buffer = np.zeros(0, dtype=np.float32)
await websocket.send_text(
json.dumps(
{
"type": "audio_ready",
"transport": "websocket_audio",
"replaced_previous": replaced_previous,
},
ensure_ascii=False,
)
)
try:
while True:
message = await websocket.receive()
message_type = message.get("type")
if message_type == "websocket.disconnect":
break
text = message.get("text")
if text is not None:
try:
payload = json.loads(text)
except json.JSONDecodeError:
payload = {}
event_type = str(payload.get("type", ""))
if event_type == "audio_config":
sample_rate = max(8000, int(payload.get("sample_rate", sample_rate) or sample_rate))
channels = max(1, int(payload.get("channels", channels) or channels))
elif event_type == "audio_reset":
_clear_buffered_audio()
runtime.pending_voice_turn = False
await audio_bus.clear()
await _broadcast_audio_reset("client-reset")
elif event_type == "ping":
await websocket.send_text(json.dumps({"type": "pong"}, ensure_ascii=False))
continue
data = message.get("bytes")
if not data:
continue
pcm = np.frombuffer(data, dtype=np.int16)
if pcm.size == 0:
continue
mono_16k = _resample_to_16k_mono(pcm, sample_rate, channels)
vad_buffer = await _process_user_audio_chunk(mono_16k, vad_buffer)
except WebSocketDisconnect:
pass
finally:
audio_clients.discard(websocket)
@app.post("/demo/run-once")
async def run_once():
fake_audio = np.zeros(16000 * 3, dtype=np.float32)
await arbitrator.on_speech_start()
result = await pipeline.run_once(fake_audio)
return result
@app.post("/chat/text")
async def chat_text(req: TextChatRequest):
text = req.text.strip()
if not text:
return JSONResponse({"ok": False, "message": "text is empty"}, status_code=400)
if runtime.busy:
return JSONResponse({"ok": False, "message": "pipeline busy"}, status_code=429)
runtime.busy = True
start_t = time.perf_counter()
try:
async def on_reply_text_now(reply_text: str) -> None:
runtime.last_reply_text = reply_text
async def on_reply_segment_now(seg_text: str, idx: int, total: int) -> None:
await _broadcast_subtitle("ai", seg_text, "text", partial=True, final=(idx >= total - 1))
await _broadcast_subtitle("user", text, "text")
await _broadcast_audio_reset("text-turn-start")
await _broadcast_animation(await avatar_service.build_reset_payload(reason="text-turn-start"))
result, audio, sr = await pipeline.process_text_turn(
text,
on_audio_chunk=_enqueue_audio_and_animation,
on_reply_text=on_reply_text_now,
on_reply_segment=on_reply_segment_now,
)
runtime.last_input_mode = "text"
runtime.last_user_text = result["user_text"]
runtime.last_reply_text = result["reply_text"]
runtime.pipeline_runs += 1
runtime.last_latency_ms = int((time.perf_counter() - start_t) * 1000)
runtime.last_asr_latency_ms = 0
runtime.last_llm_latency_ms = int(result.get("llm_latency_ms", 0))
runtime.last_llm_source = str(result.get("llm_source", "unknown"))
runtime.last_tts_latency_ms = int(result.get("tts_latency_ms", 0))
runtime.last_tts_first_chunk_ms = int(result.get("tts_first_chunk_ms", 0))
if result.get("audio_samples", 0) <= 1:
await _broadcast_audio_chunk(audio, sr)
await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
return {"ok": True, **result}
finally:
runtime.busy = False
@app.post("/chat/reset")
async def chat_reset():
llm_service.clear_history()
_clear_buffered_audio()
runtime.pending_voice_turn = False
runtime.last_user_text = ""
runtime.last_reply_text = ""
runtime.pipeline_runs = 0
runtime.last_latency_ms = 0
runtime.last_asr_latency_ms = 0
runtime.last_llm_latency_ms = 0
runtime.last_llm_source = "unknown"
runtime.last_tts_latency_ms = 0
runtime.last_tts_first_chunk_ms = 0
runtime.last_barge_in_ms = 0
runtime.barge_in_count = 0
runtime.last_input_mode = "none"
runtime.vad_start_count = 0
runtime.vad_end_count = 0
runtime.last_animation_frame_count = 0
await audio_bus.clear()
await _broadcast_audio_reset("chat-reset")
await _broadcast_animation(await avatar_service.build_reset_payload(reason="chat-reset"))
return {"ok": True}
@app.get("/demo/frame-idle-shape")
async def frame_idle_shape():
return await avatar_service.build_idle_payload(reason="demo")
@app.get("/")
async def root():
return FileResponse("web/index.html")
@app.post("/webrtc/offer")
async def webrtc_offer(request: Request):
params = await request.json()
offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
replaced_previous = len(pcs) > 0
if replaced_previous:
old_peers = list(pcs)
await asyncio.gather(*(peer.close() for peer in old_peers), return_exceptions=True)
for peer in old_peers:
pcs.discard(peer)
pc = RTCPeerConnection()
pcs.add(pc)
@pc.on("connectionstatechange")
async def on_connectionstatechange():
if pc.connectionState in {"failed", "closed", "disconnected"}:
await pc.close()
pcs.discard(pc)
@pc.on("track")
def on_track(track):
if track.kind == "audio":
asyncio.create_task(_consume_user_audio(track))
pc.addTrack(AvatarAudioTrack(audio_bus=audio_bus))
await pc.setRemoteDescription(offer)
answer = await pc.createAnswer()
await pc.setLocalDescription(answer)
return JSONResponse(
{
"sdp": pc.localDescription.sdp,
"type": pc.localDescription.type,
"replaced_previous": replaced_previous,
}
)
@app.on_event("shutdown")
async def on_shutdown():
await asyncio.gather(*(pc.close() for pc in list(pcs)), return_exceptions=True)
pcs.clear()
if __name__ == "__main__":
import uvicorn
uvicorn_kwargs = {
"app": app,
"host": settings.webrtc_host,
"port": settings.webrtc_port,
}
if settings.ssl_certfile and settings.ssl_keyfile:
uvicorn_kwargs["ssl_certfile"] = settings.ssl_certfile
uvicorn_kwargs["ssl_keyfile"] = settings.ssl_keyfile
uvicorn.run(**uvicorn_kwargs)