Spaces:
Sleeping
Sleeping
File size: 3,326 Bytes
0dc9888 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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()
|