from __future__ import annotations import asyncio import io import json import logging import time import wave from dataclasses import dataclass, field import numpy as np from dotenv import load_dotenv from fastapi import FastAPI, 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 load_dotenv() 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) 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 _replace_existing_clients(clients: set[WebSocket]) -> bool: replaced_previous = len(clients) > 0 if not replaced_previous: return False old_clients = list(clients) await asyncio.gather(*(client.close(code=1012, reason="replaced") for client in old_clients), return_exceptions=True) for client in old_clients: clients.discard(client) return True 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), "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 _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 @app.get("/health") async def health(): return _runtime_payload() @app.get("/meta") async def meta(): return { "https_enabled": bool(settings.ssl_certfile and settings.ssl_keyfile), "host": settings.http_host, "port": settings.http_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() await _replace_existing_clients(subtitle_clients) 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() await _replace_existing_clients(animation_clients) 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 = await _replace_existing_clients(audio_clients) 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 _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 _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") if __name__ == "__main__": import uvicorn uvicorn_kwargs = { "app": app, "host": settings.http_host, "port": settings.http_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)