open_ai_session.py 2.9KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. from bot.session_manager import Session
  2. from common.log import logger
  3. class OpenAISession(Session):
  4. def __init__(self, session_id, system_prompt=None, model="text-davinci-003"):
  5. super().__init__(session_id, system_prompt)
  6. self.model = model
  7. self.reset()
  8. def __str__(self):
  9. # 构造对话模型的输入
  10. """
  11. e.g. Q: xxx
  12. A: xxx
  13. Q: xxx
  14. """
  15. prompt = ""
  16. for item in self.messages:
  17. if item["role"] == "system":
  18. prompt += item["content"] + "<|endoftext|>\n\n\n"
  19. elif item["role"] == "user":
  20. prompt += "Q: " + item["content"] + "\n"
  21. elif item["role"] == "assistant":
  22. prompt += "\n\nA: " + item["content"] + "<|endoftext|>\n"
  23. if len(self.messages) > 0 and self.messages[-1]["role"] == "user":
  24. prompt += "A: "
  25. return prompt
  26. def discard_exceeding(self, max_tokens, cur_tokens=None):
  27. precise = True
  28. try:
  29. cur_tokens = self.calc_tokens()
  30. except Exception as e:
  31. precise = False
  32. if cur_tokens is None:
  33. raise e
  34. logger.debug(
  35. "Exception when counting tokens precisely for query: {}".format(e)
  36. )
  37. while cur_tokens > max_tokens:
  38. if len(self.messages) > 1:
  39. self.messages.pop(0)
  40. elif len(self.messages) == 1 and self.messages[0]["role"] == "assistant":
  41. self.messages.pop(0)
  42. if precise:
  43. cur_tokens = self.calc_tokens()
  44. else:
  45. cur_tokens = len(str(self))
  46. break
  47. elif len(self.messages) == 1 and self.messages[0]["role"] == "user":
  48. logger.warn(
  49. "user question exceed max_tokens. total_tokens={}".format(
  50. cur_tokens
  51. )
  52. )
  53. break
  54. else:
  55. logger.debug(
  56. "max_tokens={}, total_tokens={}, len(conversation)={}".format(
  57. max_tokens, cur_tokens, len(self.messages)
  58. )
  59. )
  60. break
  61. if precise:
  62. cur_tokens = self.calc_tokens()
  63. else:
  64. cur_tokens = len(str(self))
  65. return cur_tokens
  66. def calc_tokens(self):
  67. return num_tokens_from_string(str(self), self.model)
  68. # refer to https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
  69. def num_tokens_from_string(string: str, model: str) -> int:
  70. """Returns the number of tokens in a text string."""
  71. import tiktoken
  72. encoding = tiktoken.encoding_for_model(model)
  73. num_tokens = len(encoding.encode(string, disallowed_special=()))
  74. return num_tokens