xiaozhi-esp32-server/main/xiaozhi-server/core/providers/tts/base.py

import os
import re
import queue
import uuid
import asyncio
import threading
from typing import Callable, Any

from core.utils import p3
import time
from core.utils import textUtils
from abc import ABC, abstractmethod
from config.logger import setup_logging
from core.utils.audio_flow_control import FlowControlConfig, simulate_device_consumption
from core.utils.util import audio_to_data, audio_bytes_to_data, audio_bytes_to_data_stream, audio_to_data_stream
from core.utils.tts import MarkdownCleaner
from core.utils.output_counter import add_device_output
from core.handle.reportHandle import enqueue_tts_report
from core.handle.sendAudioHandle import sendAudioMessage
from core.providers.tts.dto.dto import (
    TTSMessageDTO,
    SentenceType,
    ContentType,
    InterfaceType,
)

import traceback

TAG = __name__
logger = setup_logging()


class TTSProviderBase(ABC):
    def __init__(self, config, delete_audio_file):
        self.interface_type = InterfaceType.NON_STREAM
        self.conn = None
        self.tts_timeout = 10
        self.delete_audio_file = delete_audio_file
        self.audio_file_type = "wav"
        self.output_file = config.get("output_dir", "tmp/")
        self.tts_text_queue = queue.Queue()
        self.tts_audio_queue = queue.Queue()
        self.tts_audio_first_sentence = True
        self.before_stop_play_files = []

        self.tts_text_buff = []
        self.punctuations = (
            "。",
            "？",
            "?",
            "！",
            "!",
            "；",
            ";",
            "：",
        )
        self.first_sentence_punctuations = (
            "，",
            "~",
            "、",
            ",",
            "。",
            "？",
            "?",
            "！",
            "!",
            "；",
            ";",
            "：",
        )
        self.tts_stop_request = False
        self.processed_chars = 0
        self.is_first_sentence = True
        self.flow_controller = FlowControlConfig.create_flow_controller()
        self.flow_control_enabled = config.get("enable_flow_control", True)

    def generate_filename(self, extension=".wav"):
        return os.path.join(
            self.output_file,
            f"tts-{datetime.now().date()}@{uuid.uuid4().hex}{extension}",
        )

    def handle_opus(self, opus_data: bytes):
        logger.bind(tag=TAG).debug(
            f"推送数据到队列里面帧数～～ {len(opus_data)}"
        )
        self.tts_audio_queue.put(
            (SentenceType.MIDDLE, opus_data, None)
        )

    def to_tts(self, text, opus_handler=handle_opus) -> None:
        text = MarkdownCleaner.clean_markdown(text)
        max_repeat_time = 5
        if self.delete_audio_file:
            # 需要删除文件的直接转为音频数据
            while max_repeat_time > 0:
                try:
                    audio_bytes = asyncio.run(self.text_to_speak(text, None))
                    if audio_bytes:
                        self.tts_audio_queue.put(
                            (SentenceType.FIRST, None, text)
                        )
                        audio_bytes_to_data_stream(
                            audio_bytes, file_type=self.audio_file_type, is_opus=True, callback=opus_handler
                        )
                        max_repeat_time = 0
                    else:
                        max_repeat_time -= 1
                except Exception as e:
                    logger.bind(tag=TAG).warning(
                        f"语音生成失败{5 - max_repeat_time + 1}次: {text}，错误: {e}"
                    )
                    max_repeat_time -= 1
            if max_repeat_time > 0:
                logger.bind(tag=TAG).info(
                    f"语音生成成功: {text}，重试{5 - max_repeat_time}次"
                )
            else:
                logger.bind(tag=TAG).error(
                    f"语音生成失败: {text}，请检查网络或服务是否正常"
                )
            return None
        else:
            tmp_file = self.generate_filename()
            try:
                while not os.path.exists(tmp_file) and max_repeat_time > 0:
                    try:
                        asyncio.run(self.text_to_speak(text, tmp_file))
                    except Exception as e:
                        logger.bind(tag=TAG).warning(
                            f"语音生成失败{5 - max_repeat_time + 1}次: {text}，错误: {e}"
                        )
                        # 未执行成功，删除文件
                        if os.path.exists(tmp_file):
                            os.remove(tmp_file)
                        max_repeat_time -= 1

                if max_repeat_time > 0:
                    logger.bind(tag=TAG).info(
                        f"语音生成成功: {text}:{tmp_file}，重试{5 - max_repeat_time}次"
                    )
                else:
                    logger.bind(tag=TAG).error(
                        f"语音生成失败: {text}，请检查网络或服务是否正常"
                    )
                    self.tts_audio_queue.put(
                        (SentenceType.FIRST, None, text)
                    )
                self._process_audio_file(tmp_file, callback=opus_handler)
            except Exception as e:
                logger.bind(tag=TAG).error(f"Failed to generate TTS file: {e}")
                return None

    @abstractmethod
    async def text_to_speak(self, text, output_file):
        pass

    def audio_to_pcm_data_stream(self, audio_file_path, callback: Callable[[Any], Any]=None):
        """音频文件转换为PCM编码"""
        return audio_to_data_stream(audio_file_path, is_opus=False, callback=callback)

    def audio_to_opus_data_stream(self, audio_file_path, callback: Callable[[Any], Any]=None):
        """音频文件转换为Opus编码"""
        return audio_to_data_stream(audio_file_path, is_opus=True, callback=callback)

    def tts_one_sentence(
        self,
        conn,
        content_type,
        content_detail=None,
        content_file=None,
        sentence_id=None,
    ):
        """发送一句话"""
        if not sentence_id:
            if conn.sentence_id:
                sentence_id = conn.sentence_id
            else:
                sentence_id = str(uuid.uuid4().hex)
                conn.sentence_id = sentence_id
        # 对于单句的文本，进行分段处理
        segments = re.split(r"([。！？!?；;\n])", content_detail)
        for seg in segments:
            self.tts_text_queue.put(
                TTSMessageDTO(
                    sentence_id=sentence_id,
                    sentence_type=SentenceType.MIDDLE,
                    content_type=content_type,
                    content_detail=seg,
                    content_file=content_file,
                )
            )

    async def open_audio_channels(self, conn):
        self.conn = conn
        self.tts_timeout = conn.config.get("tts_timeout", 10)
        # tts 消化线程
        self.tts_priority_thread = threading.Thread(
            target=self.tts_text_priority_thread, daemon=True
        )
        self.tts_priority_thread.start()

        # 音频播放 消化线程
        self.audio_play_priority_thread = threading.Thread(
            target=self._audio_play_priority_thread, daemon=True
        )
        self.audio_play_priority_thread.start()

    # 这里默认是非流式的处理方式
    # 流式处理方式请在子类中重写
    def tts_text_priority_thread(self):
        while not self.conn.stop_event.is_set():
            try:
                message = self.tts_text_queue.get(timeout=1)
                if message.sentence_type == SentenceType.FIRST:
                    self.conn.client_abort = False
                if self.conn.client_abort:
                    logger.bind(tag=TAG).info("收到打断信息，终止TTS文本处理线程")
                    continue
                if message.sentence_type == SentenceType.FIRST:
                    # 初始化参数
                    self.tts_stop_request = False
                    self.processed_chars = 0
                    self.tts_text_buff = []
                    self.is_first_sentence = True
                    self.tts_audio_first_sentence = True
                elif ContentType.TEXT == message.content_type:
                    self.tts_text_buff.append(message.content_detail)
                    segment_text = self._get_segment_text()
                    if segment_text:
                        self.to_tts(segment_text, opus_handler=self.handle_opus)
                elif ContentType.FILE == message.content_type:
                    self._process_remaining_text()
                    tts_file = message.content_file
                    if tts_file and os.path.exists(tts_file):
                        self._process_audio_file(tts_file, callback=self.handle_opus)
                if message.sentence_type == SentenceType.LAST:
                    self._process_remaining_text()
                    self.tts_audio_queue.put(
                        (message.sentence_type, [], message.content_detail)
                    )

            except queue.Empty:
                continue
            except Exception as e:
                logger.bind(tag=TAG).error(
                    f"处理TTS文本失败: {str(e)}, 类型: {type(e).__name__}, 堆栈: {traceback.format_exc()}"
                )
                continue

    def _audio_play_priority_thread(self):
        while not self.conn.stop_event.is_set():
            text = None
            try:
                try:
                    sentence_type, audio_datas, text = self.tts_audio_queue.get(timeout=1)
                except queue.Empty:
                    if self.conn.stop_event.is_set():
                        break
                    continue

                # 如果启用了流控
                if self.flow_control_enabled:
                    # 计算音频数据的帧数
                    if isinstance(audio_datas, bytes):
                        frame_count = 1  # 单个字节流作为一帧
                    elif isinstance(audio_datas, (list, tuple)):
                        frame_count = len(audio_datas)
                    else:
                        frame_count = 0

                    # 流控检查
                    if frame_count > 0:
                        max_wait_time = FlowControlConfig.DEFAULT_MAX_WAIT_TIME
                        wait_start_time = time.time()
                        retry_interval = FlowControlConfig.DEFAULT_RETRY_INTERVAL

                        while not self.flow_controller.can_send_frames(frame_count):
                            # 检查是否超时或需要停止
                            if (time.time() - wait_start_time > max_wait_time or
                                    self.conn.stop_event.is_set() or
                                    self.conn.client_abort):
                                logger.bind(tag=TAG).warning(
                                    f"流控等待超时或收到停止信号，跳过音频发送: {text}"
                                )
                                break

                            # 短暂等待后重试
                            time.sleep(retry_interval)

                            # 更新设备消费估计（这里假设设备以恒定速率消费）
                            # 实际应用中可能需要从设备端获取真实的消费情况
                            estimated_consumption = int(retry_interval * FlowControlConfig.DEFAULT_REFILL_RATE)
                            self.flow_controller.update_device_consumption(estimated_consumption)
                            status = self.flow_controller.get_status()
                            logger.bind(tag=TAG).debug(
                                f"流控状态: 缓冲区使用率={status['buffer_usage_percent']:.1f}%, "
                                f"可用令牌={status['available_tokens']}..."
                            )
                        else:
                            # 可以发送，记录发送的帧数
                            self.flow_controller.record_sent_frames(frame_count)

                            # 发送音频
                            future = asyncio.run_coroutine_threadsafe(
                                self._send_audio_with_flow_control(sentence_type, audio_datas, text),
                                self.conn.loop,
                            )
                            future.result()

                            # 记录输出和报告
                            if self.conn.max_output_size > 0 and text:
                                add_device_output(self.conn.headers.get("device-id"), len(text))
                            enqueue_tts_report(self.conn, text, audio_datas)

                            # 输出流控状态（调试用）
                            if frame_count > 0:  # 只在较大的音频块时输出状态
                                status = self.flow_controller.get_status()
                                logger.bind(tag=TAG).debug(
                                    f"流控状态: 缓冲区使用率={status['buffer_usage_percent']:.1f}%, "
                                    f"可用令牌={status['available_tokens']}..."
                                )
                    else:
                        # 没有音频数据，直接发送
                        future = asyncio.run_coroutine_threadsafe(
                            sendAudioMessage(self.conn, sentence_type, audio_datas, text),
                            self.conn.loop,
                        )
                        future.result()
                else:
                    # 未启用流控，直接发送
                    future = asyncio.run_coroutine_threadsafe(
                        sendAudioMessage(self.conn, sentence_type, audio_datas, text),
                        self.conn.loop,
                    )
                    future.result()

                    # 记录输出和报告
                    if self.conn.max_output_size > 0 and text:
                        add_device_output(self.conn.headers.get("device-id"), len(text))
                    enqueue_tts_report(self.conn, text, audio_datas)

            except Exception as e:
                logger.bind(tag=TAG).error(
                    f"audio_play_priority_thread: {text} {e}"
                )

    async def _send_audio_with_flow_control(self, sentence_type, audio_datas, text):
        """带流控的音频发送方法"""
        await sendAudioMessage(self.conn, sentence_type, audio_datas, text)

        # 模拟设备消费（实际应用中应该从设备获取反馈）
        if isinstance(audio_datas, bytes):
            frame_count = 1
        elif isinstance(audio_datas, (list, tuple)):
            frame_count = len(audio_datas)
        else:
            frame_count = 0

        if frame_count > 0:
            asyncio.create_task(simulate_device_consumption(self.flow_controller, frame_count))

    # 在类中添加流控制器重置方法
    def reset_flow_controller(self):
        """重置流控制器状态，通常在新会话开始时调用"""
        if hasattr(self, 'flow_controller'):
            self.flow_controller.reset()
            logger.bind(tag=TAG).info("流控制器状态已重置")

    async def start_session(self, session_id):
        pass

    async def finish_session(self, session_id):
        pass

    async def close(self):
        """资源清理方法"""
        if hasattr(self, "ws") and self.ws:
            await self.ws.close()

    def _get_segment_text(self):
        # 合并当前全部文本并处理未分割部分
        full_text = "".join(self.tts_text_buff)
        current_text = full_text[self.processed_chars :]  # 从未处理的位置开始
        last_punct_pos = -1

        # 根据是否是第一句话选择不同的标点符号集合
        punctuations_to_use = (
            self.first_sentence_punctuations
            if self.is_first_sentence
            else self.punctuations
        )

        for punct in punctuations_to_use:
            pos = current_text.rfind(punct)
            if (pos != -1 and last_punct_pos == -1) or (
                pos != -1 and pos < last_punct_pos
            ):
                last_punct_pos = pos

        if last_punct_pos != -1:
            segment_text_raw = current_text[: last_punct_pos + 1]
            segment_text = textUtils.get_string_no_punctuation_or_emoji(
                segment_text_raw
            )
            self.processed_chars += len(segment_text_raw)  # 更新已处理字符位置

            # 如果是第一句话，在找到第一个逗号后，将标志设置为False
            if self.is_first_sentence:
                self.is_first_sentence = False

            return segment_text
        elif self.tts_stop_request and current_text:
            segment_text = current_text
            self.is_first_sentence = True  # 重置标志
            return segment_text
        else:
            return None

    def _process_audio_file(self, tts_file, callback: Callable[[Any], Any]):
        """处理音频文件并转换为指定格式

        Args:
            tts_file: 音频文件路径
            content_detail: 内容详情

        Returns:
            tuple: (sentence_type, audio_datas, content_detail)
        """
        if tts_file.endswith(".p3"):
            p3.decode_opus_from_file_stream(tts_file, callback=callback)
        elif self.conn.audio_format == "pcm":
            self.audio_to_pcm_data_stream(tts_file, callback=callback)
        else:
            self.audio_to_opus_data_stream(tts_file, callback=callback)

        if (
            self.delete_audio_file
            and tts_file is not None
            and os.path.exists(tts_file)
            and tts_file.startswith(self.output_file)
        ):
            os.remove(tts_file)

    def _process_before_stop_play_files(self):
        for audio_datas, text in self.before_stop_play_files:
            self.tts_audio_queue.put((SentenceType.MIDDLE, audio_datas, text))
        self.before_stop_play_files.clear()
        self.tts_audio_queue.put((SentenceType.LAST, [], None))

    def _process_remaining_text(self):
        """处理剩余的文本并生成语音

        Returns:
            bool: 是否成功处理了文本
        """
        full_text = "".join(self.tts_text_buff)
        remaining_text = full_text[self.processed_chars :]
        if remaining_text:
            segment_text = textUtils.get_string_no_punctuation_or_emoji(remaining_text)
            if segment_text:
                if self.delete_audio_file:
                    audio_datas = self.to_tts(segment_text)
                    if audio_datas:
                        self.tts_audio_queue.put(
                            (SentenceType.MIDDLE, audio_datas, segment_text)
                        )
                else:
                    tts_file = self.to_tts(segment_text)
                    audio_datas = self._process_audio_file(tts_file, callback=self.handle_opus)
                    self.tts_audio_queue.put(
                        (SentenceType.MIDDLE, audio_datas, segment_text)
                    )
                self.processed_chars += len(full_text)
                return True
        return False