104 lines
3.6KB

  1. from common.expired_dict import ExpiredDict
  2. from common.log import logger
  3. from config import conf
  4. class Session(object):
  5. def __init__(self, session_id, system_prompt=None):
  6. self.session_id = session_id
  7. self.messages = []
  8. if system_prompt is None:
  9. self.system_prompt = conf().get("character_desc", "")
  10. else:
  11. self.system_prompt = system_prompt
  12. # 重置会话
  13. def reset(self):
  14. system_item = {"role": "system", "content": self.system_prompt}
  15. self.messages = [system_item]
  16. def set_system_prompt(self, system_prompt):
  17. self.system_prompt = system_prompt
  18. self.reset()
  19. def add_query(self, query):
  20. user_item = {"role": "user", "content": query}
  21. self.messages.append(user_item)
  22. def add_reply(self, reply):
  23. assistant_item = {"role": "assistant", "content": reply}
  24. self.messages.append(assistant_item)
  25. def discard_exceeding(self, max_tokens=None, cur_tokens=None):
  26. raise NotImplementedError
  27. def calc_tokens(self):
  28. raise NotImplementedError
  29. class SessionManager(object):
  30. def __init__(self, sessioncls, **session_args):
  31. if conf().get("expires_in_seconds"):
  32. sessions = ExpiredDict(conf().get("expires_in_seconds"))
  33. else:
  34. sessions = dict()
  35. self.sessions = sessions
  36. self.sessioncls = sessioncls
  37. self.session_args = session_args
  38. def build_session(self, session_id, system_prompt=None):
  39. """
  40. 如果session_id不在sessions中,创建一个新的session并添加到sessions中
  41. 如果system_prompt不会空,会更新session的system_prompt并重置session
  42. """
  43. if session_id is None:
  44. return self.sessioncls(session_id, system_prompt, **self.session_args)
  45. if session_id not in self.sessions:
  46. self.sessions[session_id] = self.sessioncls(
  47. session_id, system_prompt, **self.session_args
  48. )
  49. elif system_prompt is not None: # 如果有新的system_prompt,更新并重置session
  50. self.sessions[session_id].set_system_prompt(system_prompt)
  51. session = self.sessions[session_id]
  52. return session
  53. def session_query(self, query, session_id):
  54. session = self.build_session(session_id)
  55. session.add_query(query)
  56. try:
  57. max_tokens = conf().get("conversation_max_tokens", 1000)
  58. total_tokens = session.discard_exceeding(max_tokens, None)
  59. logger.debug("prompt tokens used={}".format(total_tokens))
  60. except Exception as e:
  61. logger.debug(
  62. "Exception when counting tokens precisely for prompt: {}".format(str(e))
  63. )
  64. return session
  65. def session_reply(self, reply, session_id, total_tokens=None):
  66. session = self.build_session(session_id)
  67. session.add_reply(reply)
  68. try:
  69. max_tokens = conf().get("conversation_max_tokens", 1000)
  70. tokens_cnt = session.discard_exceeding(max_tokens, total_tokens)
  71. logger.debug(
  72. "raw total_tokens={}, savesession tokens={}".format(
  73. total_tokens, tokens_cnt
  74. )
  75. )
  76. except Exception as e:
  77. logger.debug(
  78. "Exception when counting tokens precisely for session: {}".format(
  79. str(e)
  80. )
  81. )
  82. return session
  83. def clear_session(self, session_id):
  84. if session_id in self.sessions:
  85. del self.sessions[session_id]
  86. def clear_all_session(self):
  87. self.sessions.clear()