import json import os import threading import time import traceback from copy import deepcopy from typing import Any, List import websocket from pydantic_core import ValidationError from websocket import create_connection from schemas.context import ASRContext from schemas.stream import StreamDataModel, StreamResultModel from utils.logger import logger IN_TEST = os.getenv("SUBMIT_CONFIG_FILEPATH", None) is None class Client: def __init__(self, sut_url: str, context: ASRContext) -> None: # base_url = "ws://127.0.0.1:5003" self.base_url = sut_url + "/recognition" logger.info(f"{self.base_url}") self.context: ASRContext = deepcopy(context) # if not os.getenv("DATASET_FILEPATH", ""): # self.base_url = "wss://speech.4paradigm.com/aibuds/api/v1/recognition" # self.base_url = "ws://localhost:5003/recognition" self.connect_num = 0 self.exception = False self.close_time = 10**50 self.send_time: List[float] = [] self.recv_time: List[float] = [] self.predict_data: List[Any] = [] self.success = True def action(self): # 如果 5 次初始化都失败，则退出 connect_success = False for i in range(5): try: self._connect_init() connect_success = True break except Exception as e: logger.error(f"第 {i+1} 次连接失败，原因：{e}") time.sleep(int(os.getenv("connect_sleep", 10))) if not connect_success: exit(-1) self.trecv = threading.Thread(target=self._recv) self.trecv.start() self._send() self._close() return self._gen_result() def _connect_init(self): end_time = time.time() + float(os.getenv("end_time", 2)) success = False try: self.ws = create_connection(self.base_url) self.ws.send(json.dumps(self._gen_init_data())) while time.time() < end_time and not success: data = self.ws.recv() logger.info(f"data {data}") if len(data) == 0: time.sleep(1) continue if isinstance(data, str): try: data = json.loads(data) except Exception: raise Exception("初始化阶段，数据不是 json 字符串格式，终止流程") if isinstance(data, dict): success = data.get("success", False) if not success: logger.error(f"初始化失败，返回的结果为 {data}，终止流程") else: break logger.error("初始化阶段，数据不是 json 字符串格式，终止流程") exit(-1) except websocket.WebSocketConnectionClosedException or TimeoutError: raise Exception("初始化阶段连接中断，终止流程") # exit(-1) except ConnectionRefusedError: raise Exception("初始化阶段，连接失败，等待 10s 后重试，最多重试 5 次") # logger.error("初始化阶段，连接失败，等待 10s 后重试，最多重试 5 次") # self.connect_num += 1 # if self.connect_num <= 4: # time.sleep(int(os.getenv("connect_sleep", 10))) # self._connect_init() # success = True # else: # logger.error("初始化阶段连接失败多次") # exit(-1) if not success: # logger.error("初始化阶段 60s 没有返回数据，时间太长，终止流程") raise Exception("初始化阶段 60s 没有返回数据，时间太长，终止流程") else: logger.info("建立连接成功") self.connect_num = 0 def _send(self): send_ts = float(os.getenv("send_interval", 60)) if not self.success: return with open(self.context.file_path, "rb") as fp: wav_data = fp.read() meta_length = wav_data.index(b"data") + 8 try: with open(self.context.file_path, "rb") as fp: # 去掉 wav 文件的头信息 fp.read(meta_length) # 上一段音频的发送时间 last_send_time = -1 # 正文内容 while True: now_time = time.perf_counter() if last_send_time == -1: chunk = fp.read(int(self.context.chunk_size)) else: interval_cnt = max( int((now_time - last_send_time) / self.context.wait_time), 1, ) chunk = fp.read(int(self.context.chunk_size * interval_cnt)) if not chunk: break send_time_start = time.perf_counter() self.ws.send(chunk, websocket.ABNF.OPCODE_BINARY) self.send_time.append(send_time_start) last_send_time = send_time_start send_time_end = time.perf_counter() if send_time_end - send_time_start > send_ts: logger.error(f"发送延迟已经超过 {send_ts}s, 终止当前音频发送") break if (sleep_time := self.context.wait_time + now_time - send_time_end) > 0: time.sleep(sleep_time) logger.info("当条语音数据发送完成") self.ws.send(json.dumps({"end": True})) logger.info("2s 后关闭双向连接.") except BrokenPipeError: logger.error("发送数据出错，被测服务出现故障") except Exception as e: logger.error(f"Exception: {e}") logger.error(f"{traceback.print_exc()}") logger.error("发送数据失败") self.success = False # self.close_time = time.perf_counter() + int(os.getenv("api_timeout", 2)) self.close_time = time.perf_counter() + 20 * 60 def _recv(self): try: while self.ws.connected and self.success: recv_data = self.ws.recv() if isinstance(recv_data, str): if recv_data := str(recv_data): self.recv_time.append(time.perf_counter()) # 识别到最后的合并结果后再关闭 recognition_results = StreamResultModel(**json.loads(recv_data)).recognition_results if ( recognition_results.final_result and recognition_results.start_time == 0 and recognition_results.end_time == 0 and recognition_results.para_seq == 0 ): self.success = False else: self.predict_data.append(recv_data) # if recv_data.recognition_results.final_result and (IN_TEST or os.getenv('test')): # logger.info(f"recv_data {recv_data}") else: self.success = False raise Exception("返回的结果不是字符串形式") except websocket.WebSocketConnectionClosedException: logger.error("WebSocketConnectionClosedException") except ValidationError as e: logger.error("返回的结果不符合格式") logger.error(f"Exception is {e}") os._exit(1) except OSError: pass except Exception: logger.error(f"{traceback.print_exc()}") logger.error("处理被测服务返回数据时出错") self.success = False def _close(self): while time.perf_counter() < self.close_time and self.success: # while not self.success: time.sleep(1) try: self.ws.close() except Exception as e: print(e) pass def _gen_result(self) -> dict: if not self.predict_data: logger.error("没有任何数据返回") self.predict_data = [StreamResultModel(**json.loads(data)).recognition_results for data in self.predict_data] # for item in self.predict_data: # if item.final_result and (IN_TEST or os.getenv('test')): # logger.info(f"recv_data {item}") return { "fail": not self.predict_data, "send_time": self.send_time, "recv_time": self.recv_time, "predict_data": self.predict_data, } def _gen_init_data(self) -> dict: return { "parameter": { "lang": self.context.lang, "sample_rate": self.context.sample_rate, "channel": self.context.channel, "format": self.context.audio_format, "bits": self.context.bits, "enable_words": self.context.enable_words, } }