Spaces:
Sleeping
Sleeping
| from common.expired_dict import ExpiredDict | |
| from common.log import logger | |
| from config import conf | |
| class Session(object): | |
| def __init__(self, session_id, system_prompt=None): | |
| self.session_id = session_id | |
| self.messages = [] | |
| if system_prompt is None: | |
| self.system_prompt = conf().get("character_desc", "") | |
| else: | |
| self.system_prompt = system_prompt | |
| # 重置会话 | |
| def reset(self): | |
| system_item = {'role': 'system', 'content': self.system_prompt} | |
| self.messages = [system_item] | |
| def set_system_prompt(self, system_prompt): | |
| self.system_prompt = system_prompt | |
| self.reset() | |
| def add_query(self, query): | |
| user_item = {'role': 'user', 'content': query} | |
| self.messages.append(user_item) | |
| def add_reply(self, reply): | |
| assistant_item = {'role': 'assistant', 'content': reply} | |
| self.messages.append(assistant_item) | |
| def discard_exceeding(self, max_tokens=None, cur_tokens=None): | |
| raise NotImplementedError | |
| class SessionManager(object): | |
| def __init__(self, sessioncls, **session_args): | |
| if conf().get('expires_in_seconds'): | |
| sessions = ExpiredDict(conf().get('expires_in_seconds')) | |
| else: | |
| sessions = dict() | |
| self.sessions = sessions | |
| self.sessioncls = sessioncls | |
| self.session_args = session_args | |
| def build_session(self, session_id, system_prompt=None): | |
| ''' | |
| 如果session_id不在sessions中，创建一个新的session并添加到sessions中 | |
| 如果system_prompt不会空，会更新session的system_prompt并重置session | |
| ''' | |
| if session_id not in self.sessions: | |
| self.sessions[session_id] = self.sessioncls(session_id, system_prompt, **self.session_args) | |
| elif system_prompt is not None: # 如果有新的system_prompt，更新并重置session | |
| self.sessions[session_id].set_system_prompt(system_prompt) | |
| session = self.sessions[session_id] | |
| return session | |
| def session_query(self, query, session_id): | |
| session = self.build_session(session_id) | |
| session.add_query(query) | |
| try: | |
| max_tokens = conf().get("conversation_max_tokens", 1000) | |
| total_tokens = session.discard_exceeding(max_tokens, None) | |
| logger.debug("prompt tokens used={}".format(total_tokens)) | |
| except Exception as e: | |
| logger.debug("Exception when counting tokens precisely for prompt: {}".format(str(e))) | |
| return session | |
| def session_reply(self, reply, session_id, total_tokens = None): | |
| session = self.build_session(session_id) | |
| session.add_reply(reply) | |
| try: | |
| max_tokens = conf().get("conversation_max_tokens", 1000) | |
| tokens_cnt = session.discard_exceeding(max_tokens, total_tokens) | |
| logger.debug("raw total_tokens={}, savesession tokens={}".format(total_tokens, tokens_cnt)) | |
| except Exception as e: | |
| logger.debug("Exception when counting tokens precisely for session: {}".format(str(e))) | |
| return session | |
| def clear_session(self, session_id): | |
| if session_id in self.sessions: | |
| del(self.sessions[session_id]) | |
| def clear_all_session(self): | |
| self.sessions.clear() | |