Merge pull request #2948 from xinnan-tech/py_test_typing

Py test typing

Author: milieyo

main/xiaozhi-server/core/connection.py

@@ -78,7 +78,7 @@ def __init__(
 
         self.read_config_from_api = self.config.get("read_config_from_api", False)
 
-        self.websocket = None
+        self.websocket: websockets.ServerConnection | None = None
         self.headers = None
         self.device_id = None
         self.client_ip = None
@@ -169,7 +169,7 @@ def __init__(
         # 初始化提示词管理器
         self.prompt_manager = PromptManager(self.config, self.logger)
 
-    async def handle_connection(self, ws):
+    async def handle_connection(self, ws: websockets.ServerConnection):
         try:
             # 获取运行中的事件循环（必须在异步上下文中）
             self.loop = asyncio.get_running_loop()

main/xiaozhi-server/core/handle/abortHandle.py

@@ -1,9 +1,12 @@
 import json
+from typing import TYPE_CHECKING
 
+if TYPE_CHECKING:
+    from core.connection import ConnectionHandler
 TAG = __name__
 
 
-async def handleAbortMessage(conn):
+async def handleAbortMessage(conn: "ConnectionHandler"):
     conn.logger.bind(tag=TAG).info("Abort message received")
     # 设置成打断状态，会自动打断llm、tts任务
     conn.client_abort = True

main/xiaozhi-server/core/handle/helloHandle.py

@@ -3,16 +3,17 @@
 import uuid
 import random
 import asyncio
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from core.connection import ConnectionHandler
 from core.utils.dialogue import Message
 from core.utils.util import audio_to_data
 from core.providers.tts.dto.dto import SentenceType
 from core.utils.wakeup_word import WakeupWordsConfig
 from core.handle.sendAudioHandle import sendAudioMessage, send_tts_message
 from core.utils.util import remove_punctuation_and_length, opus_datas_to_wav_bytes
-from core.providers.tools.device_mcp import (
-    MCPClient,
-    send_mcp_initialize_message
-)
+from core.providers.tools.device_mcp import MCPClient, send_mcp_initialize_message
 
 TAG = __name__
 
@@ -38,7 +39,7 @@
 _wakeup_response_lock = asyncio.Lock()
 
 
-async def handleHelloMessage(conn, msg_json):
+async def handleHelloMessage(conn: "ConnectionHandler", msg_json):
     """处理hello消息"""
     audio_params = msg_json.get("audio_params")
     if audio_params:
@@ -59,7 +60,7 @@ async def handleHelloMessage(conn, msg_json):
     await conn.websocket.send(json.dumps(conn.welcome_msg))
 
 
-async def checkWakeupWords(conn, text):
+async def checkWakeupWords(conn: "ConnectionHandler", text):
     enable_wakeup_words_response_cache = conn.config[
         "enable_wakeup_words_response_cache"
     ]
@@ -120,7 +121,7 @@ async def checkWakeupWords(conn, text):
     return True
 
 
-async def wakeupWordsResponse(conn):
+async def wakeupWordsResponse(conn: "ConnectionHandler"):
     if not conn.tts:
         return
 

main/xiaozhi-server/core/handle/intentHandler.py

@@ -1,6 +1,10 @@
 import json
 import uuid
 import asyncio
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from core.connection import ConnectionHandler
 from core.utils.dialogue import Message
 from core.providers.tts.dto.dto import ContentType
 from core.handle.helloHandle import checkWakeupWords
@@ -12,10 +16,10 @@
 TAG = __name__
 
 
-async def handle_user_intent(conn, text):
+async def handle_user_intent(conn: "ConnectionHandler", text):
     # 预处理输入文本，处理可能的JSON格式
     try:
-        if text.strip().startswith('{') and text.strip().endswith('}'):
+        if text.strip().startswith("{") and text.strip().endswith("}"):
             parsed_data = json.loads(text)
             if isinstance(parsed_data, dict) and "content" in parsed_data:
                 text = parsed_data["content"]  # 提取content用于意图分析
@@ -45,7 +49,7 @@ async def handle_user_intent(conn, text):
     return await process_intent_result(conn, intent_result, text)
 
 
-async def check_direct_exit(conn, text):
+async def check_direct_exit(conn: "ConnectionHandler", text):
     """检查是否有明确的退出命令"""
     _, text = remove_punctuation_and_length(text)
     cmd_exit = conn.cmd_exit
@@ -58,7 +62,7 @@ async def check_direct_exit(conn, text):
     return False
 
 
-async def analyze_intent_with_llm(conn, text):
+async def analyze_intent_with_llm(conn: "ConnectionHandler", text):
     """使用LLM分析用户意图"""
     if not hasattr(conn, "intent") or not conn.intent:
         conn.logger.bind(tag=TAG).warning("意图识别服务未初始化")
@@ -75,7 +79,9 @@ async def analyze_intent_with_llm(conn, text):
     return None
 
 
-async def process_intent_result(conn, intent_result, original_text):
+async def process_intent_result(
+    conn: "ConnectionHandler", intent_result, original_text
+):
     """处理意图识别结果"""
     try:
         # 尝试将结果解析为JSON
@@ -94,24 +100,26 @@ async def process_intent_result(conn, intent_result, original_text):
             if function_name == "result_for_context":
                 await send_stt_message(conn, original_text)
                 conn.client_abort = False
-                
+
                 def process_context_result():
                     conn.dialogue.put(Message(role="user", content=original_text))
-                    
+
                     from core.utils.current_time import get_current_time_info
 
-                    current_time, today_date, today_weekday, lunar_date = get_current_time_info()
-                    
+                    current_time, today_date, today_weekday, lunar_date = (
+                        get_current_time_info()
+                    )
+
                     # 构建带上下文的基础提示
                     context_prompt = f"""当前时间：{current_time}
                                         今天日期：{today_date} ({today_weekday})
                                         今天农历：{lunar_date}
 
                                         请根据以上信息回答用户的问题：{original_text}"""
-                    
+
                     response = conn.intent.replyResult(context_prompt, original_text)
                     speak_txt(conn, response)
-                
+
                 conn.executor.submit(process_context_result)
                 return True
 
@@ -188,7 +196,7 @@ def process_function_call():
         return False
 
 
-def speak_txt(conn, text):
+def speak_txt(conn: "ConnectionHandler", text):
     # 记录文本
     conn.tts_MessageText = text
 

main/xiaozhi-server/core/handle/receiveAudioHandle.py

@@ -1,6 +1,10 @@
 import time
 import json
 import asyncio
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from core.connection import ConnectionHandler
 from core.utils.util import audio_to_data
 from core.handle.abortHandle import handleAbortMessage
 from core.handle.intentHandler import handle_user_intent
@@ -10,7 +14,7 @@
 TAG = __name__
 
 
-async def handleAudioMessage(conn, audio):
+async def handleAudioMessage(conn: "ConnectionHandler", audio):
     # 当前片段是否有人说话
     have_voice = conn.vad.is_vad(conn, audio)
     # 如果设备刚刚被唤醒，短暂忽略VAD检测
@@ -30,13 +34,13 @@ async def handleAudioMessage(conn, audio):
     await conn.asr.receive_audio(conn, audio, have_voice)
 
 
-async def resume_vad_detection(conn):
+async def resume_vad_detection(conn: "ConnectionHandler"):
     # 等待2秒后恢复VAD检测
     await asyncio.sleep(2)
     conn.just_woken_up = False
 
 
-async def startToChat(conn, text):
+async def startToChat(conn: "ConnectionHandler", text):
     # 检查输入是否是JSON格式（包含说话人信息）
     speaker_name = None
     language_tag = None
@@ -96,7 +100,7 @@ async def startToChat(conn, text):
     conn.executor.submit(conn.chat, actual_text)
 
 
-async def no_voice_close_connect(conn, have_voice):
+async def no_voice_close_connect(conn: "ConnectionHandler", have_voice):
     if have_voice:
         conn.last_activity_time = time.time() * 1000
         return
@@ -123,7 +127,7 @@ async def no_voice_close_connect(conn, have_voice):
             await startToChat(conn, prompt)
 
 
-async def max_out_size(conn):
+async def max_out_size(conn: "ConnectionHandler"):
     # 播放超出最大输出字数的提示
     conn.client_abort = False
     text = "不好意思，我现在有点事情要忙，明天这个时候我们再聊，约好了哦！明天不见不散，拜拜！"
@@ -134,7 +138,7 @@ async def max_out_size(conn):
     conn.close_after_chat = True
 
 
-async def check_bind_device(conn):
+async def check_bind_device(conn: "ConnectionHandler"):
     if conn.bind_code:
         # 确保bind_code是6位数字
         if len(conn.bind_code) != 6:

main/xiaozhi-server/core/handle/reportHandle.py

@@ -11,13 +11,17 @@
 
 import time
 import opuslib_next
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from core.connection import ConnectionHandler
 
 from config.manage_api_client import report as manage_report
 
 TAG = __name__
 
 
-async def report(conn, type, text, opus_data, report_time):
+async def report(conn: "ConnectionHandler", type, text, opus_data, report_time):
     """执行聊天记录上报操作
 
     Args:
@@ -45,7 +49,7 @@ async def report(conn, type, text, opus_data, report_time):
         conn.logger.bind(tag=TAG).error(f"聊天记录上报失败: {e}")
 
 
-def opus_to_wav(conn, opus_data):
+def opus_to_wav(conn: "ConnectionHandler", opus_data):
     """将Opus数据转换为WAV格式的字节流
 
     Args:
@@ -100,7 +104,7 @@ def opus_to_wav(conn, opus_data):
                 conn.logger.bind(tag=TAG).debug(f"释放decoder资源时出错: {e}")
 
 
-def enqueue_tts_report(conn, text, opus_data):
+def enqueue_tts_report(conn: "ConnectionHandler", text, opus_data):
     if not conn.read_config_from_api or conn.need_bind or not conn.report_tts_enable:
         return
     if conn.chat_history_conf == 0:
@@ -128,7 +132,7 @@ def enqueue_tts_report(conn, text, opus_data):
         conn.logger.bind(tag=TAG).error(f"加入TTS上报队列失败: {text}, {e}")
 
 
-def enqueue_asr_report(conn, text, opus_data):
+def enqueue_asr_report(conn: "ConnectionHandler", text, opus_data):
     if not conn.read_config_from_api or conn.need_bind or not conn.report_asr_enable:
         return
     if conn.chat_history_conf == 0:

main/xiaozhi-server/core/handle/sendAudioHandle.py

@@ -1,6 +1,10 @@
 import json
 import time
 import asyncio
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from core.connection import ConnectionHandler
 from core.utils import textUtils
 from core.utils.util import audio_to_data
 from core.providers.tts.dto.dto import SentenceType
@@ -13,7 +17,7 @@
 PRE_BUFFER_COUNT = 5
 
 
-async def sendAudioMessage(conn, sentenceType, audios, text):
+async def sendAudioMessage(conn: "ConnectionHandler", sentenceType, audios, text):
     if conn.tts.tts_audio_first_sentence:
         conn.logger.bind(tag=TAG).info(f"发送第一段语音: {text}")
         conn.tts.tts_audio_first_sentence = False
@@ -47,7 +51,7 @@ async def sendAudioMessage(conn, sentenceType, audios, text):
             await conn.close()
 
 
-async def _wait_for_audio_completion(conn):
+async def _wait_for_audio_completion(conn: "ConnectionHandler"):
     """
     等待音频队列清空并等待预缓冲包播放完成
 
@@ -70,7 +74,9 @@ async def _wait_for_audio_completion(conn):
         conn.logger.bind(tag=TAG).debug("音频发送完成")
 
 
-async def _send_to_mqtt_gateway(conn, opus_packet, timestamp, sequence):
+async def _send_to_mqtt_gateway(
+    conn: "ConnectionHandler", opus_packet, timestamp, sequence
+):
     """
     发送带16字节头部的opus数据包给mqtt_gateway
     Args:
@@ -92,7 +98,9 @@ async def _send_to_mqtt_gateway(conn, opus_packet, timestamp, sequence):
     await conn.websocket.send(complete_packet)
 
 
-async def sendAudio(conn, audios, frame_duration=AUDIO_FRAME_DURATION):
+async def sendAudio(
+    conn: "ConnectionHandler", audios, frame_duration=AUDIO_FRAME_DURATION
+):
     """
     发送音频包，使用 AudioRateController 进行精确的流量控制
 
@@ -121,7 +129,9 @@ async def sendAudio(conn, audios, frame_duration=AUDIO_FRAME_DURATION):
     )
 
 
-def _get_or_create_rate_controller(conn, frame_duration, is_single_packet):
+def _get_or_create_rate_controller(
+    conn: "ConnectionHandler", frame_duration, is_single_packet
+):
     """
     获取或创建 RateController 和 flow_control
 
@@ -177,7 +187,7 @@ def _get_or_create_rate_controller(conn, frame_duration, is_single_packet):
     return conn.audio_rate_controller, conn.audio_flow_control
 
 
-def _start_background_sender(conn, rate_controller, flow_control):
+def _start_background_sender(conn: "ConnectionHandler", rate_controller, flow_control):
     """
     启动后台发送循环任务
 
@@ -201,7 +211,7 @@ async def send_callback(packet):
 
 
 async def _send_audio_with_rate_control(
-    conn, audio_list, rate_controller, flow_control, send_delay
+    conn: "ConnectionHandler", audio_list, rate_controller, flow_control, send_delay
 ):
     """
     使用 rate_controller 发送音频包
@@ -233,7 +243,7 @@ async def _send_audio_with_rate_control(
             rate_controller.add_audio(packet)
 
 
-async def _do_send_audio(conn, opus_packet, flow_control):
+async def _do_send_audio(conn: "ConnectionHandler", opus_packet, flow_control):
     """
     执行实际的音频发送
     """
@@ -254,7 +264,7 @@ async def _do_send_audio(conn, opus_packet, flow_control):
     flow_control["sequence"] = sequence + 1
 
 
-async def send_tts_message(conn, state, text=None):
+async def send_tts_message(conn: "ConnectionHandler", state, text=None):
     """发送 TTS 状态消息"""
     if text is None and state == "sentence_start":
         return
@@ -281,7 +291,7 @@ async def send_tts_message(conn, state, text=None):
     await conn.websocket.send(json.dumps(message))
 
 
-async def send_stt_message(conn, text):
+async def send_stt_message(conn: "ConnectionHandler", text):
     """发送 STT 状态消息"""
     end_prompt_str = conn.config.get("end_prompt", {}).get("prompt")
     if end_prompt_str and end_prompt_str == text: