Files
banban/talkingq-url/services/text_generator.py
2026-03-24 15:04:36 +08:00

154 lines
6.1 KiB
Python

import asyncio
import time
from typing import Dict, List
from utils.logger import session_logger
from services.interrupt_handler import interrupt_handler
from utils.text_splitter import (
PUNCTUATION_MARKS,
SENTENCE_ENDINGS,
is_numbered_list_item,
split_into_sentences,
)
from services.tts_config import TTSConfig
from services.interruption_helper import InterruptionHelper
class TextGenerator:
def __init__(self, device_id: str, session_id: str):
self.device_id = device_id
self.session_id = session_id
self.session_key = (device_id, session_id)
def check_interruption(self):
return interrupt_handler.is_interrupted(self.session_key)
async def generate_text(
self,
llm_service,
transcript: str,
history: List[Dict[str, str]],
selected_role: dict,
text_queue: asyncio.Queue,
) -> str:
reply = ""
text_buffer = ""
first_text_token_time = None
sent_batches_count = 0
first_punct_sent = False
initial_batches_count = 0
llm_first_token_timeout = TTSConfig.get_llm_first_token_timeout()
first_token_timeout = False
start_time = time.perf_counter()
async def check_timeout():
nonlocal first_token_timeout
while True:
await asyncio.sleep(0.5) # 每0.5秒检查一次
if first_text_token_time is not None or first_token_timeout:
break # 已收到首个令牌或已超时,停止检查
if time.perf_counter() - start_time > llm_first_token_timeout:
session_logger.error(
self.device_id,
self.session_id,
f"LLM首个token超时({llm_first_token_timeout}秒),主动中断",
)
await InterruptionHelper.handle_llm_timeout(
self.device_id, self.session_id, llm_service
)
break
try:
token_stream = llm_service.generate_response_stream(
transcript, history, selected_role
)
try:
first_token = None
async def get_first_token():
nonlocal first_token
async for token in token_stream:
first_token = token
return token
await asyncio.wait_for(get_first_token(), llm_first_token_timeout)
if first_token:
first_text_token_time = time.perf_counter()
token_latency = first_text_token_time - start_time
session_logger.info(
self.device_id,
self.session_id,
f"收到第一个文本 Token 耗时: {token_latency:.3f}"
)
if first_token is None:
raise Exception("LLM返回的首个token为None")
reply += first_token
text_buffer += first_token
async for token in token_stream:
if self.check_interruption():
session_logger.info(
self.device_id, self.session_id, "LLM生成被中断"
)
break
if token is None:
continue
reply += token
text_buffer += token
should_submit = False
if initial_batches_count < 3 and any(
p in token for p in PUNCTUATION_MARKS
):
if not is_numbered_list_item(text_buffer):
should_submit = True
elif initial_batches_count >= 3 and any(
p in token for p in SENTENCE_ENDINGS
):
if not is_numbered_list_item(text_buffer):
sentences = split_into_sentences(text_buffer)
batch_size = 2
if len(sentences) >= batch_size:
to_synthesize = "".join(sentences[:batch_size])
text_buffer = "".join(sentences[batch_size:])
await text_queue.put(to_synthesize)
initial_batches_count += 1
sent_batches_count += 1
continue
if should_submit:
await text_queue.put(text_buffer)
text_buffer = ""
if not first_punct_sent:
first_punct_sent = True
initial_batches_count += 1
sent_batches_count += 1
except asyncio.TimeoutError:
first_token_timeout = True
session_logger.error(
self.device_id,
self.session_id,
f"LLM首个token超时({llm_first_token_timeout}秒),中断生成",
)
await InterruptionHelper.handle_llm_timeout(
self.device_id, self.session_id, llm_service
)
await text_queue.put(None)
return reply
if (
not self.check_interruption()
and not first_token_timeout
and text_buffer.strip()
):
await text_queue.put(text_buffer.strip())
await text_queue.put(None)
except Exception as e:
session_logger.error(
self.device_id, self.session_id, f"文本生成时出错: {str(e)}"
)
await InterruptionHelper.send_error_prompt_sound(
self.device_id, self.session_id
)
await text_queue.put(None)
return reply