端到端实时语音模型

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://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 = 3000       # 最多缓存约 60 秒下行音频

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.playback_backlog = 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.pending_tool_calls = {}

        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

    def dispatch_tool_call(self, event):
        call_id = event.get("call_id")
        if not isinstance(call_id, str) or not call_id:
            print("\n忽略缺少 call_id 的工具调用")
            return

        cancel_event = threading.Event()
        with self.state_lock:
            if call_id in self.pending_tool_calls:
                print(f"\n忽略重复的工具调用：{call_id}")
                return
            self.pending_tool_calls[call_id] = cancel_event

        threading.Thread(
            target=self.handle_tool_call,
            args=(event, cancel_event),
            name=f"senseaudio-tool-{call_id}",
            daemon=True,
        ).start()

    def handle_tool_call(self, event, cancel_event):
        call_id = event["call_id"]
        timeout_ms = event.get("timeout_ms", 30000)
        if not isinstance(timeout_ms, int) or timeout_ms <= 0:
            timeout_ms = 30000
        deadline = time.monotonic() + timeout_ms / 1000

        try:
            name = event.get("name")
            arguments = event.get("arguments")

            if name != "query_weather":
                result = {
                    "error": {
                        "code": "unsupported_tool",
                        "message": f"不支持工具 {name}",
                    }
                }
            elif not isinstance(arguments, dict):
                result = {
                    "error": {
                        "code": "invalid_arguments",
                        "message": "arguments 必须是 JSON 对象",
                    }
                }
            else:
                city = arguments.get("city")
                unit = arguments.get("unit", "celsius")
                if not isinstance(city, str) or not city.strip():
                    result = {
                        "error": {
                            "code": "invalid_arguments",
                            "message": "city 必须是非空字符串",
                        }
                    }
                elif unit not in {"celsius", "fahrenheit"}:
                    result = {
                        "error": {
                            "code": "invalid_arguments",
                            "message": "unit 必须是 celsius 或 fahrenheit",
                        }
                    }
                else:
                    # 模拟结果。真实工具 I/O 也应在此工作线程中执行。
                    result = {
                        "city": city,
                        "temperature": 31,
                        "unit": unit,
                    }

            if (
                cancel_event.is_set()
                or self.stop_event.is_set()
                or time.monotonic() >= deadline
            ):
                return

            self.send_json(
                {
                    "type": "tool.result",
                    "call_id": call_id,
                    "result": result,
                }
            )
        except websocket.WebSocketConnectionClosedException:
            pass
        except Exception as error:
            print(f"\n工具调用失败：{error}")
        finally:
            with self.state_lock:
                if self.pending_tool_calls.get(call_id) is cancel_event:
                    del self.pending_tool_calls[call_id]

    def handle_tool_cancelled(self, event):
        call_id = event.get("call_id")
        with self.state_lock:
            cancel_event = self.pending_tool_calls.pop(call_id, None)
        if cancel_event is not None:
            cancel_event.set()
        print(
            f"\n工具调用已取消：{call_id or 'unknown'} "
            f"({event.get('reason', '')})"
        )

    # 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:
                                self.output_stream.write(payload)
                            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.playback_backlog.set()
            try:
                self.play_queue.get_nowait()
                self.play_queue.task_done()
            except queue.Empty:
                pass
            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
            tool_cancellations = list(self.pending_tool_calls.values())
            self.pending_tool_calls.clear()

        print(f"\n结束会话：{reason}")
        for cancel_event in tool_cancellations:
            cancel_event.set()
        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,
            },
            "tools": [
                {
                    "type": "function",
                    "name": "query_weather",
                    "description": "查询指定城市的当前天气",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "city": {
                                "type": "string",
                                "description": "城市名称，例如上海",
                            },
                            "unit": {
                                "type": "string",
                                "enum": ["celsius", "fahrenheit"],
                                "description": "温度单位",
                            },
                        },
                        "required": ["city"],
                        "additionalProperties": False,
                    },
                },
            ],
        }
        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.delta":
            print(f"\r用户转写中：{event.get('text', '')}", end="", flush=True)
            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 == "tool.call":
            self.dispatch_tool_call(event)
            return

        if event_type == "tool.cancelled":
            self.handle_tool_cancelled(event)
            return

        if event_type == "session.action.completed":
            print(f"\n会话动作完成：{event.get('action', '')}")
            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 扬声器；"
                "请运行 uv run --with sounddevice python -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.playback_backlog.is_set():
                    self.playback_backlog.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()
                try:
                    self.ws.close()
                except Exception:
                    pass

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

            for stream in (self.input_stream, self.output_stream):
                if stream is not None:
                    try:
                        stream.abort()
                    except Exception:
                        pass

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

            for stream in (self.input_stream, self.output_stream):
                if stream is not None:
                    try:
                        stream.close()
                    except Exception:
                        pass

            if (
                self.websocket_thread is not None
                and self.websocket_thread.ident is not None
            ):
                self.websocket_thread.join(timeout=1)


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