端到端实时语音模型

import json
import os
import queue
import sys
import threading
import time

import sounddevice as sd
import websocket


def parse_device(name):
    value = os.environ.get(name)
    if not value:
        return None
    try:
        return int(value)
    except ValueError:
        return value


if "--list-devices" in sys.argv:
    print(sd.query_devices())
    raise SystemExit(0)


API_KEY = os.environ.get("SENSEAUDIO_API_KEY")
WS_URL = os.environ.get(
    "SENSEAUDIO_WS_URL",
    "wss://preview.api.senseaudio.cn/ws/v1/realtime/voice-dialog",
)
GREETING = os.environ.get("SENSEAUDIO_GREETING", "")
VOICE = os.environ.get("SENSEAUDIO_VOICE", "")
INPUT_DEVICE = parse_device("SENSEAUDIO_INPUT_DEVICE")
OUTPUT_DEVICE = parse_device("SENSEAUDIO_OUTPUT_DEVICE")

INPUT_SAMPLE_RATE = 16000
INPUT_BLOCK_FRAMES = 640       # 40ms
OUTPUT_SAMPLE_RATE = 24000
OUTPUT_BLOCK_FRAMES = 480      # 20ms，降低打断后的播放尾音
OUTPUT_BLOCK_BYTES = OUTPUT_BLOCK_FRAMES * 2
MIC_QUEUE_BLOCKS = 25          # 最多缓存约 1 秒上行音频
PLAY_QUEUE_BLOCKS = 500        # 最多缓存约 10 秒下行音频

if not API_KEY:
    raise RuntimeError("请先设置 SENSEAUDIO_API_KEY 环境变量")


class RealtimeVoiceClient:
    def __init__(self):
        self.stop_event = threading.Event()
        self.closed_event = threading.Event()
        self.mic_enabled = threading.Event()
        self.start_sent = threading.Event()
        self.fatal_event = threading.Event()
        self.mic_dropped = threading.Event()
        self.input_warning = threading.Event()
        self.output_warning = threading.Event()
        self.output_busy = threading.Event()

        self.mic_queue = queue.Queue(maxsize=MIC_QUEUE_BLOCKS)
        self.play_queue = queue.Queue(maxsize=PLAY_QUEUE_BLOCKS)
        self.state_lock = threading.Lock()
        self.send_lock = threading.Lock()
        self.output_stream_lock = threading.Lock()

        self.waiting_greeting = bool(GREETING)
        self.response_active = False
        self.downlink_audio_active = False
        self.cancel_sent = False
        self.drop_current_response = False
        self.assistant_text_started = False
        self.playback_epoch = 0
        self.downlink_buffer = bytearray()
        self.ending = False
        self.fatal_reason = None

        self.input_stream = None
        self.output_stream = None
        self.uplink_thread = None
        self.output_thread = None
        self.websocket_thread = None

        self.ws = websocket.WebSocketApp(
            WS_URL,
            header=[f"Authorization: Bearer {API_KEY}"],
            on_open=self.on_open,
            on_message=self.on_message,
            on_close=self.on_close,
            on_error=self.on_error,
        )

    def set_fatal(self, reason):
        with self.state_lock:
            if self.fatal_reason is None:
                self.fatal_reason = reason
        self.fatal_event.set()

    def send_json(self, payload):
        message = json.dumps(payload, ensure_ascii=False)
        with self.send_lock:
            self.ws.send(message)

    def send_binary(self, payload):
        with self.send_lock:
            if self.stop_event.is_set():
                return False
            self.ws.send(payload, opcode=websocket.ABNF.OPCODE_BINARY)
            return True

    # PortAudio 实时回调中不能执行网络请求、打印或其他阻塞操作。
    def input_callback(self, indata, frames, time_info, status):
        if status:
            self.input_warning.set()
        if self.stop_event.is_set() or not self.mic_enabled.is_set():
            return

        # indata 的内存在回调返回后会被复用，必须复制。
        payload = bytes(indata)
        try:
            self.mic_queue.put_nowait(payload)
        except queue.Full:
            # 实时对话优先保留最新音频，避免积压带来数秒延迟。
            self.mic_dropped.set()
            try:
                self.mic_queue.get_nowait()
                self.mic_queue.task_done()
            except queue.Empty:
                pass
            try:
                self.mic_queue.put_nowait(payload)
            except queue.Full:
                pass

    def uplink_worker(self):
        try:
            next_send_at = time.monotonic()
            while not self.stop_event.is_set():
                try:
                    payload = self.mic_queue.get(timeout=0.1)
                except queue.Empty:
                    continue
                try:
                    if self.mic_enabled.is_set():
                        # 即使虚拟设备突发产帧，也按 PCM 实际时长发送。
                        now = time.monotonic()
                        if next_send_at < now - 1:
                            next_send_at = now
                        if self.stop_event.wait(max(0, next_send_at - now)):
                            return
                        if not self.send_binary(payload):
                            return
                        next_send_at = max(next_send_at, now) + (
                            len(payload) / (INPUT_SAMPLE_RATE * 2)
                        )
                finally:
                    self.mic_queue.task_done()
        except websocket.WebSocketConnectionClosedException:
            pass
        except Exception as error:
            self.set_fatal(f"发送麦克风音频失败：{error}")

    def output_worker(self):
        try:
            while not self.stop_event.is_set():
                try:
                    epoch, payload = self.play_queue.get(timeout=0.1)
                except queue.Empty:
                    continue
                try:
                    with self.output_stream_lock:
                        with self.state_lock:
                            valid = (
                                epoch == self.playback_epoch
                                and not self.drop_current_response
                            )
                        if valid and not self.stop_event.is_set():
                            self.output_busy.set()
                            try:
                                if self.output_stream.write(payload):
                                    self.output_warning.set()
                            finally:
                                self.output_busy.clear()
                finally:
                    self.play_queue.task_done()
        except Exception as error:
            if not self.stop_event.is_set():
                self.set_fatal(f"播放模型音频失败：{error}")

    def clear_playback_queue(self):
        while True:
            try:
                self.play_queue.get_nowait()
            except queue.Empty:
                return
            else:
                self.play_queue.task_done()

    def enqueue_playback(self, epoch, payload):
        try:
            self.play_queue.put_nowait((epoch, payload))
        except queue.Full:
            self.set_fatal("扬声器播放队列已满，请检查输出设备")

    def handle_binary_audio(self, message):
        payload = bytes(message)
        with self.state_lock:
            if (
                not self.downlink_audio_active
                or self.drop_current_response
                or self.ending
            ):
                return
            epoch = self.playback_epoch

        self.downlink_buffer.extend(payload)
        while len(self.downlink_buffer) >= OUTPUT_BLOCK_BYTES:
            block = bytes(self.downlink_buffer[:OUTPUT_BLOCK_BYTES])
            del self.downlink_buffer[:OUTPUT_BLOCK_BYTES]
            self.enqueue_playback(epoch, block)

    def finish_downlink_audio(self):
        with self.state_lock:
            valid = (
                self.downlink_audio_active
                and not self.drop_current_response
                and not self.ending
            )
            epoch = self.playback_epoch
            self.downlink_audio_active = False

        if valid and self.downlink_buffer:
            if len(self.downlink_buffer) % 2:
                self.set_fatal("服务端返回了长度不是 16-bit 对齐的 PCM 音频")
            else:
                self.enqueue_playback(epoch, bytes(self.downlink_buffer))
        self.downlink_buffer.clear()

    def interrupt_assistant(self):
        with self.state_lock:
            had_audio = (
                self.downlink_audio_active
                or not self.play_queue.empty()
                or self.output_busy.is_set()
            )
            self.playback_epoch += 1
            should_cancel = (
                self.response_active
                and not self.cancel_sent
                and not self.ending
            )
            self.response_active = False
            self.downlink_audio_active = False
            if should_cancel:
                self.cancel_sent = True
                self.drop_current_response = True

        self.downlink_buffer.clear()
        self.clear_playback_queue()

        if had_audio and self.output_stream is not None:
            try:
                with self.output_stream_lock:
                    self.output_stream.abort(ignore_errors=False)
                    if not self.stop_event.is_set():
                        self.output_stream.start()
            except Exception as error:
                self.set_fatal(f"重置扬声器失败：{error}")

        if should_cancel:
            print("\n检测到用户打断，已取消当前模型回复")
            try:
                self.send_json({"type": "cancel"})
            except websocket.WebSocketConnectionClosedException:
                pass

    def enable_mic_after_greeting(self):
        # turn.done 到达时开场白二进制帧已经全部入队，排空后再开麦避免回声。
        self.play_queue.join()
        if self.stop_event.is_set():
            return
        try:
            with self.output_stream_lock:
                # stop() 会等待 PortAudio 和硬件中的待播缓冲全部完成。
                self.output_stream.stop(ignore_errors=False)
                if self.stop_event.is_set():
                    return
                self.output_stream.start()
        except Exception as error:
            self.set_fatal(f"等待开场白播放完成失败：{error}")
            return
        self.mic_enabled.set()
        print("开场白播放完成，可以开始说话；按 Ctrl+C 结束会话")

    def request_end(self, reason):
        with self.state_lock:
            if self.ending:
                return
            self.ending = True
            self.playback_epoch += 1

        print(f"\n结束会话：{reason}")
        self.mic_enabled.clear()
        self.stop_event.set()
        self.downlink_buffer.clear()
        self.clear_playback_queue()

        if self.start_sent.is_set():
            threading.Thread(
                target=self.send_end,
                name="senseaudio-end",
                daemon=True,
            ).start()
        else:
            self.abort_websocket()

    def send_end(self):
        try:
            self.send_json({"type": "end"})
        except Exception:
            self.abort_websocket()

    def abort_websocket(self):
        # 低层 abort 可唤醒卡在 send()/recv() 中的线程。
        sock = self.ws.sock
        if sock is not None:
            try:
                sock.abort()
            except Exception:
                pass

    def on_open(self, ws):
        if self.stop_event.is_set():
            self.abort_websocket()
            return

        start = {
            "type": "start",
            "model": "senseaudio-realtime-1.0",
            "instructions": "你是SenseAudio助手",
            "greeting": GREETING,
            "audio_setting": {
                "sample_rate": INPUT_SAMPLE_RATE,
                "format": "pcm_s16le",
                "channel": 1,
            },
        }
        if VOICE:
            start["voice"] = VOICE

        try:
            self.send_json(start)
            self.start_sent.set()
        except Exception as error:
            self.set_fatal(f"发送 start 失败：{error}")

    def on_message(self, ws, message):
        if self.stop_event.is_set():
            return

        if isinstance(message, (bytes, bytearray, memoryview)):
            self.handle_binary_audio(message)
            return

        try:
            event = json.loads(message)
        except json.JSONDecodeError as error:
            self.set_fatal(f"服务端返回了无效 JSON：{error}")
            return

        event_type = event.get("type")

        if event_type == "ready":
            if self.waiting_greeting:
                print("连接已就绪，正在播放开场白")
            else:
                self.mic_enabled.set()
                print("可以开始说话；按 Ctrl+C 结束会话")
            return

        if event_type == "speech.started":
            self.interrupt_assistant()
            return

        if event_type == "user.transcript.done":
            with self.state_lock:
                # 新用户轮次开始生成回复，解除上一轮 cancel 的丢弃状态。
                self.cancel_sent = False
                self.drop_current_response = False
                self.response_active = True
                self.assistant_text_started = False
            print(f"\n你：{event.get('text', '')}")
            return

        if event_type == "assistant.text.delta":
            with self.state_lock:
                if self.drop_current_response:
                    return
                self.response_active = True
                first_delta = not self.assistant_text_started
                self.assistant_text_started = True
            if first_delta:
                print("\n助手：", end="", flush=True)
            print(event.get("text", ""), end="", flush=True)
            return

        if event_type == "assistant.text.done":
            with self.state_lock:
                suppressed = self.drop_current_response
                text_started = self.assistant_text_started
                if not suppressed:
                    self.response_active = True
                    self.assistant_text_started = False
            if suppressed:
                return
            if text_started:
                print()
            elif event.get("text"):
                print(f"\n助手：{event['text']}")
            return

        if event_type == "assistant.audio.start":
            if (
                event.get("sample_rate") != OUTPUT_SAMPLE_RATE
                or event.get("format") != "pcm_s16le"
            ):
                self.set_fatal("服务端输出音频不是预期的 pcm_s16le / 24000Hz")
                return
            self.downlink_buffer.clear()
            with self.state_lock:
                self.downlink_audio_active = True
                if not self.drop_current_response:
                    self.response_active = True
            return

        if event_type == "assistant.audio.done":
            self.finish_downlink_audio()
            return

        if event_type == "turn.done":
            # 正常情况下 audio.done 已先到达；这里也做一次幂等收尾。
            self.finish_downlink_audio()
            with self.state_lock:
                greeting_done = self.waiting_greeting
                text_incomplete = self.assistant_text_started
                self.waiting_greeting = False
                self.response_active = False
                self.downlink_audio_active = False
                self.cancel_sent = False
                self.drop_current_response = False
                self.assistant_text_started = False

            if text_incomplete:
                print()
            if greeting_done:
                threading.Thread(
                    target=self.enable_mic_after_greeting,
                    name="senseaudio-greeting-drain",
                    daemon=True,
                ).start()
            else:
                print("本轮结束，继续监听麦克风")
            return

        if event_type == "error":
            print(
                "\n服务端错误："
                f"{event.get('code', 'unknown')} - {event.get('message', '')}"
            )
            if event.get("code") == "no_active_response":
                with self.state_lock:
                    self.cancel_sent = False
                    self.drop_current_response = False
                    self.response_active = False
                    self.assistant_text_started = False
            return

        print("收到未知服务端事件：", event)

    def on_close(self, ws, code, reason):
        self.mic_enabled.clear()
        self.stop_event.set()
        self.closed_event.set()
        print(f"\n连接关闭：{code} {reason or ''}")

    def on_error(self, ws, error):
        if not self.ending:
            print(f"\n连接错误：{error}")

    def open_audio_devices(self):
        try:
            sd.check_input_settings(
                device=INPUT_DEVICE,
                channels=1,
                dtype="int16",
                samplerate=INPUT_SAMPLE_RATE,
            )
            sd.check_output_settings(
                device=OUTPUT_DEVICE,
                channels=1,
                dtype="int16",
                samplerate=OUTPUT_SAMPLE_RATE,
            )
            self.input_stream = sd.RawInputStream(
                device=INPUT_DEVICE,
                samplerate=INPUT_SAMPLE_RATE,
                channels=1,
                dtype="int16",
                blocksize=INPUT_BLOCK_FRAMES,
                latency="low",
                callback=self.input_callback,
            )
            self.output_stream = sd.RawOutputStream(
                device=OUTPUT_DEVICE,
                samplerate=OUTPUT_SAMPLE_RATE,
                channels=1,
                dtype="int16",
                blocksize=OUTPUT_BLOCK_FRAMES,
                latency="low",
            )
        except Exception as error:
            if self.input_stream is not None:
                self.input_stream.close()
                self.input_stream = None
            if self.output_stream is not None:
                self.output_stream.close()
                self.output_stream = None
            raise RuntimeError(
                "无法打开 16kHz 麦克风或 24kHz 扬声器；"
                "请运行 python3 -m sounddevice 检查设备"
            ) from error

    def run(self):
        self.open_audio_devices()
        self.uplink_thread = threading.Thread(
            target=self.uplink_worker,
            name="senseaudio-uplink",
            daemon=True,
        )
        self.output_thread = threading.Thread(
            target=self.output_worker,
            name="senseaudio-playback",
            daemon=True,
        )
        self.websocket_thread = threading.Thread(
            target=lambda: self.ws.run_forever(
                ping_interval=20,
                ping_timeout=10,
            ),
            name="senseaudio-websocket",
            daemon=True,
        )

        try:
            self.output_stream.start()
            self.input_stream.start()
            self.uplink_thread.start()
            self.output_thread.start()
            self.websocket_thread.start()

            while self.websocket_thread.is_alive():
                if self.mic_dropped.is_set():
                    self.mic_dropped.clear()
                    print("\n警告：网络发送较慢，已丢弃过期麦克风帧")
                if self.input_warning.is_set():
                    self.input_warning.clear()
                    print("\n警告：麦克风发生丢帧或输入溢出")
                if self.output_warning.is_set():
                    self.output_warning.clear()
                    print("\n警告：扬声器发生输出欠载")
                if self.fatal_event.is_set():
                    self.request_end(self.fatal_reason or "本地音频错误")
                    break
                self.websocket_thread.join(0.1)

        except KeyboardInterrupt:
            self.request_end("用户按下 Ctrl+C")
        finally:
            if self.websocket_thread.is_alive() and not self.closed_event.wait(2):
                self.abort_websocket()
                self.ws.close()

            self.stop_event.set()
            self.mic_enabled.clear()
            self.clear_playback_queue()

            if self.input_stream is not None:
                self.input_stream.abort()
                self.input_stream.close()

            if self.uplink_thread is not None:
                self.uplink_thread.join(timeout=1)
            if self.output_thread is not None:
                self.output_thread.join(timeout=1)

            if self.output_stream is not None:
                self.output_stream.abort()
                self.output_stream.close()
            if self.websocket_thread is not None:
                self.websocket_thread.join(timeout=1)


if __name__ == "__main__":
    RealtimeVoiceClient().run()

python3 -m sounddevice
export SENSEAUDIO_INPUT_DEVICE='输入设备名称或编号'
export SENSEAUDIO_OUTPUT_DEVICE='输出设备名称或编号'
python3 realtime_voice.py