chat.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  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. class Chat:
  11. history = [] # 对话历史上下文
  12. save_first_len = 1 # 保留起始的几条上下文,小于等于 0 设置无效
  13. save_last_len = 10 # 保留末尾的几条上下文,小于等于 0 为不限制
  14. def __init__(self, history=history, save_first_len=1, save_last_len=10) -> None:
  15. self.history = history
  16. self.save_first_len = save_first_len if save_first_len >= 0 else 0
  17. self.save_last_len = save_last_len if save_last_len >= 0 else 0
  18. def chat(self, query):
  19. self.history.append({
  20. "role": "user",
  21. "content": query
  22. })
  23. completion = client.chat.completions.create(
  24. model="moonshot-v1-8k",
  25. messages=self.history,
  26. temperature=0.3,
  27. )
  28. result = completion.choices[0].message.content
  29. self.history.append({
  30. "role": "assistant",
  31. "content": result
  32. })
  33. # 随着对话的进行,模型每次需要传入的 token 都会线性增加,只保留特定几轮对话
  34. 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):
  35. self.history = self.history[0:self.save_first_len] + self.history[len(self.history) - self.save_last_len:]
  36. return result