178 lines
6.0KB

  1. # encoding:utf-8
  2. from bot.bot import Bot
  3. from config import conf
  4. from common.log import logger
  5. import openai
  6. from datetime import date
  7. user_session = dict()
  8. # OpenAI对话模型API (可用)
  9. class OpenAIBot(Bot):
  10. def __init__(self):
  11. openai.api_key = conf().get('open_ai_api_key')
  12. def reply(self, query, context=None):
  13. # acquire reply content
  14. if not context or not context.get('type') or context.get('type') == 'TEXT':
  15. logger.info("[OPEN_AI] query={}".format(query))
  16. from_user_id = context['from_user_id']
  17. if query == '#清除记忆':
  18. Session.clear_session(from_user_id)
  19. return '记忆已清除'
  20. new_query = Session.build_session_query(query, from_user_id)
  21. logger.debug("[OPEN_AI] session query={}".format(new_query))
  22. reply_content = self.reply_text(new_query, from_user_id)
  23. logger.debug("[OPEN_AI] new_query={}, user={}".format(new_query, from_user_id))
  24. Session.save_session(query, reply_content, from_user_id)
  25. return reply_content
  26. elif context.get('type', None) == 'IMAGE_CREATE':
  27. return self.create_img(query)
  28. def reply_text(self, query, user_id):
  29. try:
  30. response = openai.Completion.create(
  31. model="text-davinci-003", # 对话模型的名称
  32. prompt=query,
  33. temperature=0.9, # 值在[0,1]之间,越大表示回复越具有不确定性
  34. max_tokens=1200, # 回复最大的字符数
  35. top_p=1,
  36. frequency_penalty=0.0, # [-2,2]之间,该值越大则更倾向于产生不同的内容
  37. presence_penalty=0.0, # [-2,2]之间,该值越大则更倾向于产生不同的内容
  38. stop=["#"]
  39. )
  40. res_content = response.choices[0]["text"].strip().rstrip("<|im_end|>")
  41. except Exception as e:
  42. logger.exception(e)
  43. Session.clear_session(user_id)
  44. return None
  45. logger.info("[OPEN_AI] reply={}".format(res_content))
  46. return res_content
  47. def create_img(self, query):
  48. try:
  49. logger.info("[OPEN_AI] image_query={}".format(query))
  50. response = openai.Image.create(
  51. prompt=query, #图片描述
  52. n=1, #每次生成图片的数量
  53. size="256x256" #图片大小,可选有 256x256, 512x512, 1024x1024
  54. )
  55. image_url = response['data'][0]['url']
  56. logger.info("[OPEN_AI] image_url={}".format(image_url))
  57. except Exception as e:
  58. logger.exception(e)
  59. return None
  60. return image_url
  61. def edit_img(self, query, src_img):
  62. try:
  63. response = openai.Image.create_edit(
  64. image=open(src_img, 'rb'),
  65. mask=open('cat-mask.png', 'rb'),
  66. prompt=query,
  67. n=1,
  68. size='512x512'
  69. )
  70. image_url = response['data'][0]['url']
  71. logger.info("[OPEN_AI] image_url={}".format(image_url))
  72. except Exception as e:
  73. logger.exception(e)
  74. return None
  75. return image_url
  76. def migration_img(self, query, src_img):
  77. try:
  78. response = openai.Image.create_variation(
  79. image=open(src_img, 'rb'),
  80. n=1,
  81. size="512x512"
  82. )
  83. image_url = response['data'][0]['url']
  84. logger.info("[OPEN_AI] image_url={}".format(image_url))
  85. except Exception as e:
  86. logger.exception(e)
  87. return None
  88. return image_url
  89. def append_question_mark(self, query):
  90. end_symbols = [".", "。", "?", "?", "!", "!"]
  91. for symbol in end_symbols:
  92. if query.endswith(symbol):
  93. return query
  94. return query + "?"
  95. class Session(object):
  96. @staticmethod
  97. def build_session_query(query, user_id):
  98. '''
  99. build query with conversation history
  100. e.g. Q: xxx
  101. A: xxx
  102. Q: xxx
  103. :param query: query content
  104. :param user_id: from user id
  105. :return: query content with conversaction
  106. '''
  107. prompt = conf().get("character_desc", "")
  108. if prompt:
  109. prompt += "\n\n"
  110. session = user_session.get(user_id, None)
  111. if session:
  112. for conversation in session:
  113. prompt += "Q: " + conversation["question"] + "\n\n\nA: " + conversation["answer"] + "<|im_end|>\n"
  114. prompt += "Q: " + query + "\nA: "
  115. return prompt
  116. else:
  117. return prompt + "Q: " + query + "\nA: "
  118. @staticmethod
  119. def save_session(query, answer, user_id):
  120. max_tokens = conf().get("conversation_max_tokens")
  121. if not max_tokens:
  122. # default 3000
  123. max_tokens = 1000
  124. conversation = dict()
  125. conversation["question"] = query
  126. conversation["answer"] = answer
  127. session = user_session.get(user_id)
  128. logger.debug(conversation)
  129. logger.debug(session)
  130. if session:
  131. # append conversation
  132. session.append(conversation)
  133. else:
  134. # create session
  135. queue = list()
  136. queue.append(conversation)
  137. user_session[user_id] = queue
  138. # discard exceed limit conversation
  139. Session.discard_exceed_conversation(user_session[user_id], max_tokens)
  140. @staticmethod
  141. def discard_exceed_conversation(session, max_tokens):
  142. count = 0
  143. count_list = list()
  144. for i in range(len(session)-1, -1, -1):
  145. # count tokens of conversation list
  146. history_conv = session[i]
  147. count += len(history_conv["question"]) + len(history_conv["answer"])
  148. count_list.append(count)
  149. for c in count_list:
  150. if c > max_tokens:
  151. # pop first conversation
  152. session.pop(0)
  153. @staticmethod
  154. def clear_session(user_id):
  155. user_session[user_id] = []