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 db.sop import sop_task_get_list, sop_task_status, sop_task_type, sop_task_query, sop_task_create, \
    sop_task_update_status, sop_stage_create_many, sop_node_create_many
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 plugins.sop import Sop
from service.batch_task import stop_batch_task, start_batch_task
from service.sop_service import start_sop_task, stop_sop_task
from service.robot import get_robot, init_robot
from service.sop_service import get_sop_task_list, get_sop_stage_list, get_sop_node_list
from ui.ui import WinGUI
from common.log import logger

class Win(WinGUI):
    def __init__(self):
        super().__init__()
        self.__event_bind()
        self.__style_config()
        self.__is_started = False
        self.selection_batch_task_id = None
        self.selection_task_id = None
        self.wx_wxid = None
        self.robot = 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)

        # SOP管理
        self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list.bind('<Button-3>', self.sop_task_action_event)
        self.tk_tabs_main_tabs.tk_tabs_sop.tk_button_sync.bind('<Button-1>', self.sop_task_sync)

        pass

    def __style_config(self):
        pass

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

    def stop_event(self, event):
        stop_batch_task()
        stop_sop_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()
        self.robot = get_robot()
        self.robot.enableReceivingMsg()  # 加队列

        if self.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()
            start_sop_task()

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

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

            # 注册SOP插件
            self.robot.register_plugin(Sop())

            self.wx_wxid = self.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, self.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)

            # sop_list 获取列表内容
            self.sop_task_get_list_all()

    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

    # 点击同步后台SOP到本地DB（每隔3s刷新一次列表内容）
    def sop_task_get_list_all(self):
        def refresh_sop_task():
            # 清空列表
            self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list.delete(
                *self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list.get_children()
            )

            # 查询数据
            connection = get_global_db_connection()
            cursor = connection.cursor()
            try:
                results = sop_task_get_list(cursor, self.wx_wxid)
                # logger.info(f"get sop_task list results:{results}")
            finally:
                cursor.close()

            # 把查询到的数据显示在table里
            for result in results:
                created_at = datetime.fromtimestamp(result["created_at"]).strftime("%Y-%m-%d %H:%M:%S")
                status = sop_task_status[result["status"]]
                sop_type = sop_task_type[result["type"]]

                values = (result['id'], result['name'], result['bot_wxid_list'], status, created_at, sop_type)
                self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list.insert('', "end",
                                                                  iid=result['id'],
                                                                  values=values)
                # 维持选中状态
                if result['id'] == self.selection_task_id:
                    self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list.selection_set(result['id'])

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

        self.tk_tabs_main_tabs.after(0, refresh_sop_task)

    # 点击同步后台SOP到本地DB
    def sop_task_sync(self, event):
        if not self.__is_started:
            return

        # logger.info(f"[SOP] 点击同步 {event}")
        connection = get_global_db_connection()
        cursor = connection.cursor()
        sop_task_list = get_sop_task_list()

        if sop_task_list is not None:
            for k in sop_task_list:
                try:
                    sop_task = sop_task_query(cursor, k['id'])
                    # logger.info(f"sop_task={sop_task}")
                    if sop_task is None:
                        _ = sop_task_create(cursor, k, self.wx_wxid)

                        # 根据 sop_task id 获取 sop_stage 列表
                        sop_stage_list = get_sop_stage_list([k['id']])
                        # logger.info(f"sop_stage_list: {sop_stage_list}")
                        if sop_stage_list is not None:
                            sop_stage_create_many(cursor, sop_stage_list)
                            sop_stage_ids = []
                            for val in sop_stage_list:
                                sop_stage_ids.append(val['id'])
                            # logger.info(f"sop_stage_ids={sop_stage_ids}")

                            # 根据 sop_stage_ids 获取 sop_node 列表
                            if len(sop_stage_ids) > 0:
                                sop_node_list = get_sop_node_list(sop_stage_ids)
                                # logger.info(f"sop_stage_ids={sop_stage_ids}")
                                # logger.info(f"sop_node_list={sop_node_list}")
                                if sop_node_list is not None:
                                    sop_node_create_many(cursor, sop_node_list)

                        connection.commit()

                except Exception as e:
                    # 回滚事务
                    connection.rollback()
                    logger.error(f"SQLite DB发生错误: {e}")
                finally:
                    # 确保资源被正确释放
                    cursor.close()

        messagebox.showinfo('SOP管理', '同步SOP成功!')
        return

    # 同步后台设置的SOP列表到本地，根据租户ID获取
    def sop_task_action_event(self, event):
        # logger.info(f"event={event}, started={self.__is_started}")

        if not self.__is_started:
            return

        # 获取右键项的标识符
        item_id = self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list.identify_row(event.y)
        # logger.info(f"Right-clicked item_id: {item_id}")
        self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list.selection_set(item_id)
        # 获取该项的值
        item_values = self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list.item(item_id, "values")
        if len(item_values) > 0:
            self.selection_task_id = int(item_id)
            # logger.info(f"Right-clicked item_values:{item_values}")

            # 创建菜单
            menu = tkinter.Menu(self.tk_tabs_main_tabs.tk_tabs_sop.tk_table_sop_list, tearoff=0)
            if item_values[3] == "已开始":
                menu.add_command(label="禁用", command=lambda: self.sop_forbid(event))
            elif item_values[3] == "已禁用":
                menu.add_command(label="开始", command=lambda: self.sop_resume(event))
            menu.post(event.x_root, event.y_root)
        else:
            self.selection_task_id = None
            # logger.info("no item is selected")

    # sop_task 禁用
    def sop_forbid(self, event):
        if not self.__is_started:
            return

        if self.selection_task_id is None:
            return

        connection = get_global_db_connection()
        cursor = connection.cursor()
        try:
            sop_task_update_status(cursor, self.selection_task_id, 5)
            connection.commit()
        except Exception as e:
            logger.error(f"set sop_task forbid exception:{e}")
        finally:
            cursor.close()

    # sop_task 重新启用
    def sop_resume(self, event):
        if not self.__is_started:
            return

        if self.selection_task_id is None:
            return

        connection = get_global_db_connection()
        cursor = connection.cursor()
        try:
            sop_task_update_status(cursor, self.selection_task_id, 3)
            connection.commit()
        except Exception as e:
            logger.error(f"set sop_task resume exception:{e}")
        finally:
            cursor.close()