Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.

baidu_wenxin.py 5.0KB

před 5 měsíci
před 5 měsíci
před 5 měsíci
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  1. # encoding:utf-8
  2. import requests
  3. import json
  4. from common import const
  5. from bot.bot import Bot
  6. from bot.session_manager import SessionManager
  7. from bridge.context import ContextType
  8. from bridge.reply import Reply, ReplyType
  9. from common.log import logger
  10. from config import conf
  11. from bot.baidu.baidu_wenxin_session import BaiduWenxinSession
  12. BAIDU_API_KEY = conf().get("baidu_wenxin_api_key")
  13. BAIDU_SECRET_KEY = conf().get("baidu_wenxin_secret_key")
  14. class BaiduWenxinBot(Bot):
  15. def __init__(self):
  16. super().__init__()
  17. wenxin_model = conf().get("baidu_wenxin_model")
  18. if wenxin_model is not None:
  19. wenxin_model = conf().get("baidu_wenxin_model") or "eb-instant"
  20. else:
  21. if conf().get("model") and conf().get("model") == const.WEN_XIN:
  22. wenxin_model = "completions"
  23. elif conf().get("model") and conf().get("model") == const.WEN_XIN_4:
  24. wenxin_model = "completions_pro"
  25. self.sessions = SessionManager(BaiduWenxinSession, model=wenxin_model)
  26. def reply(self, query, context=None):
  27. # acquire reply content
  28. if context and context.type:
  29. if context.type == ContextType.TEXT:
  30. logger.info("[BAIDU] query={}".format(query))
  31. session_id = context["session_id"]
  32. reply = None
  33. if query == "#清除记忆":
  34. self.sessions.clear_session(session_id)
  35. reply = Reply(ReplyType.INFO, "记忆已清除")
  36. elif query == "#清除所有":
  37. self.sessions.clear_all_session()
  38. reply = Reply(ReplyType.INFO, "所有人记忆已清除")
  39. else:
  40. session = self.sessions.session_query(query, session_id)
  41. result = self.reply_text(session)
  42. total_tokens, completion_tokens, reply_content = (
  43. result["total_tokens"],
  44. result["completion_tokens"],
  45. result["content"],
  46. )
  47. logger.debug(
  48. "[BAIDU] new_query={}, session_id={}, reply_cont={}, completion_tokens={}".format(session.messages, session_id, reply_content, completion_tokens)
  49. )
  50. if total_tokens == 0:
  51. reply = Reply(ReplyType.ERROR, reply_content)
  52. else:
  53. self.sessions.session_reply(reply_content, session_id, total_tokens)
  54. reply = Reply(ReplyType.TEXT, reply_content)
  55. return reply
  56. elif context.type == ContextType.IMAGE_CREATE:
  57. ok, retstring = self.create_img(query, 0)
  58. reply = None
  59. if ok:
  60. reply = Reply(ReplyType.IMAGE_URL, retstring)
  61. else:
  62. reply = Reply(ReplyType.ERROR, retstring)
  63. return reply
  64. def reply_text(self, session: BaiduWenxinSession, retry_count=0):
  65. try:
  66. logger.info("[BAIDU] model={}".format(session.model))
  67. access_token = self.get_access_token()
  68. if access_token == 'None':
  69. logger.warn("[BAIDU] access token 获取失败")
  70. return {
  71. "total_tokens": 0,
  72. "completion_tokens": 0,
  73. "content": 0,
  74. }
  75. url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/" + session.model + "?access_token=" + access_token
  76. headers = {
  77. 'Content-Type': 'application/json'
  78. }
  79. payload = {'messages': session.messages}
  80. response = requests.request("POST", url, headers=headers, data=json.dumps(payload))
  81. response_text = json.loads(response.text)
  82. logger.info(f"[BAIDU] response text={response_text}")
  83. res_content = response_text["result"]
  84. total_tokens = response_text["usage"]["total_tokens"]
  85. completion_tokens = response_text["usage"]["completion_tokens"]
  86. logger.info("[BAIDU] reply={}".format(res_content))
  87. return {
  88. "total_tokens": total_tokens,
  89. "completion_tokens": completion_tokens,
  90. "content": res_content,
  91. }
  92. except Exception as e:
  93. need_retry = retry_count < 2
  94. logger.warn("[BAIDU] Exception: {}".format(e))
  95. need_retry = False
  96. self.sessions.clear_session(session.session_id)
  97. result = {"total_tokens": 0, "completion_tokens": 0, "content": "出错了: {}".format(e)}
  98. return result
  99. def get_access_token(self):
  100. """
  101. 使用 AK,SK 生成鉴权签名(Access Token)
  102. :return: access_token,或是None(如果错误)
  103. """
  104. url = "https://aip.baidubce.com/oauth/2.0/token"
  105. params = {"grant_type": "client_credentials", "client_id": BAIDU_API_KEY, "client_secret": BAIDU_SECRET_KEY}
  106. return str(requests.post(url, params=params).json().get("access_token"))