main.py

import asyncio
import logging
import re
from collections import deque
from contextlib import asynccontextmanager, suppress
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from mlx_audio.stt.utils import load_model as load_asr_model
from mlx_audio.tts.utils import load_model as load_tts_model
from mlx_vlm import apply_chat_template, load as load_vlm, stream_generate
from silero_vad import VADIterator, load_silero_vad

# Constants for Models & Voices
BASE_DIR = Path(__file__).resolve().parent
WEB_DIR = BASE_DIR / "web"
INDEX_HTML_PATH = WEB_DIR / "index.html"
TTS_MODEL_PATH = "mlx-community/Qwen3-TTS-12Hz-0.6B-Base-8bit"
ASR_MODEL_PATH = "mlx-community/Qwen3-ASR-0.6B-4bit"
VLM_MODEL_PATH = "mlx-community/gemma-4-e2b-it-4bit"
REF_AUDIO_PATH = BASE_DIR / "custom_voice.wav"
REF_TEXT = "为您提供高效准确的信息是我的职责。请您随时提出需求，我将会为您进行最深度的分析和解答。"
INPUT_SAMPLE_RATE = 16000
MAX_CHAT_MESSAGES = 10
BYTES_PER_PCM16_SAMPLE = 2
VAD_FRAME_SAMPLES = 512
VAD_FRAME_BYTES = VAD_FRAME_SAMPLES * BYTES_PER_PCM16_SAMPLE
PARTIAL_TRANSCRIBE_INTERVAL_BYTES = INPUT_SAMPLE_RATE  # 16-bit mono PCM at 16 kHz: ~0.5s

# Global context variables
tts_model = None
asr_model = None
vad_model = None
vlm_model = None
vlm_processor = None
output_sample_rate = INPUT_SAMPLE_RATE

logger = logging.getLogger("mlx_live")
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(
        logging.Formatter(
            fmt="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
            datefmt="%H:%M:%S",
        )
    )
    logger.addHandler(handler)
logger.setLevel(logging.INFO)
logger.propagate = False


@dataclass
class ConnectionSession:
    websocket: WebSocket
    send_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    chat_history: list[dict[str, str]] = field(default_factory=list)
    response_task: asyncio.Task | None = None
    response_token: int = 0
    interrupt_requested: bool = False

    def ensure_response_active(self, response_token: int | None = None) -> None:
        if response_token is not None and response_token != self.response_token:
            raise asyncio.CancelledError

    async def send_json(self, payload: Any, response_token: int | None = None) -> None:
        self.ensure_response_active(response_token)
        async with self.send_lock:
            self.ensure_response_active(response_token)
            await self.websocket.send_json(payload)
        self.ensure_response_active(response_token)

    async def send_bytes(self, payload: bytes, response_token: int | None = None) -> None:
        self.ensure_response_active(response_token)
        async with self.send_lock:
            self.ensure_response_active(response_token)
            await self.websocket.send_bytes(payload)
        self.ensure_response_active(response_token)


def trim_chat_history(chat_history: list[dict[str, str]]) -> None:
    while len(chat_history) > MAX_CHAT_MESSAGES:
        del chat_history[: min(2, len(chat_history))]


def preserve_interrupted_turn(
    chat_history: list[dict[str, str]], recognized_text: str, llm_response_text: str
) -> None:
    partial_response = llm_response_text.strip()
    if partial_response:
        if not chat_history or chat_history[-1].get("role") != "assistant":
            chat_history.append({"role": "assistant", "content": partial_response})
            trim_chat_history(chat_history)
        return

    if (
        chat_history
        and chat_history[-1].get("role") == "user"
        and chat_history[-1].get("content") == recognized_text
    ):
        chat_history.pop()


async def interrupt_active_response(session: ConnectionSession) -> None:
    task = session.response_task
    if task is None or task.done() or session.interrupt_requested:
        return

    session.interrupt_requested = True
    session.response_token += 1
    logger.info("🛑 检测到用户插话，打断当前回复...")

    with suppress(RuntimeError, WebSocketDisconnect):
        await session.send_json({"type": "interrupt"})
        await session.send_json({"type": "state", "phase": "idle"})

    task.cancel()


def launch_response_task(session: ConnectionSession, recognized_text: str) -> None:
    session.interrupt_requested = False
    session.response_token += 1
    response_token = session.response_token
    session.response_task = asyncio.create_task(
        run_response_task(session, recognized_text, response_token)
    )


@asynccontextmanager
async def lifespan(app: FastAPI):
    global tts_model, asr_model, vad_model, vlm_model, vlm_processor, output_sample_rate

    if not REF_AUDIO_PATH.exists():
        raise FileNotFoundError(f"Reference audio not found: {REF_AUDIO_PATH}")

    logger.info("⏳ 正在加载 TTS 模型...")
    tts_model = load_tts_model(TTS_MODEL_PATH)
    output_sample_rate = tts_model.sample_rate
    logger.info("✅ TTS 模型加载完成 (sample_rate=%s)", output_sample_rate)

    logger.info("⏳ 正在加载 ASR 模型...")
    asr_model = load_asr_model(ASR_MODEL_PATH)
    logger.info("✅ ASR 模型加载完成")

    logger.info("⏳ 正在加载 VAD 模型...")
    vad_model = load_silero_vad()
    logger.info("✅ VAD 模型加载完成")

    logger.info("⏳ 正在加载 VLM 模型...")
    vlm_model, vlm_processor = load_vlm(VLM_MODEL_PATH)
    logger.info("✅ VLM 模型加载完成")
    logger.info("🚀 所有模型准备完毕！服务端已就绪。")

    yield


app = FastAPI(lifespan=lifespan)

# Ensure web directory exists
WEB_DIR.mkdir(exist_ok=True)
app.mount("/static", StaticFiles(directory=str(WEB_DIR)), name="static")


@app.get("/")
async def get():
    with INDEX_HTML_PATH.open("r", encoding="utf-8") as f:
        return HTMLResponse(f.read())


async def process_llm_and_tts(
    session: ConnectionSession, recognized_text: str, response_token: int
) -> None:
    """处理 LLM 推理与 TTS 流式合成序列"""
    chat_history = session.chat_history
    llm_response_text = ""
    assistant_committed = False

    logger.info("📝 用户: %s", recognized_text)
    await session.send_json({"type": "asr", "text": "🙎‍♂️ " + recognized_text}, response_token)
    await session.send_json({"type": "state", "phase": "generating"}, response_token)

    logger.info("🧠 正在生成回复...")
    chat_history.append({"role": "user", "content": recognized_text})
    trim_chat_history(chat_history)

    prompt = apply_chat_template(vlm_processor, getattr(vlm_model, "config", None), chat_history)

    await session.send_json({"type": "asr", "text": "\n🤖 "}, response_token)
    sentence_buffer = ""
    audio_queue = asyncio.Queue()

    async def tts_worker() -> None:
        while True:
            sentence = await audio_queue.get()
            if sentence is None:
                break

            filtered_sentence = re.sub(
                r'[^\u4e00-\u9fa5a-zA-Z0-9\s，。！？、；：“”‘’（）《》·.,!?:;\'"()\-&@$%+=/]',
                "",
                sentence,
            ).strip()
            if not filtered_sentence:
                continue

            logger.info("🔊 推送合成片段: %s", filtered_sentence)
            generator = tts_model.generate(
                text=filtered_sentence,
                lang_code="zh",
                ref_audio=str(REF_AUDIO_PATH),
                ref_text=REF_TEXT,
                stream=True,
                streaming_interval=0.5,
                verbose=False,
            )
            for result in generator:
                audio_np = np.asarray(result.audio, dtype=np.float32)
                await session.send_bytes(audio_np.tobytes(), response_token)
                await asyncio.sleep(0.001)

    tts_task = asyncio.create_task(tts_worker())
    delimiters = r'([,，。！？.!?\n]+)'
    aggregate_buffer = ""
    min_tts_chunk_len = 25

    try:
        for chunk in stream_generate(vlm_model, vlm_processor, prompt, max_tokens=200):
            session.ensure_response_active(response_token)
            chunk_text = getattr(chunk, "text", str(chunk))
            llm_response_text += chunk_text
            sentence_buffer += chunk_text

            await session.send_json({"type": "llm_partial", "text": llm_response_text}, response_token)
            await asyncio.sleep(0.01)

            parts = re.split(delimiters, sentence_buffer)
            if len(parts) > 1:
                for i in range(0, len(parts) - 1, 2):
                    sent = parts[i]
                    delim = parts[i + 1] if i + 1 < len(parts) else ""
                    complete_sentence = sent + delim
                    aggregate_buffer += complete_sentence

                    if len(aggregate_buffer) >= min_tts_chunk_len:
                        if aggregate_buffer.strip():
                            audio_queue.put_nowait(aggregate_buffer)
                        aggregate_buffer = ""
                sentence_buffer = parts[-1]

        logger.info("🤖 回复内容: %s", llm_response_text)
        logger.info("✅ 回复生成完毕")
        await session.send_json({"type": "asr", "text": llm_response_text + "\n"}, response_token)

        if aggregate_buffer.strip():
            sentence_buffer = aggregate_buffer + sentence_buffer
        if sentence_buffer.strip():
            audio_queue.put_nowait(sentence_buffer)

        audio_queue.put_nowait(None)
        await tts_task

        if llm_response_text.strip():
            chat_history.append({"role": "assistant", "content": llm_response_text})
            trim_chat_history(chat_history)
            assistant_committed = True

        await session.send_json({"type": "state", "phase": "idle"}, response_token)
        await session.send_json({"type": "done"}, response_token)
        logger.info("✅ 语音交互回合结束")
    except asyncio.CancelledError:
        if not assistant_committed:
            preserve_interrupted_turn(chat_history, recognized_text, llm_response_text)
        raise
    except Exception:
        if not assistant_committed:
            preserve_interrupted_turn(chat_history, recognized_text, llm_response_text)
        raise
    finally:
        if not tts_task.done():
            tts_task.cancel()
            with suppress(asyncio.CancelledError):
                await tts_task


async def run_response_task(
    session: ConnectionSession, recognized_text: str, response_token: int
) -> None:
    try:
        await process_llm_and_tts(session, recognized_text, response_token)
    except asyncio.CancelledError:
        logger.info("⏹ 当前回复已取消")
        raise
    except (RuntimeError, WebSocketDisconnect):
        raise
    except Exception as e:
        logger.exception("❌ 回复任务异常: %s", e)
        if response_token == session.response_token:
            with suppress(RuntimeError, WebSocketDisconnect):
                await session.send_json({"type": "error", "message": "识别或生成失败，请重试。"})
                await session.send_json({"type": "state", "phase": "idle"})
                await session.send_json({"type": "done"})
    finally:
        if session.response_task is asyncio.current_task():
            session.response_task = None
            session.interrupt_requested = False


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    session = ConnectionSession(websocket=websocket)
    logger.info("🔌 客户端已连接")

    await session.send_json(
        {
            "type": "init",
            "outputSampleRate": output_sample_rate,
            "inputSampleRate": INPUT_SAMPLE_RATE,
        }
    )

    vad_iterator = VADIterator(vad_model, sampling_rate=INPUT_SAMPLE_RATE)
    rolling_buffer = deque(maxlen=20)  # Keep a small lead-in to avoid clipping word starts.
    audio_buffer = bytearray()
    vad_pending = bytearray()
    last_partial_len = 0

    try:
        while True:
            message = await websocket.receive()

            if message.get("type") == "websocket.disconnect":
                logger.info("👋 客户端正常挂断了连接")
                break

            if "bytes" not in message or not message["bytes"]:
                continue

            vad_pending.extend(message["bytes"])

            while len(vad_pending) >= VAD_FRAME_BYTES:
                vad_chunk = bytes(vad_pending[:VAD_FRAME_BYTES])
                del vad_pending[:VAD_FRAME_BYTES]

                audio_np = np.frombuffer(vad_chunk, dtype=np.int16).astype(np.float32) / 32768.0
                tensor_chunk = torch.from_numpy(audio_np)

                try:
                    speech_dict = vad_iterator(tensor_chunk)
                    rolling_buffer.append(vad_chunk)

                    if not vad_iterator.triggered:
                        if speech_dict and "end" in speech_dict:
                            if len(audio_buffer) == 0:
                                continue

                            logger.info("🎙 语音输入结束，正在执行最终转写...")
                            full_audio_np = (
                                np.frombuffer(audio_buffer, dtype=np.int16).astype(np.float32) / 32768.0
                            )

                            try:
                                asr_result = asr_model.generate(full_audio_np)
                                recognized_text = getattr(asr_result, "text", str(asr_result)).strip()
                                if recognized_text:
                                    if session.response_task is not None and not session.response_task.done():
                                        await interrupt_active_response(session)
                                    launch_response_task(session, recognized_text)
                            except Exception as e:
                                logger.exception("❌ 最终转写异常: %s", e)
                                with suppress(RuntimeError, WebSocketDisconnect):
                                    await session.send_json({"type": "error", "message": "识别失败，请重试。"})
                                    await session.send_json({"type": "state", "phase": "idle"})
                                    await session.send_json({"type": "done"})

                            audio_buffer.clear()
                            rolling_buffer.clear()
                            last_partial_len = 0
                    else:
                        if speech_dict and "start" in speech_dict:
                            if session.response_task is not None and not session.response_task.done():
                                await interrupt_active_response(session)
                            logger.info("🎙 监听到语音，开始收音并流式听写...")
                            audio_buffer.clear()
                            for buffered_chunk in rolling_buffer:
                                audio_buffer.extend(buffered_chunk)
                            last_partial_len = len(audio_buffer)
                        else:
                            audio_buffer.extend(vad_chunk)

                            if len(audio_buffer) - last_partial_len > PARTIAL_TRANSCRIBE_INTERVAL_BYTES:
                                last_partial_len = len(audio_buffer)
                                partial_audio_np = (
                                    np.frombuffer(audio_buffer, dtype=np.int16).astype(np.float32) / 32768.0
                                )

                                try:
                                    asr_result = asr_model.generate(partial_audio_np)
                                    partial_text = getattr(asr_result, "text", str(asr_result)).strip()
                                    if partial_text:
                                        await session.send_json({"type": "user_partial", "text": partial_text})
                                except Exception:
                                    pass

                except Exception as e:
                    logger.exception("VAD 运行时遇到数据块异常: %s", e)

    except WebSocketDisconnect:
        logger.info("👋 客户端连接断开")
    finally:
        if session.response_task is not None and not session.response_task.done():
            session.response_task.cancel()
            with suppress(asyncio.CancelledError):
                await session.response_task


if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)