seg_yinzy
/
dify-on-wechat
mirror of https://github.com/fwasdfw/dify-on-wechat.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: Add configuration items to support custom data directories and facilitate the storage of itchat.pkl

master
李超 2023-04-09 11:59:31 +08:00
parent 720de9d73f
commit 29fbf69945
2 changed files with 15 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
channel/wechat/wechat_channel.py
View File

@ -18,7 +18,7 @@ from lib import itchat
from lib.itchat.content import * from lib.itchat.content import *
from bridge.reply import * from bridge.reply import *
from bridge.context import * from bridge.context import *
from config import conf from config import conf, get_data_path
from common.time_check import time_checker from common.time_check import time_checker
from common.expired_dict import ExpiredDict from common.expired_dict import ExpiredDict
from plugins import * from plugins import *
@ -97,13 +97,14 @@ class WechatChannel(ChatChannel):
itchat.instance.receivingRetryCount = 600 # 修改断线超时时间 itchat.instance.receivingRetryCount = 600 # 修改断线超时时间
# login by scan QRCode # login by scan QRCode
hotReload = conf().get('hot_reload', False) hotReload = conf().get('hot_reload', False)
status_path = os.path.join(get_data_path(), "itchat.pkl")
try: try:
itchat.auto_login(enableCmdQR=2, hotReload=hotReload, qrCallback=qrCallback) itchat.auto_login(enableCmdQR=2, hotReload=hotReload, statusStorageDir=status_path, qrCallback=qrCallback)
except Exception as e: except Exception as e:
if hotReload: if hotReload:
logger.error("Hot reload failed, try to login without hot reload") logger.error("Hot reload failed, try to login without hot reload")
itchat.logout() itchat.logout()
os.remove("itchat.pkl") os.remove(status_path)
itchat.auto_login(enableCmdQR=2, hotReload=hotReload, qrCallback=qrCallback) itchat.auto_login(enableCmdQR=2, hotReload=hotReload, qrCallback=qrCallback)
else: else:
raise e raise e

14
config.py
View File

@ -93,6 +93,8 @@ available_setting = {
"debug": False, # 是否开启debug模式，开启后会打印更多日志 "debug": False, # 是否开启debug模式，开启后会打印更多日志
"config_data_path": "", # 数据目录
# 插件配置 # 插件配置
"plugin_trigger_prefix": "$", # 规范插件提供聊天相关指令的前缀，建议不要和管理员指令前缀"#"冲突 "plugin_trigger_prefix": "$", # 规范插件提供聊天相关指令的前缀，建议不要和管理员指令前缀"#"冲突
} }
@ -130,7 +132,7 @@ class Config(dict):
def load_user_datas(self): def load_user_datas(self):
try: try:
with open('user_datas.pkl', 'rb') as f: with open(os.path.join(get_data_path(), 'user_datas.pkl'), 'rb') as f:
self.user_datas = pickle.load(f) self.user_datas = pickle.load(f)
logger.info("[Config] User datas loaded.") logger.info("[Config] User datas loaded.")
except FileNotFoundError as e: except FileNotFoundError as e:
@ -141,7 +143,7 @@ class Config(dict):
def save_user_datas(self): def save_user_datas(self):
try: try:
with open('user_datas.pkl', 'wb') as f: with open(os.path.join(get_data_path(), 'user_datas.pkl'), 'wb') as f:
pickle.dump(self.user_datas, f) pickle.dump(self.user_datas, f)
logger.info("[Config] User datas saved.") logger.info("[Config] User datas saved.")
except Exception as e: except Exception as e:
@ -196,6 +198,12 @@ def read_file(path):
with open(path, mode='r', encoding='utf-8') as f: with open(path, mode='r', encoding='utf-8') as f:
return f.read() return f.read()
def conf(): def conf():
return config return config
def get_data_path():
data_path = os.path.join(get_root(), conf().get('config_data_path', ""))
if not os.path.exists(data_path):
logger.info("[INIT] data path not exists, create it: {}".format(data_path))
os.makedirs(data_path)
return data_path