458 lines
16 KiB
Python
458 lines
16 KiB
Python
import socketio
|
||
import time
|
||
import base64
|
||
import os
|
||
from datetime import datetime
|
||
from typing import Dict, Callable, Any, Optional
|
||
from functools import partial
|
||
from threading import Event
|
||
|
||
class ASEngineClient:
|
||
def __init__(self, server_url: str, namespace: str, id:str):
|
||
"""
|
||
初始化客户端
|
||
:param server_url: 服务器URL，例如 'https://asengine.net'
|
||
:param namespace: 命名空间，例如 '/socket.io/7100eebf7a909e1ee3f4e06766f512f0358ecb6ac51d5434c17d31dd70890dc1'
|
||
"""
|
||
self.server_url = server_url
|
||
self.namespace = namespace
|
||
self.id = id
|
||
self.sid = None
|
||
self.connected = False
|
||
|
||
# 创建SocketIO客户端
|
||
self.sio = socketio.Client()
|
||
|
||
# 存储路由和对应的回调函数
|
||
self.routes: Dict[str, Dict[str, Callable]] = {}
|
||
|
||
self._on_connect_callback = None
|
||
self._on_connect_callback_entry=None
|
||
self._on_connect_callback_initData=None
|
||
|
||
# ACK相关
|
||
self.ack_events: Dict[str, Event] = {} # 存储ack_id和对应的Event对象
|
||
|
||
# 注册基础事件处理函数
|
||
self._register_base_events()
|
||
|
||
def _register_base_events(self):
|
||
"""注册基础事件处理函数"""
|
||
# 连接事件
|
||
@self.sio.on('connect', namespace=self.namespace)
|
||
def on_connect():
|
||
self.connected = True
|
||
self.sid = self.sio.get_sid(namespace=self.namespace)
|
||
print(f"Connected to server. SID: {self.sid}")
|
||
if self._on_connect_callback is not None:
|
||
self._on_connect_callback(
|
||
self,
|
||
self._on_connect_callback_entry,
|
||
init_data=self._on_connect_callback_initData
|
||
)
|
||
|
||
# 断开连接事件
|
||
@self.sio.on('disconnect', namespace=self.namespace)
|
||
def on_disconnect():
|
||
self.connected = False
|
||
self.sid = None
|
||
print("Disconnected from server")
|
||
|
||
@self.sio.on('error', namespace=self.namespace)
|
||
def on_error(data):
|
||
print(f"Received error: {data}")
|
||
self.connected = False
|
||
self.sid = None
|
||
|
||
# 响应事件
|
||
@self.sio.on('response', namespace=self.namespace)
|
||
def on_response(data):
|
||
print(f"Received response: {data}")
|
||
|
||
# ACK事件处理
|
||
@self.sio.on('ack', namespace=self.namespace)
|
||
def on_ack(data):
|
||
ack_id = data.get('ack_id')
|
||
if ack_id and ack_id in self.ack_events:
|
||
print(f"Received ACK for message: {ack_id}")
|
||
event = self.ack_events[ack_id]
|
||
event.set() # 触发事件，通知等待的线程
|
||
|
||
def register_on_connect_entry(self, callback: Callable, entry: Any, init_data: Any)->bool:
|
||
"""注册连接事件回调函数
|
||
|
||
参数:
|
||
- callback: 连接事件回调函数，接受 **kwargs 参数
|
||
- entry: 回调关联的条目对象
|
||
- init_data: 初始化数据
|
||
|
||
回调关键字参数:
|
||
- entry: 回调关联的条目对象
|
||
- init_data: 初始化数据
|
||
"""
|
||
if not callable(callback):
|
||
return False
|
||
self._on_connect_callback = callback
|
||
self._on_connect_callback_entry=entry
|
||
self._on_connect_callback_initData=init_data
|
||
return True
|
||
|
||
|
||
def register_route_with_entry(self, route: str,
|
||
callback: Optional[Callable] = None,
|
||
entry: Any = None,
|
||
init_data: Optional[Any] = None) -> bool:
|
||
"""注册带条目的路由回调函数
|
||
|
||
参数:
|
||
- route: 路由名称
|
||
- callback: 回调函数，接受 **kwargs 参数
|
||
- entry: 回调关联的条目对象
|
||
- init_data: 初始化数据
|
||
|
||
回调关键字参数:
|
||
- entry: 回调关联的条目对象
|
||
- data: 接收到的消息数据
|
||
- init_data: 初始化数据
|
||
"""
|
||
if not route:
|
||
print("Route cannot be empty")
|
||
return False
|
||
|
||
if callback is None:
|
||
bound_cb = self._default_callback
|
||
else:bound_cb = callback
|
||
self.routes[route] = {
|
||
'callback': bound_cb,
|
||
'entry': entry,
|
||
'init_data': init_data,
|
||
}
|
||
|
||
@self.sio.on(route, namespace=self.namespace)
|
||
def on_route_message(data, _route=route):
|
||
if not self.connected: return
|
||
route_info = self.routes.get(_route)
|
||
if not route_info:
|
||
return
|
||
cb = route_info['callback'] # 已经绑定过 entry
|
||
init = route_info.get('init_data')
|
||
entry = route_info.get('entry')
|
||
try:
|
||
cb(self, entry, data=data, init_data=init) # 传递关键字参数
|
||
except Exception as e:
|
||
print(f"Callback error on route '{_route}': {e}")
|
||
|
||
return True
|
||
|
||
|
||
def register_route(self, route: str, callback: Optional[Callable] = None, init_data: Optional[Any] = None):
|
||
"""注册路由回调函数
|
||
|
||
参数:
|
||
- route: 路由名称
|
||
- callback: 回调函数，接受 **kwargs 参数
|
||
- init_data: 初始化数据
|
||
|
||
回调关键字参数:
|
||
- entry: 回调关联的条目对象（此处为 None）
|
||
- data: 接收到的消息数据
|
||
- init_data: 初始化数据
|
||
"""
|
||
if not route:
|
||
print("Route cannot be empty")
|
||
return False
|
||
|
||
# 存储路由和回调
|
||
self.routes[route] = {
|
||
'callback': callback if callback else self._default_callback,
|
||
'entry': None,
|
||
'init_data': init_data
|
||
}
|
||
|
||
# 注册事件处理器
|
||
@self.sio.on(route, namespace=self.namespace)
|
||
def on_route_message(data):
|
||
print(f"Received message on route '{route}': {data}")
|
||
if 'callback' in self.routes[route]:
|
||
try:
|
||
self.routes[route]['callback'](
|
||
self,
|
||
entry=None,
|
||
data=data,
|
||
init_data=init_data
|
||
)
|
||
except Exception as e:
|
||
print(f"Error in callback for route '{route}': {e}")
|
||
|
||
print(f"Route '{route}' registered successfully")
|
||
return True
|
||
|
||
def _default_callback(self, **kwargs):
|
||
"""默认回调函数
|
||
|
||
关键字参数:
|
||
- entry: 回调关联的条目对象
|
||
- data: 接收到的消息数据
|
||
- init_data: 初始化数据
|
||
"""
|
||
data = kwargs.get('data', {})
|
||
print(f"Default callback received: {data}")
|
||
|
||
def connect(self, timeout: int = 5) -> bool:
|
||
"""
|
||
连接到服务器
|
||
:param timeout: 超时时间(秒)
|
||
:return: 是否连接成功
|
||
"""
|
||
try:
|
||
self.sio.connect(
|
||
self.server_url+"?id="+self.id,
|
||
namespaces=[self.namespace],
|
||
wait_timeout=timeout
|
||
)
|
||
# 等待连接建立
|
||
start_time = time.time()
|
||
while not self.connected and (time.time() - start_time) < timeout:
|
||
print(f"Connecting To {self.server_url}")
|
||
time.sleep(0.2)
|
||
return self.connected
|
||
except Exception as e:
|
||
print(f"Connection failed: {e}")
|
||
return False
|
||
|
||
def disconnect(self):
|
||
"""断开与服务器的连接"""
|
||
if self.connected:
|
||
self.sio.disconnect()
|
||
self.connected = False
|
||
self.sid = None
|
||
print("Disconnected from server")
|
||
|
||
def send_text(self, route: str, text: str, id:str="", wait_for_ack: bool = False, timeout: int = 5, return_event: bool = False) -> Any:
|
||
"""
|
||
发送文本消息
|
||
:param route: 目标路由
|
||
:param text: 要发送的文本
|
||
:param id: 发送者ID
|
||
:param wait_for_ack: 是否等待服务器ACK确认
|
||
:param timeout: 等待ACK的超时时间（秒）
|
||
:param return_event: 是否返回事件对象和ack_id而不是等待ACK
|
||
:return: 如果wait_for_ack为True且return_event为False，返回是否收到ACK；如果return_event为True，返回(event, ack_id)；否则返回是否发送成功
|
||
"""
|
||
if not self.connected:
|
||
print("Not connected to server")
|
||
return False
|
||
if id == "":
|
||
id = self.id
|
||
|
||
try:
|
||
# 生成唯一ack_id
|
||
ack_id = f"{id}_{datetime.now().strftime('%Y%m%d%H%M%S%f')}"
|
||
|
||
# 准备发送数据
|
||
data = {
|
||
'type': 'text',
|
||
'data': text,
|
||
'id': id,
|
||
'ack_id': ack_id
|
||
}
|
||
|
||
# 发送数据
|
||
self.sio.emit(route, data, namespace=self.namespace)
|
||
print(f"Sent text to route '{route}': {text[:20]}")
|
||
|
||
# 如果不需要等待ACK，直接返回True
|
||
if not wait_for_ack:
|
||
return True
|
||
|
||
# 初始化ACK事件
|
||
event = Event()
|
||
self.ack_events[ack_id] = event
|
||
|
||
# 如果需要返回事件对象，直接返回
|
||
if return_event:
|
||
return event, ack_id
|
||
|
||
# 等待ACK或超时
|
||
received_ack = event.wait(timeout=timeout)
|
||
|
||
# 清理ACK事件
|
||
del self.ack_events[ack_id]
|
||
|
||
if received_ack:
|
||
return True
|
||
else:
|
||
print(f"ACK timeout for message: {ack_id}")
|
||
return False
|
||
except Exception as e:
|
||
print(f"Failed to send text to route '{route}': {e}")
|
||
return False
|
||
|
||
def send_voice(self, route: str, voice_data: bytes, id:str=None, wait_for_ack: bool = False, timeout: int = 5) -> bool:
|
||
"""
|
||
发送语音数据
|
||
:param route: 目标路由
|
||
:param voice_data: 语音数据字节流
|
||
:param id: 发送者ID
|
||
:param wait_for_ack: 是否等待服务器ACK确认
|
||
:param timeout: 等待ACK的超时时间（秒）
|
||
:return: 是否发送成功（若wait_for_ack为True，则返回是否收到ACK）
|
||
"""
|
||
if id is None:
|
||
id = self.id
|
||
if not self.connected:
|
||
print("Not connected to server")
|
||
return False
|
||
|
||
try:
|
||
# 生成唯一ack_id
|
||
ack_id = f"{id}_{datetime.now().strftime('%Y%m%d%H%M%S%f')}"
|
||
|
||
# 发送数据
|
||
self.sio.emit(
|
||
route,
|
||
{
|
||
'type': 'audio',
|
||
'data': voice_data,
|
||
'timestamp': int(time.time() * 1000),
|
||
'sampleRate': 16000,
|
||
'id': id,
|
||
'ack_id': ack_id
|
||
},
|
||
namespace=self.namespace
|
||
)
|
||
|
||
print(f"Sent voice file to route '{route}'")
|
||
|
||
# 如果不需要等待ACK，直接返回True
|
||
if not wait_for_ack:
|
||
return True
|
||
|
||
# 初始化ACK事件
|
||
event = Event()
|
||
self.ack_events[ack_id] = event
|
||
|
||
# 等待ACK或超时
|
||
received_ack = event.wait(timeout=timeout)
|
||
|
||
# 清理ACK事件
|
||
del self.ack_events[ack_id]
|
||
|
||
if received_ack:
|
||
return True
|
||
else:
|
||
print(f"ACK timeout for message: {ack_id}")
|
||
return False
|
||
except Exception as e:
|
||
print(f"Failed to send voice to route '{route}': {e}")
|
||
return False
|
||
|
||
def send_image(self, route: str, file_path: str, id:str, wait_for_ack: bool = False, timeout: int = 5) -> bool:
|
||
"""
|
||
发送图片文件
|
||
:param route: 目标路由
|
||
:param file_path: 图片文件路径
|
||
:param id: 发送者ID
|
||
:param wait_for_ack: 是否等待服务器ACK确认
|
||
:param timeout: 等待ACK的超时时间（秒）
|
||
:return: 是否发送成功（若wait_for_ack为True，则返回是否收到ACK）
|
||
"""
|
||
if not self.connected:
|
||
print("Not connected to server")
|
||
return False
|
||
|
||
|
||
try:
|
||
# 检查文件是否存在
|
||
if not os.path.exists(file_path):
|
||
print(f"File not found: {file_path}")
|
||
return False
|
||
|
||
# 读取文件并进行base64编码
|
||
with open(file_path, 'rb') as f:
|
||
image_data = 'data:image/jpeg;base64,' + base64.b64encode(f.read()).decode('utf-8')
|
||
|
||
# 获取文件名
|
||
filename = os.path.basename(file_path)
|
||
|
||
# 生成唯一ack_id
|
||
ack_id = f"{id}_{datetime.now().strftime('%Y%m%d%H%M%S%f')}"
|
||
|
||
# 发送数据
|
||
self.sio.emit(
|
||
route,
|
||
{
|
||
'type': 'image',
|
||
'data': image_data,
|
||
'id': id,
|
||
'ack_id': ack_id
|
||
},
|
||
namespace=self.namespace
|
||
)
|
||
print(f"Sent image file to route '{route}': {filename}")
|
||
|
||
# 如果不需要等待ACK，直接返回True
|
||
if not wait_for_ack:
|
||
return True
|
||
|
||
# 初始化ACK事件
|
||
event = Event()
|
||
self.ack_events[ack_id] = event
|
||
|
||
# 等待ACK或超时
|
||
received_ack = event.wait(timeout=timeout)
|
||
|
||
# 清理ACK事件
|
||
del self.ack_events[ack_id]
|
||
|
||
if received_ack:
|
||
return True
|
||
else:
|
||
print(f"ACK timeout for message: {ack_id}")
|
||
return False
|
||
except Exception as e:
|
||
print(f"Failed to send image to route '{route}': {e}")
|
||
return False
|
||
|
||
def run_forever(self):
|
||
"""
|
||
保持客户端运行"""
|
||
if self.connected:
|
||
try:
|
||
while True:
|
||
time.sleep(1)
|
||
except KeyboardInterrupt:
|
||
print("Client stopped by user")
|
||
self.disconnect()
|
||
else:
|
||
print("Not connected to server")
|
||
|
||
|
||
class HSAEngineClient(ASEngineClient):
|
||
def __init__(self, server_url: str, namespace: str, id:str, chatmanager):
|
||
super().__init__(server_url, namespace, id)
|
||
self.chatmanager = chatmanager
|
||
def loadmarkdown(self, fmd, fmdp, fsp):
|
||
now = time.time()
|
||
# 并行发送三个消息，获取事件对象和ack_id
|
||
event1, ack_id1 = self.send_text('markdown-in', fmd, self.id, wait_for_ack=True, return_event=True)
|
||
event2, ack_id2 = self.send_text('markdown-prompt-in', fmdp, self.id, wait_for_ack=True, return_event=True)
|
||
event3, ack_id3 = self.send_text('score-prompt-in', fsp, self.id, wait_for_ack=True, return_event=True)
|
||
self.chatmanager.app.logger.info(f"Load markdown time cost: {time.time() - now}")
|
||
now = time.time()
|
||
# 等待所有事件完成
|
||
timeout = 5 # 超时时间
|
||
success1 = event1.wait(timeout=timeout)
|
||
success2 = event2.wait(timeout=timeout)
|
||
success3 = event3.wait(timeout=timeout)
|
||
self.chatmanager.app.logger.info(f"Wait markdown time cost: {time.time() - now}")
|
||
now = time.time()
|
||
# 清理ACK事件
|
||
if ack_id1 in self.ack_events:
|
||
del self.ack_events[ack_id1]
|
||
if ack_id2 in self.ack_events:
|
||
del self.ack_events[ack_id2]
|
||
if ack_id3 in self.ack_events:
|
||
del self.ack_events[ack_id3]
|
||
return success1 and success2 and success3
|