QChatGPT/pkg/openai/session.py

import logging
import threading
import time
import json

import pkg.openai.manager
import pkg.openai.modelmgr
import pkg.database.manager
import pkg.utils.context

import pkg.plugin.host as plugin_host
import pkg.plugin.models as plugin_models

# 运行时保存的所有session
sessions = {}


class SessionOfflineStatus:
    ON_GOING = 'on_going'
    EXPLICITLY_CLOSED = 'explicitly_closed'


# 从数据加载session
def load_sessions():
    global sessions

    db_inst = pkg.utils.context.get_database_manager()

    session_data = db_inst.load_valid_sessions()

    for session_name in session_data:
        logging.info('加载session: {}'.format(session_name))

        temp_session = Session(session_name)
        temp_session.name = session_name
        temp_session.create_timestamp = session_data[session_name]['create_timestamp']
        temp_session.last_interact_timestamp = session_data[session_name]['last_interact_timestamp']
        temp_session.prompt = json.loads(session_data[session_name]['prompt'])

        sessions[session_name] = temp_session


# 获取指定名称的session，如果不存在则创建一个新的
def get_session(session_name: str):
    global sessions
    if session_name not in sessions:
        sessions[session_name] = Session(session_name)
    return sessions[session_name]


def dump_session(session_name: str):
    global sessions
    if session_name in sessions:
        assert isinstance(sessions[session_name], Session)
        sessions[session_name].persistence()
        del sessions[session_name]


# 通用的OpenAI API交互session
# session内部保留了对话的上下文，
# 收到用户消息后，将上下文提交给OpenAI API生成回复
class Session:
    name = ''

    prompt = {}

    import config

    create_timestamp = 0

    last_interact_timestamp = 0

    just_switched_to_exist_session = False

    response_lock = None

    # 加锁
    def acquire_response_lock(self):
        logging.debug('{},lock acquire,{}'.format(self.name, self.response_lock))
        self.response_lock.acquire()
        logging.debug('{},lock acquire successfully,{}'.format(self.name, self.response_lock))

    # 释放锁
    def release_response_lock(self):
        if self.response_lock.locked():
            logging.debug('{},lock release,{}'.format(self.name, self.response_lock))
            self.response_lock.release()
            logging.debug('{},lock release successfully,{}'.format(self.name, self.response_lock))

    # 从配置文件获取会话预设信息
    def get_default_prompt(self, use_default: str=None):
        config = pkg.utils.context.get_config()

        import pkg.openai.dprompt as dprompt

        if use_default is None:
            current_default_prompt = dprompt.get_prompt(dprompt.get_current())
        else:
            current_default_prompt = dprompt.get_prompt(use_default)

        return [{
            'role': 'system',
            'content': current_default_prompt
        }]

    def __init__(self, name: str):
        self.name = name
        self.create_timestamp = int(time.time())
        self.last_interact_timestamp = int(time.time())
        self.schedule()

        self.response_lock = threading.Lock()
        self.prompt = self.get_default_prompt()

    # 设定检查session最后一次对话是否超过过期时间的计时器
    def schedule(self):
        threading.Thread(target=self.expire_check_timer_loop, args=(self.create_timestamp,)).start()

    # 检查session是否已经过期
    def expire_check_timer_loop(self, create_timestamp: int):
        global sessions
        while True:
            time.sleep(60)

            # 不是此session已更换，退出
            if self.create_timestamp != create_timestamp or self not in sessions.values():
                return

            config = pkg.utils.context.get_config()
            if int(time.time()) - self.last_interact_timestamp > config.session_expire_time:
                logging.info('session {} 已过期'.format(self.name))

                # 触发插件事件
                args = {
                    'session_name': self.name,
                    'session': self,
                    'session_expire_time': config.session_expire_time
                }
                event = pkg.plugin.host.emit(plugin_models.SessionExpired, **args)
                if event.is_prevented_default():
                    return

                self.reset(expired=True, schedule_new=False)

                # 删除此session
                del sessions[self.name]
                return

    # 请求回复
    # 这个函数是阻塞的
    def append(self, text: str) -> str:
        self.last_interact_timestamp = int(time.time())

        # 触发插件事件
        if self.prompt == self.get_default_prompt():
            args = {
                'session_name': self.name,
                'session': self,
                'default_prompt': self.prompt,
            }

            event = pkg.plugin.host.emit(plugin_models.SessionFirstMessageReceived, **args)
            if event.is_prevented_default():
                return None

        config = pkg.utils.context.get_config()
        max_length = config.prompt_submit_length if hasattr(config, "prompt_submit_length") else 1024

        # 向API请求补全
        self.cut_out(text, max_length)
        message = pkg.utils.context.get_openai_manager().request_completion(
            self.prompt
        )

        # 处理回复
        res_test = message
        res_ans = res_test

        # 去除开头可能的提示
        res_ans_spt = res_test.split("\n\n")
        if len(res_ans_spt) > 1:
            del (res_ans_spt[0])
            res_ans = '\n\n'.join(res_ans_spt)

        if config.completion_api_params['model'] in pkg.openai.modelmgr.CHAT_COMPLETION_MODELS:
            self.prompt.append({'role':'assistant', 'content':res_ans})
        elif config.completion_api_params['model'] in pkg.openai.modelmgr.COMPLETION_MODELS:
            self.prompt.append({'role':'', 'content':res_ans})

        if self.just_switched_to_exist_session:
            self.just_switched_to_exist_session = False
            self.set_ongoing()

        return res_ans if res_ans[0]!='\n' else res_ans[1:]

    # 删除上一回合并返回上一回合的问题
    def undo(self) -> str:
        self.last_interact_timestamp = int(time.time())

        # 删除上一回合
        if self.prompt[-1]['role'] != 'user':
            res = self.prompt[-1]['content']
            self.prompt.remove(self.prompt[-2])
        else:
            res = self.prompt[-2]['content']
        self.prompt.remove(self.prompt[-1])

        # 返回上一回合的问题
        return res

    # 构建对话体
    def cut_out(self, msg: str, max_tokens: int) -> str:

        if len(msg) > max_tokens:
            msg = msg[:max_tokens]

        self.prompt.append({
            'role': 'user',
            'content': msg
        })

        logging.debug('cut_out: {}'.format(msg))

    # 持久化session
    def persistence(self):
        if self.prompt == self.get_default_prompt():
            return

        db_inst = pkg.utils.context.get_database_manager()

        name_spt = self.name.split('_')

        subject_type = name_spt[0]
        subject_number = int(name_spt[1])

        db_inst.persistence_session(subject_type, subject_number, self.create_timestamp, self.last_interact_timestamp,
                                    json.dumps(self.prompt))

    # 重置session
    def reset(self, explicit: bool = False, expired: bool = False, schedule_new: bool = True, use_prompt: str = None):
        if self.prompt[-1]['role'] != "system":
            self.persistence()
            if explicit:
                # 触发插件事件
                args = {
                    'session_name': self.name,
                    'session': self
                }

                # 此事件不支持阻止默认行为
                _ = pkg.plugin.host.emit(plugin_models.SessionExplicitReset, **args)

                pkg.utils.context.get_database_manager().explicit_close_session(self.name, self.create_timestamp)

            if expired:
                pkg.utils.context.get_database_manager().set_session_expired(self.name, self.create_timestamp)
        self.prompt = self.get_default_prompt(use_prompt)
        self.create_timestamp = int(time.time())
        self.last_interact_timestamp = int(time.time())
        self.just_switched_to_exist_session = False

        # self.response_lock = threading.Lock()

        if schedule_new:
            self.schedule()

    # 将本session的数据库状态设置为on_going
    def set_ongoing(self):
        pkg.utils.context.get_database_manager().set_session_ongoing(self.name, self.create_timestamp)

    # 切换到上一个session
    def last_session(self):
        last_one = pkg.utils.context.get_database_manager().last_session(self.name, self.last_interact_timestamp)
        if last_one is None:
            return None
        else:
            self.persistence()

            self.create_timestamp = last_one['create_timestamp']
            self.last_interact_timestamp = last_one['last_interact_timestamp']
            self.prompt = json.loads(last_one['prompt'])

            self.just_switched_to_exist_session = True
            return self

    # 切换到下一个session
    def next_session(self):
        next_one = pkg.utils.context.get_database_manager().next_session(self.name, self.last_interact_timestamp)
        if next_one is None:
            return None
        else:
            self.persistence()

            self.create_timestamp = next_one['create_timestamp']
            self.last_interact_timestamp = next_one['last_interact_timestamp']
            self.prompt = json.loads(next_one['prompt'])

            self.just_switched_to_exist_session = True
            return self

    def list_history(self, capacity: int = 10, page: int = 0):
        return pkg.utils.context.get_database_manager().list_history(self.name, capacity, page)

    def draw_image(self, prompt: str):
        return pkg.utils.context.get_openai_manager().request_image(prompt)