Files
banban/talkingq-url/handlers/service_connection_handler.py
T

237 lines
9.3 KiB
Python

import asyncio
from services.audio_session import AudioSession
from services.task_manager import task_manager
from utils.logger import session_logger
from services.interrupt_handler import interrupt_handler
import time
from config import settings
async def _connect_with_interrupt_cleanup(service, session_key):
"""Own warmup connections even when interruption races with connect()."""
if interrupt_handler.is_interrupted(session_key):
return False
await interrupt_handler.register_cleanup_handler(
session_key, service.close,
priority=interrupt_handler.PRIORITY["NETWORK_CONNECTION"],
)
try:
await service.connect()
if interrupt_handler.is_interrupted(session_key):
await service.close()
return False
return True
except asyncio.CancelledError:
await service.close()
raise
except Exception:
await service.close()
raise
async def prepare_asr_service(session, device_id, session_id):
session.asr_error = None
session.asr_service.sample_rate = session.sample_rate
if hasattr(session.asr_service, "device_id"):
session.asr_service.device_id = device_id
if hasattr(session.asr_service, "session_id"):
session.asr_service.session_id = session_id
await establish_asr_connection(session)
async def prepare_llm_service(session, device_id, session_id, selected_role):
if hasattr(session.llm_service, "device_id"):
session.llm_service.device_id = device_id
if hasattr(session.llm_service, "session_id"):
session.llm_service.session_id = session_id
if hasattr(session.llm_service, "connect"):
if not await _connect_with_interrupt_cleanup(session.llm_service, (device_id, session_id)):
return
from services.conversation_history import conversation_history_manager
from config import settings
device_history = await conversation_history_manager.get_history(device_id)
if device_history:
session.llm_service.history = device_history.history[
-settings.max_conversation_history :
]
async def prepare_tts_service(session, device_id, session_id, selected_role):
session.tts_service.selected_role = selected_role
if hasattr(session.tts_service, "device_id"):
session.tts_service.device_id = device_id
if hasattr(session.tts_service, "session_id"):
session.tts_service.session_id = session_id
await establish_tts_connection(session)
async def establish_asr_connection(session: AudioSession):
try:
if session.device_id and session.session_id:
session_key = (session.device_id, session.session_id)
if interrupt_handler.is_interrupted(session_key):
session_logger.info(
session.device_id, session.session_id, "会话已中断,跳过建立ASR连接"
)
return
session.send_task = await task_manager.create_task(
send_audio_task(session),
device_id=session.device_id,
session_key=(session.device_id, session.session_id) if session.device_id and session.session_id else None,
task_type="audio_send"
)
await session.asr_service.connect()
if not session.asr_service.connected:
raise RuntimeError("语音识别 WebSocket 连接未建立")
if session.device_id and session.session_id:
if interrupt_handler.is_interrupted((session.device_id, session.session_id)):
raise RuntimeError("连接期间会话已中断")
await interrupt_handler.register_cleanup_handler(
(session.device_id, session.session_id), session.asr_service.close,
priority=interrupt_handler.PRIORITY["NETWORK_CONNECTION"],
)
await interrupt_handler.register_cleanup_handler(
(session.device_id, session.session_id), session.cancel_send_task,
priority=interrupt_handler.PRIORITY["TASK_CANCELLATION"],
)
session_logger.info(
"unknown", "unknown", f"语音识别连接已建立,采样率: {session.sample_rate}Hz"
)
except asyncio.CancelledError:
await session.cancel_send_task()
await session.asr_service.close()
raise
except Exception as e:
session.asr_error = f"{type(e).__name__}: {e}"
await session.cancel_send_task()
await session.asr_service.close()
session_logger.error(
"unknown", "unknown", f"建立语音识别 WebSocket 连接失败: {e}", exc_info=True
)
async def establish_tts_connection(session: AudioSession):
try:
if session.device_id and session.session_id:
session_key = (session.device_id, session.session_id)
if interrupt_handler.is_interrupted(session_key):
session_logger.info(
session.device_id, session.session_id, "会话已中断,跳过建立TTS连接"
)
return
if not await _connect_with_interrupt_cleanup(
session.tts_service, (session.device_id, session.session_id)
):
return
if not session.tts_service.connected:
session_logger.error(
session.device_id or "unknown",
session.session_id or "unknown",
"语音合成连接未建立",
)
return
session_logger.info(
session.device_id or "unknown",
session.session_id or "unknown",
"语音合成连接已建立",
)
except Exception as e:
session_logger.error(
session.device_id or "unknown",
session.session_id or "unknown",
f"建立语音合成连接失败: {e}",
)
async def send_audio_task(session: AudioSession):
try:
while True:
if session.device_id and session.session_id:
session_key = (session.device_id, session.session_id)
if interrupt_handler.is_interrupted(session_key):
session_logger.info(
session.device_id, session.session_id, "发送音频任务被中断"
)
break
queue = session.send_queue
audio_data = await queue.get()
try:
if audio_data is None: # Drain all preceding audio before finish-task.
break
await session.asr_service.send_audio(audio_data)
finally:
queue.task_done()
except asyncio.CancelledError:
session.asr_error = session.asr_error or "发送音频任务被取消"
session_logger.info(
session.device_id or "unknown",
session.session_id or "unknown",
"发送音频数据任务被取消",
)
except Exception as e:
session.asr_error = str(e)
session_logger.error(
session.device_id or "unknown",
session.session_id or "unknown",
f"发送音频数据时出错: {str(e)}",
exc_info=True,
)
async def end_asr_session(session: AudioSession):
"""Drain device audio, collect the complete transcript, always release ASR."""
session.transcript = ""
try:
session_key = (session.device_id, session.session_id)
if interrupt_handler.is_interrupted(session_key):
raise RuntimeError("识别会话已中断")
if session.send_task and not session.send_task.done():
await session.send_queue.put(None)
await asyncio.wait_for(
asyncio.shield(session.send_task),
timeout=settings.asr_queue_drain_timeout_seconds,
)
if session.asr_error:
raise RuntimeError(session.asr_error)
if session.send_queue and not session.send_queue.empty():
raise RuntimeError("发送任务提前退出,音频未完整发送")
if interrupt_handler.is_interrupted(session_key):
raise RuntimeError("识别会话已中断")
asr_end_start_time = time.perf_counter()
await session.asr_service.send_end()
if interrupt_handler.is_interrupted(session_key):
raise RuntimeError("识别会话已中断")
session.transcript = await session.asr_service.receive_results()
session_logger.info(
session.device_id or "unknown", session.session_id or "unknown",
f"ASR结束延迟: {time.perf_counter() - asr_end_start_time:.2f} 秒",
)
except Exception as e:
session.asr_error = f"{type(e).__name__}: {e}"
session.transcript = ""
session_logger.error(
session.device_id or "unknown", session.session_id or "unknown",
f"ASR结束处理出错: {session.asr_error}",
)
finally:
if session.send_task:
session.send_task.cancel()
try:
await asyncio.wait_for(session.send_task, timeout=2.0)
except (asyncio.CancelledError, asyncio.TimeoutError):
pass
finally:
session.send_task = None
await session.asr_service.close()
session_logger.info(
session.device_id or "unknown", session.session_id or "unknown",
"语音识别 WebSocket 连接和 ClientSession 已关闭",
)