formatting code

This commit is contained in:
lanvent
2023-04-17 01:00:08 +08:00
parent 3b8972ce1f
commit 8f72e8c3e6
92 changed files with 1850 additions and 1188 deletions

View File

@@ -2,6 +2,7 @@ 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
@@ -13,7 +14,7 @@ class Session(object):
# 重置会话
def reset(self):
system_item = {'role': 'system', 'content': self.system_prompt}
system_item = {"role": "system", "content": self.system_prompt}
self.messages = [system_item]
def set_system_prompt(self, system_prompt):
@@ -21,13 +22,13 @@ class Session(object):
self.reset()
def add_query(self, query):
user_item = {'role': 'user', 'content': query}
user_item = {"role": "user", "content": query}
self.messages.append(user_item)
def add_reply(self, reply):
assistant_item = {'role': 'assistant', 'content': reply}
assistant_item = {"role": "assistant", "content": reply}
self.messages.append(assistant_item)
def discard_exceeding(self, max_tokens=None, cur_tokens=None):
raise NotImplementedError
@@ -37,8 +38,8 @@ class Session(object):
class SessionManager(object):
def __init__(self, sessioncls, **session_args):
if conf().get('expires_in_seconds'):
sessions = ExpiredDict(conf().get('expires_in_seconds'))
if conf().get("expires_in_seconds"):
sessions = ExpiredDict(conf().get("expires_in_seconds"))
else:
sessions = dict()
self.sessions = sessions
@@ -46,20 +47,22 @@ class SessionManager(object):
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
'''
"""
如果session_id不在sessions中创建一个新的session并添加到sessions中
如果system_prompt不会空会更新session的system_prompt并重置session
"""
if session_id is None:
return self.sessioncls(session_id, system_prompt, **self.session_args)
if session_id not in self.sessions:
self.sessions[session_id] = self.sessioncls(session_id, system_prompt, **self.session_args)
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)
@@ -68,23 +71,33 @@ class SessionManager(object):
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)))
logger.debug(
"Exception when counting tokens precisely for prompt: {}".format(str(e))
)
return session
def session_reply(self, reply, session_id, total_tokens = None):
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))
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)))
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])
del self.sessions[session_id]
def clear_all_session(self):
self.sessions.clear()