chat.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. from openai import OpenAI
  2. client = OpenAI(
  3. api_key = "sk-g4Vp3q4nSGPGNalQ04sbkiSeJX95PAZeYMv8ENB7mVq2Iu8j",
  4. base_url = "https://api.moonshot.cn/v1",
  5. )
  6. # 默认人设上下文
  7. history = [
  8. {"role": "system", "content": "你是 Coo,由汇智提供的人工智能助手,你更擅长中文和英文的对话。你会为用户提供安全,有帮助,准确的回答。同时,你会拒绝一切涉及恐怖主义,种族歧视,黄色暴力等问题的回答。使用检索到的上下文片段来回答问题。如果你不知道答案,就说你不知道。"}
  9. ]
  10. # c = Chat()
  11. # print(c.chat('东京奥运会中国得了多少金牌?'))
  12. # print(c.chat('巴黎呢?'))
  13. class Chat:
  14. history = [] # 对话历史上下文
  15. save_first_len = 1 # 保留起始的几条上下文,小于等于 0 设置无效
  16. save_last_len = 10 # 保留末尾的几条上下文,小于等于 0 为不限制
  17. def __init__(self, history=history, save_first_len=1, save_last_len=10) -> None:
  18. self.history = history
  19. self.save_first_len = save_first_len if save_first_len >= 0 else 0
  20. self.save_last_len = save_last_len if save_last_len >= 0 else 0
  21. def chat(self, query):
  22. self.history.append({
  23. "role": "user",
  24. "content": query
  25. })
  26. completion = client.chat.completions.create(
  27. model="moonshot-v1-8k",
  28. messages=self.history,
  29. temperature=0.3,
  30. )
  31. result = completion.choices[0].message.content
  32. self.history.append({
  33. "role": "assistant",
  34. "content": result
  35. })
  36. # 随着对话的进行,模型每次需要传入的 token 都会线性增加,只保留特定几轮对话
  37. if len(self.save_first_len) > 0 and len(self.save_last_len) > 0 and self.save_first_len + self.save_last_len < len(self.history):
  38. self.history = self.history[0:self.save_first_len] + self.history[len(self.history) - self.save_last_len:]
  39. return result