126 lines
3.9 KiB
Python
126 lines
3.9 KiB
Python
import websocket
|
|
import time
|
|
import json
|
|
import fileChroma
|
|
import threading
|
|
import fileCon
|
|
|
|
config = {}
|
|
stop_event = False
|
|
|
|
def on_open(ws):
|
|
print("连接已建立")
|
|
|
|
def on_message(ws, message):
|
|
#print(f"接收到服务器消息: {message}")
|
|
print("接收到服务器消息大小：", len(message))
|
|
res = json.loads(message)
|
|
#print("req:", res)
|
|
results = fileChroma.KBaseChroma(res)
|
|
resp ={}
|
|
resp["im_context"] = res["im_context"]
|
|
resp["knowledge_base"] = res["knowledge_base"]
|
|
resp["query_select"] = results
|
|
#print("resp:",resp)
|
|
respStr = json.dumps(resp, ensure_ascii=False)
|
|
print("发送到服务器消息大小：", len(respStr))
|
|
send_message(ws, respStr) # 发送消息
|
|
|
|
#发送消息
|
|
def send_message(ws, message):
|
|
try:
|
|
ws.send(message)
|
|
#print(f"发送消息: {message}")
|
|
except Exception as e:
|
|
print(f"发送消息失败: {e}")
|
|
|
|
def on_close(ws, close_status_code, close_msg):
|
|
if stop_event:
|
|
print("WebSocket 已关闭")
|
|
return
|
|
global config
|
|
print("连接已关闭，尝试重新连接...")
|
|
MAX_RETRIES = config.get("ws_max_retries", 5) # 最大重连次数
|
|
RETRY_DELAY = config.get("ws_retry_delay", 5) # 重连延迟时间（秒）
|
|
retries = 0
|
|
while retries < MAX_RETRIES:
|
|
try:
|
|
time.sleep(RETRY_DELAY)
|
|
ws = websocket.WebSocketApp(config["ws_url"]+"&kbase_server_id="+config["kbase_server_id"],
|
|
on_open=on_open,
|
|
on_message=on_message,
|
|
on_close=on_close,
|
|
# send_message=send_message,
|
|
on_error=on_error)
|
|
ws.run_forever()
|
|
break
|
|
except Exception as e:
|
|
retries += 1
|
|
print(f"重连尝试 {retries} 失败: {e}")
|
|
if retries == MAX_RETRIES:
|
|
print("达到最大重连次数，停止尝试。")
|
|
|
|
def on_error(ws, error):
|
|
print(f"发生错误: {error}")
|
|
|
|
def readConfig():
|
|
"""读取配置文件"""
|
|
try:
|
|
with open(r'E:\Code\saw-ai\python-sag\saw-ai-py.json', 'r', encoding='utf-8') as file:
|
|
config = json.load(file)
|
|
return config
|
|
except FileNotFoundError:
|
|
print("配置文件未找到")
|
|
return None
|
|
except json.JSONDecodeError:
|
|
print("配置文件格式错误")
|
|
return None
|
|
|
|
def writeConfig(config):
|
|
"""写入配置文件"""
|
|
try:
|
|
with open(r'E:\Code\saw-ai\python-sag\saw-ai-py.json', 'w', encoding='utf-8') as file:
|
|
json.dump(config, file, ensure_ascii=False, indent=4)
|
|
except Exception as e:
|
|
print(f"写入配置文件失败: {e}")
|
|
|
|
#监听用户输入
|
|
def input_listener(ws):
|
|
while True:
|
|
user_input = input()
|
|
if user_input.lower() == "stop":
|
|
stop_event = True
|
|
ws.close()
|
|
break
|
|
|
|
def sfile_content():
|
|
#每10秒获取转换一次
|
|
while True:
|
|
time.sleep(10)
|
|
# 文件内容转换
|
|
fileCon.get_file_add_convert()
|
|
if stop_event:
|
|
break
|
|
|
|
if __name__ == "__main__":
|
|
config = readConfig()
|
|
if config is None:
|
|
print("无法读取配置文件")
|
|
exit(1)
|
|
print("config:",config)
|
|
# 创建 WebSocket 连接
|
|
ws = websocket.WebSocketApp(config["ws_url"]+"&kbase_server_id="+config["kbase_server_id"],
|
|
on_open=on_open,
|
|
on_message=on_message,
|
|
on_close=on_close,
|
|
# send_message=send_message,
|
|
on_error=on_error)
|
|
|
|
input_thread = threading.Thread(target=input_listener, args=(ws,))
|
|
input_thread.start()
|
|
if config["file_convert"]:
|
|
file_thread = threading.Thread(target=sfile_content)
|
|
file_thread.start()
|
|
# 运行 WebSocket 连接
|
|
ws.run_forever()
|
|
|