mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-06-02 00:57:41 +08:00
formatting code
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user