You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

link_ai_bot.py 21KB

1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
1 vuosi sitten
11 kuukautta sitten
1 vuosi sitten
11 kuukautta sitten
10 kuukautta sitten
11 kuukautta sitten
1 vuosi sitten
7 kuukautta sitten
8 kuukautta sitten
1 vuosi sitten
8 kuukautta sitten
1 vuosi sitten
11 kuukautta sitten
7 kuukautta sitten
7 kuukautta sitten
7 kuukautta sitten
10 kuukautta sitten
10 kuukautta sitten
10 kuukautta sitten
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475
  1. # access LinkAI knowledge base platform
  2. # docs: https://link-ai.tech/platform/link-app/wechat
  3. import re
  4. import time
  5. import requests
  6. import config
  7. from bot.bot import Bot
  8. from bot.chatgpt.chat_gpt_session import ChatGPTSession
  9. from bot.session_manager import SessionManager
  10. from bridge.context import Context, ContextType
  11. from bridge.reply import Reply, ReplyType
  12. from common.log import logger
  13. from config import conf, pconf
  14. import threading
  15. from common import memory, utils
  16. import base64
  17. import os
  18. class LinkAIBot(Bot):
  19. # authentication failed
  20. AUTH_FAILED_CODE = 401
  21. NO_QUOTA_CODE = 406
  22. def __init__(self):
  23. super().__init__()
  24. self.sessions = LinkAISessionManager(LinkAISession, model=conf().get("model") or "gpt-3.5-turbo")
  25. self.args = {}
  26. def reply(self, query, context: Context = None) -> Reply:
  27. if context.type == ContextType.TEXT:
  28. return self._chat(query, context)
  29. elif context.type == ContextType.IMAGE_CREATE:
  30. if not conf().get("text_to_image"):
  31. logger.warn("[LinkAI] text_to_image is not enabled, ignore the IMAGE_CREATE request")
  32. return Reply(ReplyType.TEXT, "")
  33. ok, res = self.create_img(query, 0)
  34. if ok:
  35. reply = Reply(ReplyType.IMAGE_URL, res)
  36. else:
  37. reply = Reply(ReplyType.ERROR, res)
  38. return reply
  39. else:
  40. reply = Reply(ReplyType.ERROR, "Bot不支持处理{}类型的消息".format(context.type))
  41. return reply
  42. def _chat(self, query, context, retry_count=0) -> Reply:
  43. """
  44. 发起对话请求
  45. :param query: 请求提示词
  46. :param context: 对话上下文
  47. :param retry_count: 当前递归重试次数
  48. :return: 回复
  49. """
  50. if retry_count > 2:
  51. # exit from retry 2 times
  52. logger.warn("[LINKAI] failed after maximum number of retry times")
  53. return Reply(ReplyType.TEXT, "请再问我一次吧")
  54. try:
  55. # load config
  56. if context.get("generate_breaked_by"):
  57. logger.info(f"[LINKAI] won't set appcode because a plugin ({context['generate_breaked_by']}) affected the context")
  58. app_code = None
  59. else:
  60. plugin_app_code = self._find_group_mapping_code(context)
  61. app_code = context.kwargs.get("app_code") or plugin_app_code or conf().get("linkai_app_code")
  62. linkai_api_key = conf().get("linkai_api_key")
  63. session_id = context["session_id"]
  64. session_message = self.sessions.session_msg_query(query, session_id)
  65. logger.debug(f"[LinkAI] session={session_message}, session_id={session_id}")
  66. # image process
  67. img_cache = memory.USER_IMAGE_CACHE.get(session_id)
  68. if img_cache:
  69. messages = self._process_image_msg(app_code=app_code, session_id=session_id, query=query, img_cache=img_cache)
  70. if messages:
  71. session_message = messages
  72. model = conf().get("model")
  73. # remove system message
  74. if session_message[0].get("role") == "system":
  75. if app_code or model == "wenxin":
  76. session_message.pop(0)
  77. body = {
  78. "app_code": app_code,
  79. "messages": session_message,
  80. "model": model, # 对话模型的名称, 支持 gpt-3.5-turbo, gpt-3.5-turbo-16k, gpt-4, wenxin, xunfei
  81. "temperature": conf().get("temperature"),
  82. "top_p": conf().get("top_p", 1),
  83. "frequency_penalty": conf().get("frequency_penalty", 0.0), # [-2,2]之间,该值越大则更倾向于产生不同的内容
  84. "presence_penalty": conf().get("presence_penalty", 0.0), # [-2,2]之间,该值越大则更倾向于产生不同的内容
  85. "session_id": session_id,
  86. "sender_id": session_id,
  87. "channel_type": conf().get("channel_type", "wx")
  88. }
  89. try:
  90. from linkai import LinkAIClient
  91. client_id = LinkAIClient.fetch_client_id()
  92. if client_id:
  93. body["client_id"] = client_id
  94. # start: client info deliver
  95. if context.kwargs.get("msg"):
  96. body["session_id"] = context.kwargs.get("msg").from_user_id
  97. if context.kwargs.get("msg").is_group:
  98. body["is_group"] = True
  99. body["group_name"] = context.kwargs.get("msg").from_user_nickname
  100. body["sender_name"] = context.kwargs.get("msg").actual_user_nickname
  101. else:
  102. if body.get("channel_type") in ["wechatcom_app"]:
  103. body["sender_name"] = context.kwargs.get("msg").from_user_id
  104. else:
  105. body["sender_name"] = context.kwargs.get("msg").from_user_nickname
  106. except Exception as e:
  107. pass
  108. file_id = context.kwargs.get("file_id")
  109. if file_id:
  110. body["file_id"] = file_id
  111. logger.info(f"[LINKAI] query={query}, app_code={app_code}, model={body.get('model')}, file_id={file_id}")
  112. headers = {"Authorization": "Bearer " + linkai_api_key}
  113. # do http request
  114. base_url = conf().get("linkai_api_base", "https://api.link-ai.tech")
  115. res = requests.post(url=base_url + "/v1/chat/completions", json=body, headers=headers,
  116. timeout=conf().get("request_timeout", 180))
  117. if res.status_code == 200:
  118. # execute success
  119. response = res.json()
  120. reply_content = response["choices"][0]["message"]["content"]
  121. total_tokens = response["usage"]["total_tokens"]
  122. res_code = response.get('code')
  123. logger.info(f"[LINKAI] reply={reply_content}, total_tokens={total_tokens}, res_code={res_code}")
  124. if res_code == 429:
  125. logger.warn(f"[LINKAI] 用户访问超出限流配置,sender_id={body.get('sender_id')}")
  126. else:
  127. self.sessions.session_reply(reply_content, session_id, total_tokens, query=query)
  128. agent_suffix = self._fetch_agent_suffix(response)
  129. if agent_suffix:
  130. reply_content += agent_suffix
  131. if not agent_suffix:
  132. knowledge_suffix = self._fetch_knowledge_search_suffix(response)
  133. if knowledge_suffix:
  134. reply_content += knowledge_suffix
  135. # image process
  136. if response["choices"][0].get("img_urls"):
  137. thread = threading.Thread(target=self._send_image, args=(context.get("channel"), context, response["choices"][0].get("img_urls")))
  138. thread.start()
  139. if response["choices"][0].get("text_content"):
  140. reply_content = response["choices"][0].get("text_content")
  141. reply_content = self._process_url(reply_content)
  142. return Reply(ReplyType.TEXT, reply_content)
  143. else:
  144. response = res.json()
  145. error = response.get("error")
  146. logger.error(f"[LINKAI] chat failed, status_code={res.status_code}, "
  147. f"msg={error.get('message')}, type={error.get('type')}")
  148. if res.status_code >= 500:
  149. # server error, need retry
  150. time.sleep(2)
  151. logger.warn(f"[LINKAI] do retry, times={retry_count}")
  152. return self._chat(query, context, retry_count + 1)
  153. error_reply = "提问太快啦,请休息一下再问我吧"
  154. if res.status_code == 409:
  155. error_reply = "这个问题我还没有学会,请问我其它问题吧"
  156. return Reply(ReplyType.TEXT, error_reply)
  157. except Exception as e:
  158. logger.exception(e)
  159. # retry
  160. time.sleep(2)
  161. logger.warn(f"[LINKAI] do retry, times={retry_count}")
  162. return self._chat(query, context, retry_count + 1)
  163. def _process_image_msg(self, app_code: str, session_id: str, query:str, img_cache: dict):
  164. try:
  165. enable_image_input = False
  166. app_info = self._fetch_app_info(app_code)
  167. if not app_info:
  168. logger.debug(f"[LinkAI] not found app, can't process images, app_code={app_code}")
  169. return None
  170. plugins = app_info.get("data").get("plugins")
  171. for plugin in plugins:
  172. if plugin.get("input_type") and "IMAGE" in plugin.get("input_type"):
  173. enable_image_input = True
  174. if not enable_image_input:
  175. return
  176. msg = img_cache.get("msg")
  177. path = img_cache.get("path")
  178. msg.prepare()
  179. logger.info(f"[LinkAI] query with images, path={path}")
  180. messages = self._build_vision_msg(query, path)
  181. memory.USER_IMAGE_CACHE[session_id] = None
  182. return messages
  183. except Exception as e:
  184. logger.exception(e)
  185. def _find_group_mapping_code(self, context):
  186. try:
  187. if context.kwargs.get("isgroup"):
  188. group_name = context.kwargs.get("msg").from_user_nickname
  189. if config.plugin_config and config.plugin_config.get("linkai"):
  190. linkai_config = config.plugin_config.get("linkai")
  191. group_mapping = linkai_config.get("group_app_map")
  192. if group_mapping and group_name:
  193. return group_mapping.get(group_name)
  194. except Exception as e:
  195. logger.exception(e)
  196. return None
  197. def _build_vision_msg(self, query: str, path: str):
  198. try:
  199. suffix = utils.get_path_suffix(path)
  200. with open(path, "rb") as file:
  201. base64_str = base64.b64encode(file.read()).decode('utf-8')
  202. messages = [{
  203. "role": "user",
  204. "content": [
  205. {
  206. "type": "text",
  207. "text": query
  208. },
  209. {
  210. "type": "image_url",
  211. "image_url": {
  212. "url": f"data:image/{suffix};base64,{base64_str}"
  213. }
  214. }
  215. ]
  216. }]
  217. return messages
  218. except Exception as e:
  219. logger.exception(e)
  220. def reply_text(self, session: ChatGPTSession, app_code="", retry_count=0) -> dict:
  221. if retry_count >= 2:
  222. # exit from retry 2 times
  223. logger.warn("[LINKAI] failed after maximum number of retry times")
  224. return {
  225. "total_tokens": 0,
  226. "completion_tokens": 0,
  227. "content": "请再问我一次吧"
  228. }
  229. try:
  230. body = {
  231. "app_code": app_code,
  232. "messages": session.messages,
  233. "model": conf().get("model") or "gpt-3.5-turbo", # 对话模型的名称, 支持 gpt-3.5-turbo, gpt-3.5-turbo-16k, gpt-4, wenxin, xunfei
  234. "temperature": conf().get("temperature"),
  235. "top_p": conf().get("top_p", 1),
  236. "frequency_penalty": conf().get("frequency_penalty", 0.0), # [-2,2]之间,该值越大则更倾向于产生不同的内容
  237. "presence_penalty": conf().get("presence_penalty", 0.0), # [-2,2]之间,该值越大则更倾向于产生不同的内容
  238. }
  239. if self.args.get("max_tokens"):
  240. body["max_tokens"] = self.args.get("max_tokens")
  241. headers = {"Authorization": "Bearer " + conf().get("linkai_api_key")}
  242. # do http request
  243. base_url = conf().get("linkai_api_base", "https://api.link-ai.tech")
  244. res = requests.post(url=base_url + "/v1/chat/completions", json=body, headers=headers,
  245. timeout=conf().get("request_timeout", 180))
  246. if res.status_code == 200:
  247. # execute success
  248. response = res.json()
  249. reply_content = response["choices"][0]["message"]["content"]
  250. total_tokens = response["usage"]["total_tokens"]
  251. logger.info(f"[LINKAI] reply={reply_content}, total_tokens={total_tokens}")
  252. return {
  253. "total_tokens": total_tokens,
  254. "completion_tokens": response["usage"]["completion_tokens"],
  255. "content": reply_content,
  256. }
  257. else:
  258. response = res.json()
  259. error = response.get("error")
  260. logger.error(f"[LINKAI] chat failed, status_code={res.status_code}, "
  261. f"msg={error.get('message')}, type={error.get('type')}")
  262. if res.status_code >= 500:
  263. # server error, need retry
  264. time.sleep(2)
  265. logger.warn(f"[LINKAI] do retry, times={retry_count}")
  266. return self.reply_text(session, app_code, retry_count + 1)
  267. return {
  268. "total_tokens": 0,
  269. "completion_tokens": 0,
  270. "content": "提问太快啦,请休息一下再问我吧"
  271. }
  272. except Exception as e:
  273. logger.exception(e)
  274. # retry
  275. time.sleep(2)
  276. logger.warn(f"[LINKAI] do retry, times={retry_count}")
  277. return self.reply_text(session, app_code, retry_count + 1)
  278. def _fetch_app_info(self, app_code: str):
  279. headers = {"Authorization": "Bearer " + conf().get("linkai_api_key")}
  280. # do http request
  281. base_url = conf().get("linkai_api_base", "https://api.link-ai.tech")
  282. params = {"app_code": app_code}
  283. res = requests.get(url=base_url + "/v1/app/info", params=params, headers=headers, timeout=(5, 10))
  284. if res.status_code == 200:
  285. return res.json()
  286. else:
  287. logger.warning(f"[LinkAI] find app info exception, res={res}")
  288. def create_img(self, query, retry_count=0, api_key=None):
  289. try:
  290. logger.info("[LinkImage] image_query={}".format(query))
  291. headers = {
  292. "Content-Type": "application/json",
  293. "Authorization": f"Bearer {conf().get('linkai_api_key')}"
  294. }
  295. data = {
  296. "prompt": query,
  297. "n": 1,
  298. "model": conf().get("text_to_image") or "dall-e-2",
  299. "response_format": "url",
  300. "img_proxy": conf().get("image_proxy")
  301. }
  302. url = conf().get("linkai_api_base", "https://api.link-ai.tech") + "/v1/images/generations"
  303. res = requests.post(url, headers=headers, json=data, timeout=(5, 90))
  304. t2 = time.time()
  305. image_url = res.json()["data"][0]["url"]
  306. logger.info("[OPEN_AI] image_url={}".format(image_url))
  307. return True, image_url
  308. except Exception as e:
  309. logger.error(format(e))
  310. return False, "画图出现问题,请休息一下再问我吧"
  311. def _fetch_knowledge_search_suffix(self, response) -> str:
  312. try:
  313. if response.get("knowledge_base"):
  314. search_hit = response.get("knowledge_base").get("search_hit")
  315. first_similarity = response.get("knowledge_base").get("first_similarity")
  316. logger.info(f"[LINKAI] knowledge base, search_hit={search_hit}, first_similarity={first_similarity}")
  317. plugin_config = pconf("linkai")
  318. if plugin_config and plugin_config.get("knowledge_base") and plugin_config.get("knowledge_base").get("search_miss_text_enabled"):
  319. search_miss_similarity = plugin_config.get("knowledge_base").get("search_miss_similarity")
  320. search_miss_text = plugin_config.get("knowledge_base").get("search_miss_suffix")
  321. if not search_hit:
  322. return search_miss_text
  323. if search_miss_similarity and float(search_miss_similarity) > first_similarity:
  324. return search_miss_text
  325. except Exception as e:
  326. logger.exception(e)
  327. def _fetch_agent_suffix(self, response):
  328. try:
  329. plugin_list = []
  330. logger.debug(f"[LinkAgent] res={response}")
  331. if response.get("agent") and response.get("agent").get("chain") and response.get("agent").get("need_show_plugin"):
  332. chain = response.get("agent").get("chain")
  333. suffix = "\n\n- - - - - - - - - - - -"
  334. i = 0
  335. for turn in chain:
  336. plugin_name = turn.get('plugin_name')
  337. suffix += "\n"
  338. need_show_thought = response.get("agent").get("need_show_thought")
  339. if turn.get("thought") and plugin_name and need_show_thought:
  340. suffix += f"{turn.get('thought')}\n"
  341. if plugin_name:
  342. plugin_list.append(turn.get('plugin_name'))
  343. if turn.get('plugin_icon'):
  344. suffix += f"{turn.get('plugin_icon')} "
  345. suffix += f"{turn.get('plugin_name')}"
  346. if turn.get('plugin_input'):
  347. suffix += f":{turn.get('plugin_input')}"
  348. if i < len(chain) - 1:
  349. suffix += "\n"
  350. i += 1
  351. logger.info(f"[LinkAgent] use plugins: {plugin_list}")
  352. return suffix
  353. except Exception as e:
  354. logger.exception(e)
  355. def _process_url(self, text):
  356. try:
  357. url_pattern = re.compile(r'\[(.*?)\]\((http[s]?://.*?)\)')
  358. def replace_markdown_url(match):
  359. return f"{match.group(2)}"
  360. return url_pattern.sub(replace_markdown_url, text)
  361. except Exception as e:
  362. logger.error(e)
  363. def _send_image(self, channel, context, image_urls):
  364. if not image_urls:
  365. return
  366. max_send_num = conf().get("max_media_send_count")
  367. send_interval = conf().get("media_send_interval")
  368. file_type = (".pdf", ".doc", ".docx", ".csv", ".xls", ".xlsx", ".txt", ".rtf", ".ppt", ".pptx")
  369. try:
  370. i = 0
  371. for url in image_urls:
  372. if max_send_num and i >= max_send_num:
  373. continue
  374. i += 1
  375. if url.endswith(".mp4"):
  376. reply_type = ReplyType.VIDEO_URL
  377. elif url.endswith(file_type):
  378. reply_type = ReplyType.FILE
  379. url = _download_file(url)
  380. if not url:
  381. continue
  382. else:
  383. reply_type = ReplyType.IMAGE_URL
  384. reply = Reply(reply_type, url)
  385. channel.send(reply, context)
  386. if send_interval:
  387. time.sleep(send_interval)
  388. except Exception as e:
  389. logger.error(e)
  390. def _download_file(url: str):
  391. try:
  392. file_path = "tmp"
  393. if not os.path.exists(file_path):
  394. os.makedirs(file_path)
  395. file_name = url.split("/")[-1] # 获取文件名
  396. file_path = os.path.join(file_path, file_name)
  397. response = requests.get(url)
  398. with open(file_path, "wb") as f:
  399. f.write(response.content)
  400. return file_path
  401. except Exception as e:
  402. logger.warn(e)
  403. class LinkAISessionManager(SessionManager):
  404. def session_msg_query(self, query, session_id):
  405. session = self.build_session(session_id)
  406. messages = session.messages + [{"role": "user", "content": query}]
  407. return messages
  408. def session_reply(self, reply, session_id, total_tokens=None, query=None):
  409. session = self.build_session(session_id)
  410. if query:
  411. session.add_query(query)
  412. session.add_reply(reply)
  413. try:
  414. max_tokens = conf().get("conversation_max_tokens", 2500)
  415. tokens_cnt = session.discard_exceeding(max_tokens, total_tokens)
  416. logger.debug(f"[LinkAI] chat history, before tokens={total_tokens}, now tokens={tokens_cnt}")
  417. except Exception as e:
  418. logger.warning("Exception when counting tokens precisely for session: {}".format(str(e)))
  419. return session
  420. class LinkAISession(ChatGPTSession):
  421. def calc_tokens(self):
  422. if not self.messages:
  423. return 0
  424. return len(str(self.messages))
  425. def discard_exceeding(self, max_tokens, cur_tokens=None):
  426. cur_tokens = self.calc_tokens()
  427. if cur_tokens > max_tokens:
  428. for i in range(0, len(self.messages)):
  429. if i > 0 and self.messages[i].get("role") == "assistant" and self.messages[i - 1].get("role") == "user":
  430. self.messages.pop(i)
  431. self.messages.pop(i - 1)
  432. return self.calc_tokens()
  433. return cur_tokens