newbot/logic/logic_ui.py

import hashlib
import json
import time
import tkinter
import uuid
from datetime import datetime
from tkinter import messagebox

import requests

from common.sql_lite import get_global_db_connection
from config import conf, set_config
from db.batch_task import batch_task_get_list, batch_task_status, batch_task_update_status
from logic.logic_batch_task_create import open_batch_task_create_win
from logic.logic_batch_task_detail import open_batch_task_detail_win
from plugins.agent import Agent
from plugins.custom_agent import CustomAgent
from service.batch_task import stop_batch_task, start_batch_task
from service.robot import get_robot, init_robot
from ui.ui import WinGUI


class Win(WinGUI):
    def __init__(self):
        super().__init__()
        self.__event_bind()
        self.__style_config()
        self.__is_started = False
        self.selection_batch_task_id = None

    def __event_bind(self):
        self.tk_tabs_main_tabs.tk_tabs_start.tk_button_save.bind('<Button-1>', self.save_event)
        self.tk_tabs_main_tabs.tk_tabs_start.tk_button_start.bind('<Button-1>', self.start_event)
        self.tk_tabs_main_tabs.tk_tabs_start.tk_button_pause.bind('<Button-1>', self.stop_event)
        self.tk_tabs_main_tabs.tk_tabs_start.tk_button_version.bind('<Button-1>', self.version_event)
        self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_button_create.bind('<Button-1>', self.batch_task_create_event)
        self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.bind('<Double-1>',
                                                                                self.batch_task_detail_event)
        self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.bind('<Button-1>', self.just_click_event)
        self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.bind('<Button-3>',
                                                                                self.batch_task_action_event)
        pass

    def __style_config(self):
        pass

    def version_event(self, event):
        messagebox.showinfo('版本信息', '当前版本：v1.0.0')

    def stop_event(self, event):
        stop_batch_task()
        if self.tk_tabs_main_tabs.tk_tabs_start.tk_button_pause.cget('state').__str__() == tkinter.DISABLED:
            return
        robot = get_robot()
        if robot is not None:
            robot.wcf.cleanup()
            robot = None
        self.tk_tabs_main_tabs.tk_tabs_start.tk_button_start.config(state=tkinter.NORMAL)
        self.tk_tabs_main_tabs.tk_tabs_start.tk_button_pause.config(state=tkinter.DISABLED)
        self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_button_create.config(state=tkinter.DISABLED)
        self.__is_started = False
        messagebox.showinfo('提示', '助手已停止运行！')

    @staticmethod
    def check_token():
        # 获取当前网卡mac地址
        mac = uuid.getnode()
        url = "https://wxadminapi.gkscrm.com/wechat-api/token/check"

        token = str(conf().get("token"))

        payload = json.dumps({
            "token": token,
            "mac": str(mac)
        })

        headers = {
            'Content-Type': 'application/json'
        }

        response = requests.request("POST", url, headers=headers, data=payload)

        if response.status_code == 200:
            resp = json.loads(response.text)

            if resp['valid'] == True:
                sign = resp['sign']
                timestamp = resp['timestamp']
                # 获取当前unix时间戳
                current_timestamp = int(time.time())

                if abs(current_timestamp - int(timestamp)) > 180:
                    messagebox.showerror('错误', '验证失败！请重试')
                    return False
                # 创建一个MD5哈希对象
                md5_hash = hashlib.md5()
                # 更新哈希对象
                md5_hash.update(token.encode('utf-8'))
                # 获取十六进制格式的MD5哈希值
                md5_token = md5_hash.hexdigest()

                md5_hash = hashlib.md5()
                md5_hash.update(str(mac).encode('utf-8'))
                md5_mac = md5_hash.hexdigest()

                md5_hash = hashlib.md5()
                md5_hash.update(str(timestamp).encode('utf-8'))
                md5_timestamp = md5_hash.hexdigest()

                md5_hash = hashlib.md5()
                md5_hash.update((md5_token + md5_mac + md5_timestamp).encode('utf-8'))
                md5_sign = md5_hash.hexdigest()

                if md5_sign != sign:
                    messagebox.showerror('错误', 'token验证失败！')
                    return False

                agent_info = resp['agent_info']
                if agent_info is None:
                    return False

                set_config("agent_id", agent_info.get("id", 0))
                set_config("role", agent_info.get("role", ""))
                set_config("background", agent_info.get("background", ""))
                set_config("examples", agent_info.get("examples", ""))
                set_config("dataset_id", agent_info.get("dataset_id", ""))
                set_config("collection_id", agent_info.get("collection_id", ""))
                set_config("custom_agent_base", resp['custom_agent_base'])
                set_config("custom_agent_key", resp['custom_agent_key'])
                set_config("openai_base", resp['openai_base'])
                set_config("openai_key", resp['openai_key'])
                set_config("dataset_base", resp['dataset_base'])
                set_config("dataset_key", resp['dataset_key'])

                return True

            else:
                messagebox.showerror('错误',
                                     'token已失效 或 与当前设备未绑定，每个token只能绑定一台设备！请填入和当前设备绑定的token，或者获取新的token！')
                return False
        else:
            messagebox.showerror('错误', '您的网络状态异常！请稍候重试')
            return False

    def start_event(self, event):
        if self.tk_tabs_main_tabs.tk_tabs_start.tk_button_start.cget('state').__str__() == tkinter.DISABLED:
            return

        if not self.check_token():
            return

        # 接收消息
        # robot.enableRecvMsg()     # 可能会丢消息？
        # init_robot()
        robot = get_robot()
        robot.enableReceivingMsg()  # 加队列

        if robot.wcf.is_login():
            self.__is_started = True
            self.tk_tabs_main_tabs.tk_tabs_start.tk_button_start.config(state='disabled')
            self.tk_tabs_main_tabs.tk_tabs_start.tk_button_pause.config(state='normal')
            self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_button_create.config(state='normal')
            messagebox.showinfo('提示', '助手开始运行！')
            start_batch_task()

            # 注册插件
            if conf().get("agent_id") != 0 and conf().get("openai_base") and conf().get("openai_key"):
                robot.register_plugin(Agent())

            if conf().get("agent_id") == 0 and conf().get("custom_agent_base") and conf().get("custom_agent_key"):
                robot.register_plugin(CustomAgent())

            wx_wxid = robot.wcf.get_self_wxid()
            connection = get_global_db_connection()

            def refresh_list():
                # 清空列表
                self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.delete(
                    *self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.get_children())

                # 查询数据
                cursor = connection.cursor()
                try:
                    results = batch_task_get_list(cursor, wx_wxid)
                finally:
                    cursor.close()
                for result in results:
                    created_at = datetime.fromtimestamp(result["created_at"]).strftime("%Y-%m-%d %H:%M:%S")
                    status = batch_task_status[result["status"]]
                    values = (
                    created_at, status, f"{result['success'] + result['fail']}/{result['total']}", result['fail'],
                    result['content'])
                    self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.insert('', "end",
                                                                                              iid=result['id'],
                                                                                              values=values)
                    # 维持选中状态
                    if result['id'] == self.selection_batch_task_id:
                        self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.selection_set(result['id'])

                # 每 3秒调用一次这个函数
                self.tk_tabs_main_tabs.after(3000, refresh_list)

            self.tk_tabs_main_tabs.after(0, refresh_list)

    def save_event(self, event):
        conf().update({
            # "api_base": self.tk_tabs_main_tabs.tk_tabs_start.tk_input_api_base.get(),
            # "api_key": self.tk_tabs_main_tabs.tk_tabs_start.tk_input_api_key.get(),
            "token": self.tk_tabs_main_tabs.tk_tabs_start.tk_input_token.get()
        })
        # 将字典写入 JSON 文件
        with open('config.json', 'w') as json_file:
            json.dump(conf(), json_file, indent=4)

        messagebox.showinfo('提示', '保存文件成功！')

    def batch_task_create_event(self, event):
        if self.__is_started:
            open_batch_task_create_win()

    def batch_task_detail_event(self, event):
        # 获取双击项的标识符
        item_id = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.identify_row(event.y)
        # 获取该项的值
        item_values = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.item(item_id, "values")
        # print("Double-clicked item:", item_values)
        if len(item_values) > 0:
            self.selection_batch_task_id = int(item_id)
            open_batch_task_detail_win(int(item_id), item_values)
        else:
            self.selection_batch_task_id = None

    def batch_task_start_event(self, event):
        # 获取双击项的标识符
        item_id = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.identify_row(event.y)
        # 获取该项的值
        item_values = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.item(item_id, "values")
        # print("Double-clicked item:", item_values)
        if len(item_values) > 0:
            self.selection_batch_task_id = int(item_id)
            connection = get_global_db_connection()
            cursor = connection.cursor()
            try:
                batch_task_update_status(cursor, int(item_id), 1)
                connection.commit()
            except Exception as e:
                # 回滚事务
                connection.rollback()
                print(f"发生错误: {e}")
            finally:
                # 确保资源被正确释放
                cursor.close()
        else:
            self.selection_batch_task_id = None

    def batch_task_stop_event(self, event):
        # 获取双击项的标识符
        item_id = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.identify_row(event.y)
        # 获取该项的值
        item_values = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.item(item_id, "values")
        # print("Double-clicked item:", item_values)
        if len(item_values) > 0:
            self.selection_batch_task_id = int(item_id)
            connection = get_global_db_connection()
            cursor = connection.cursor()
            try:
                batch_task_update_status(cursor, int(item_id), 3)
                connection.commit()
            except Exception as e:
                # 回滚事务
                connection.rollback()
                print(f"发生错误: {e}")
            finally:
                # 确保资源被正确释放
                cursor.close()
        else:
            self.selection_batch_task_id = None

    def batch_task_action_event(self, event):
        # 获取右键项的标识符
        item_id = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.identify_row(event.y)
        print(f"Right-clicked item: {item_id}")
        self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.selection_set(item_id)
        # 获取该项的值
        item_values = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.item(item_id, "values")
        if len(item_values) > 0:
            self.selection_batch_task_id = int(item_id)

            # print("Right-clicked item:", item_values)

            # 创建菜单
            menu = tkinter.Menu(self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list, tearoff=0)
            menu.add_command(label="查看详情", command=lambda: self.batch_task_detail_event(event))
            if item_values[1] == "等待中" or item_values[1] == "已开始":
                menu.add_command(label="停止", command=lambda: self.batch_task_stop_event(event))
            elif item_values[1] == "已停止":
                menu.add_command(label="开始", command=lambda: self.batch_task_start_event(event))
            menu.post(event.x_root, event.y_root)
        else:
            self.selection_batch_task_id = None

    def just_click_event(self, event):
        item_id = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.identify_row(event.y)
        self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.selection_set(item_id)
        item_values = self.tk_tabs_main_tabs.tk_tabs_batch_task.tk_table_batch_task_list.item(item_id, "values")
        if len(item_values) > 0:
            self.selection_batch_task_id = int(item_id)
        else:
            self.selection_batch_task_id = None