73 lines
2.7KB

  1. from bot.bot import Bot
  2. from bridge.context import ContextType
  3. from bridge.reply import Reply, ReplyType
  4. from common.log import logger
  5. from bridge.context import Context
  6. from bot.chatgpt.chat_gpt_session import ChatGPTSession
  7. from bot.session_manager import SessionManager
  8. from config import conf
  9. import requests
  10. import time
  11. class LinkAIBot(Bot):
  12. # authentication failed
  13. AUTH_FAILED_CODE = 401
  14. def __init__(self):
  15. self.base_url = "https://api.link-ai.chat/v1"
  16. self.sessions = SessionManager(ChatGPTSession, model=conf().get("model") or "gpt-3.5-turbo")
  17. def reply(self, query, context: Context = None) -> Reply:
  18. return self._chat(query, context)
  19. def _chat(self, query, context, retry_count=0):
  20. if retry_count >= 2:
  21. # exit from retry 2 times
  22. logger.warn("[LINKAI] failed after maximum number of retry times")
  23. return Reply(ReplyType.ERROR, "请再问我一次吧")
  24. try:
  25. session_id = context["session_id"]
  26. session = self.sessions.session_query(query, session_id)
  27. # remove system message
  28. if session.messages[0].get("role") == "system":
  29. session.messages.pop(0)
  30. # load config
  31. app_code = conf().get("app_code")
  32. app_market_api_key = conf().get("app_market_api_key")
  33. logger.info(f"[LINKAI] query={query}, app_code={app_code}")
  34. body = {
  35. "appCode": app_code,
  36. "messages": session.messages
  37. }
  38. headers = {"Authorization": "Bearer " + app_market_api_key}
  39. # do http request
  40. res = requests.post(url=self.base_url + "/chat/completion", json=body, headers=headers).json()
  41. if not res or not res["success"]:
  42. if res.get("code") == self.AUTH_FAILED_CODE:
  43. logger.exception(f"[LINKAI] please check your app_market_api_key, res={res}")
  44. return Reply(ReplyType.ERROR, "请再问我一次吧")
  45. else:
  46. # retry
  47. time.sleep(2)
  48. logger.warn(f"[LINKAI] do retry, times={retry_count}")
  49. return self._chat(query, context, retry_count + 1)
  50. # execute success
  51. reply_content = res["data"]["content"]
  52. logger.info(f"[LINKAI] reply={reply_content}")
  53. self.sessions.session_reply(reply_content, session_id)
  54. return Reply(ReplyType.TEXT, reply_content)
  55. except Exception as e:
  56. logger.exception(e)
  57. # retry
  58. time.sleep(2)
  59. logger.warn(f"[LINKAI] do retry, times={retry_count}")
  60. return self._chat(query, context, retry_count + 1)