saw-go/kbase_py/im_ws.py

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()