J1147
/
ai-on-wechat
Sledovat 1
Oblíbit 0
Rozštěpit 0
Zdrojový kód Úkoly 0 Požadavky na natažení 0 Vydání 0 Wiki Aktivita
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ů.
919 Revize
3 Větve
10MB
Strom: be0bb591e7
Větve Značky
${ item.name }
Vytvořit větev ${ searchTerm }
z „be0bb591e7“
${ noResults }
ai-on-wechat/bot/linkai/link_ai_bot.py

405 lines
18KB
Surový Blame Historie 

  1. # access LinkAI knowledge base platform

  2. # docs: https://link-ai.tech/platform/link-app/wechat

  3. 

  4. import time

  5. 

  6. import requests

  7. 

  8. import config

  9. from bot.bot import Bot

  10. from bot.chatgpt.chat_gpt_session import ChatGPTSession

  11. from bot.session_manager import SessionManager

  12. from bridge.context import Context, ContextType

  13. from bridge.reply import Reply, ReplyType

  14. from common.log import logger

  15. from config import conf, pconf

  16. import threading

  17. from common import memory, utils

  18. import base64

  19. 

  20. 

  21. class LinkAIBot(Bot):

  22.     # authentication failed

  23.     AUTH_FAILED_CODE = 401

  24.     NO_QUOTA_CODE = 406

  25. 

  26.     def __init__(self):

  27.         super().__init__()

  28.         self.sessions = LinkAISessionManager(LinkAISession, model=conf().get("model") or "gpt-3.5-turbo")

  29.         self.args = {}

  30. 

  31.     def reply(self, query, context: Context = None) -> Reply:

  32.         if context.type == ContextType.TEXT:

  33.             return self._chat(query, context)

  34.         elif context.type == ContextType.IMAGE_CREATE:

  35.             if not conf().get("text_to_image"):

  36.                 logger.warn("[LinkAI] text_to_image is not enabled, ignore the IMAGE_CREATE request")

  37.                 return Reply(ReplyType.TEXT, "")

  38.             ok, res = self.create_img(query, 0)

  39.             if ok:

  40.                 reply = Reply(ReplyType.IMAGE_URL, res)

  41.             else:

  42.                 reply = Reply(ReplyType.ERROR, res)

  43.             return reply

  44.         else:

  45.             reply = Reply(ReplyType.ERROR, "Bot不支持处理{}类型的消息".format(context.type))

  46.             return reply

  47. 

  48.     def _chat(self, query, context, retry_count=0) -> Reply:

  49.         """

  50.         发起对话请求

  51.         :param query: 请求提示词

  52.         :param context: 对话上下文

  53.         :param retry_count: 当前递归重试次数

  54.         :return: 回复

  55.         """

  56.         if retry_count > 2:

  57.             # exit from retry 2 times

  58.             logger.warn("[LINKAI] failed after maximum number of retry times")

  59.             return Reply(ReplyType.TEXT, "请再问我一次吧")

  60. 

  61.         try:

  62.             # load config

  63.             if context.get("generate_breaked_by"):

  64.                 logger.info(f"[LINKAI] won't set appcode because a plugin ({context['generate_breaked_by']}) affected the context")

  65.                 app_code = None

  66.             else:

  67.                 plugin_app_code = self._find_group_mapping_code(context)

  68.                 app_code = context.kwargs.get("app_code") or plugin_app_code or conf().get("linkai_app_code")

  69.             linkai_api_key = conf().get("linkai_api_key")

  70. 

  71.             session_id = context["session_id"]

  72.             session_message = self.sessions.session_msg_query(query, session_id)

  73.             logger.debug(f"[LinkAI] session={session_message}, session_id={session_id}")

  74. 

  75.             # image process

  76.             img_cache = memory.USER_IMAGE_CACHE.get(session_id)

  77.             if img_cache:

  78.                 messages = self._process_image_msg(app_code=app_code, session_id=session_id, query=query, img_cache=img_cache)

  79.                 if messages:

  80.                     session_message = messages

  81. 

  82.             model = conf().get("model")

  83.             # remove system message

  84.             if session_message[0].get("role") == "system":

  85.                 if app_code or model == "wenxin":

  86.                     session_message.pop(0)

  87. 

  88.             body = {

  89.                 "app_code": app_code,

  90.                 "messages": session_message,

  91.                 "model": model,     # 对话模型的名称, 支持 gpt-3.5-turbo, gpt-3.5-turbo-16k, gpt-4, wenxin, xunfei

  92.                 "temperature": conf().get("temperature"),

  93.                 "top_p": conf().get("top_p", 1),

  94.                 "frequency_penalty": conf().get("frequency_penalty", 0.0),  # [-2,2]之间，该值越大则更倾向于产生不同的内容

  95.                 "presence_penalty": conf().get("presence_penalty", 0.0),  # [-2,2]之间，该值越大则更倾向于产生不同的内容

  96.             }

  97.             file_id = context.kwargs.get("file_id")

  98.             if file_id:

  99.                 body["file_id"] = file_id

  100.             logger.info(f"[LINKAI] query={query}, app_code={app_code}, model={body.get('model')}, file_id={file_id}")

  101.             headers = {"Authorization": "Bearer " + linkai_api_key}

  102. 

  103.             # do http request

  104.             base_url = conf().get("linkai_api_base", "https://api.link-ai.chat")

  105.             res = requests.post(url=base_url + "/v1/chat/completions", json=body, headers=headers,

  106.                                 timeout=conf().get("request_timeout", 180))

  107.             if res.status_code == 200:

  108.                 # execute success

  109.                 response = res.json()

  110.                 reply_content = response["choices"][0]["message"]["content"]

  111.                 total_tokens = response["usage"]["total_tokens"]

  112.                 logger.info(f"[LINKAI] reply={reply_content}, total_tokens={total_tokens}")

  113.                 self.sessions.session_reply(reply_content, session_id, total_tokens, query=query)

  114.     

  115.                 agent_suffix = self._fetch_agent_suffix(response)

  116.                 if agent_suffix:

  117.                     reply_content += agent_suffix

  118.                 if not agent_suffix:

  119.                     knowledge_suffix = self._fetch_knowledge_search_suffix(response)

  120.                     if knowledge_suffix:

  121.                         reply_content += knowledge_suffix

  122.                 # image process

  123.                 if response["choices"][0].get("img_urls"):

  124.                     thread = threading.Thread(target=self._send_image, args=(context.get("channel"), context, response["choices"][0].get("img_urls")))

  125.                     thread.start()

  126.                     if response["choices"][0].get("text_content"):

  127.                         reply_content = response["choices"][0].get("text_content")

  128.                 return Reply(ReplyType.TEXT, reply_content)

  129. 

  130.             else:

  131.                 response = res.json()

  132.                 error = response.get("error")

  133.                 logger.error(f"[LINKAI] chat failed, status_code={res.status_code}, "

  134.                              f"msg={error.get('message')}, type={error.get('type')}")

  135. 

  136.                 if res.status_code >= 500:

  137.                     # server error, need retry

  138.                     time.sleep(2)

  139.                     logger.warn(f"[LINKAI] do retry, times={retry_count}")

  140.                     return self._chat(query, context, retry_count + 1)

  141. 

  142.                 return Reply(ReplyType.TEXT, "提问太快啦，请休息一下再问我吧")

  143. 

  144.         except Exception as e:

  145.             logger.exception(e)

  146.             # retry

  147.             time.sleep(2)

  148.             logger.warn(f"[LINKAI] do retry, times={retry_count}")

  149.             return self._chat(query, context, retry_count + 1)

  150. 

  151.     def _process_image_msg(self, app_code: str, session_id: str, query:str, img_cache: dict):

  152.         try:

  153.             enable_image_input = False

  154.             app_info = self._fetch_app_info(app_code)

  155.             if not app_info:

  156.                 logger.debug(f"[LinkAI] not found app, can't process images, app_code={app_code}")

  157.                 return None

  158.             plugins = app_info.get("data").get("plugins")

  159.             for plugin in plugins:

  160.                 if plugin.get("input_type") and "IMAGE" in plugin.get("input_type"):

  161.                     enable_image_input = True

  162.             if not enable_image_input:

  163.                 return

  164.             msg = img_cache.get("msg")

  165.             path = img_cache.get("path")

  166.             msg.prepare()

  167.             logger.info(f"[LinkAI] query with images, path={path}")

  168.             messages = self._build_vision_msg(query, path)

  169.             memory.USER_IMAGE_CACHE[session_id] = None

  170.             return messages

  171.         except Exception as e:

  172.             logger.exception(e)

  173. 

  174.     def _find_group_mapping_code(self, context):

  175.         try:

  176.             if context.kwargs.get("isgroup"):

  177.                 group_name = context.kwargs.get("msg").from_user_nickname

  178.                 if config.plugin_config and config.plugin_config.get("linkai"):

  179.                     linkai_config = config.plugin_config.get("linkai")

  180.                     group_mapping = linkai_config.get("group_app_map")

  181.                     if group_mapping and group_name:

  182.                         return group_mapping.get(group_name)

  183.         except Exception as e:

  184.             logger.exception(e)

  185.             return None

  186. 

  187.     def _build_vision_msg(self, query: str, path: str):

  188.         try:

  189.             suffix = utils.get_path_suffix(path)

  190.             with open(path, "rb") as file:

  191.                 base64_str = base64.b64encode(file.read()).decode('utf-8')

  192.                 messages = [{

  193.                     "role": "user",

  194.                     "content": [

  195.                         {

  196.                             "type": "text",

  197.                             "text": query

  198.                         },

  199.                         {

  200.                             "type": "image_url",

  201.                             "image_url": {

  202.                                 "url": f"data:image/{suffix};base64,{base64_str}"

  203.                             }

  204.                         }

  205.                     ]

  206.                 }]

  207.                 return messages

  208.         except Exception as e:

  209.             logger.exception(e)

  210. 

  211.     def reply_text(self, session: ChatGPTSession, app_code="", retry_count=0) -> dict:

  212.         if retry_count >= 2:

  213.             # exit from retry 2 times

  214.             logger.warn("[LINKAI] failed after maximum number of retry times")

  215.             return {

  216.                 "total_tokens": 0,

  217.                 "completion_tokens": 0,

  218.                 "content": "请再问我一次吧"

  219.             }

  220. 

  221.         try:

  222.             body = {

  223.                 "app_code": app_code,

  224.                 "messages": session.messages,

  225.                 "model": conf().get("model") or "gpt-3.5-turbo",  # 对话模型的名称, 支持 gpt-3.5-turbo, gpt-3.5-turbo-16k, gpt-4, wenxin, xunfei

  226.                 "temperature": conf().get("temperature"),

  227.                 "top_p": conf().get("top_p", 1),

  228.                 "frequency_penalty": conf().get("frequency_penalty", 0.0),  # [-2,2]之间，该值越大则更倾向于产生不同的内容

  229.                 "presence_penalty": conf().get("presence_penalty", 0.0),  # [-2,2]之间，该值越大则更倾向于产生不同的内容

  230.             }

  231.             if self.args.get("max_tokens"):

  232.                 body["max_tokens"] = self.args.get("max_tokens")

  233.             headers = {"Authorization": "Bearer " +  conf().get("linkai_api_key")}

  234. 

  235.             # do http request

  236.             base_url = conf().get("linkai_api_base", "https://api.link-ai.chat")

  237.             res = requests.post(url=base_url + "/v1/chat/completions", json=body, headers=headers,

  238.                                 timeout=conf().get("request_timeout", 180))

  239.             if res.status_code == 200:

  240.                 # execute success

  241.                 response = res.json()

  242.                 reply_content = response["choices"][0]["message"]["content"]

  243.                 total_tokens = response["usage"]["total_tokens"]

  244.                 logger.info(f"[LINKAI] reply={reply_content}, total_tokens={total_tokens}")

  245.                 return {

  246.                     "total_tokens": total_tokens,

  247.                     "completion_tokens": response["usage"]["completion_tokens"],

  248.                     "content": reply_content,

  249.                 }

  250. 

  251.             else:

  252.                 response = res.json()

  253.                 error = response.get("error")

  254.                 logger.error(f"[LINKAI] chat failed, status_code={res.status_code}, "

  255.                              f"msg={error.get('message')}, type={error.get('type')}")

  256. 

  257.                 if res.status_code >= 500:

  258.                     # server error, need retry

  259.                     time.sleep(2)

  260.                     logger.warn(f"[LINKAI] do retry, times={retry_count}")

  261.                     return self.reply_text(session, app_code, retry_count + 1)

  262. 

  263.                 return {

  264.                     "total_tokens": 0,

  265.                     "completion_tokens": 0,

  266.                     "content": "提问太快啦，请休息一下再问我吧"

  267.                 }

  268. 

  269.         except Exception as e:

  270.             logger.exception(e)

  271.             # retry

  272.             time.sleep(2)

  273.             logger.warn(f"[LINKAI] do retry, times={retry_count}")

  274.             return self.reply_text(session, app_code, retry_count + 1)

  275. 

  276.     def _fetch_app_info(self, app_code: str):

  277.         headers = {"Authorization": "Bearer " + conf().get("linkai_api_key")}

  278.         # do http request

  279.         base_url = conf().get("linkai_api_base", "https://api.link-ai.chat")

  280.         params = {"app_code": app_code}

  281.         res = requests.get(url=base_url + "/v1/app/info", params=params, headers=headers, timeout=(5, 10))

  282.         if res.status_code == 200:

  283.             return res.json()

  284.         else:

  285.             logger.warning(f"[LinkAI] find app info exception, res={res}")

  286. 

  287.     def create_img(self, query, retry_count=0, api_key=None):

  288.         try:

  289.             logger.info("[LinkImage] image_query={}".format(query))

  290.             headers = {

  291.                 "Content-Type": "application/json",

  292.                 "Authorization": f"Bearer {conf().get('linkai_api_key')}"

  293.             }

  294.             data = {

  295.                 "prompt": query,

  296.                 "n": 1,

  297.                 "model": conf().get("text_to_image") or "dall-e-2",

  298.                 "response_format": "url",

  299.                 "img_proxy": conf().get("image_proxy")

  300.             }

  301.             url = conf().get("linkai_api_base", "https://api.link-ai.chat") + "/v1/images/generations"

  302.             res = requests.post(url, headers=headers, json=data, timeout=(5, 90))

  303.             t2 = time.time()

  304.             image_url = res.json()["data"][0]["url"]

  305.             logger.info("[OPEN_AI] image_url={}".format(image_url))

  306.             return True, image_url

  307. 

  308.         except Exception as e:

  309.             logger.error(format(e))

  310.             return False, "画图出现问题，请休息一下再问我吧"

  311. 

  312. 

  313.     def _fetch_knowledge_search_suffix(self, response) -> str:

  314.         try:

  315.             if response.get("knowledge_base"):

  316.                 search_hit = response.get("knowledge_base").get("search_hit")

  317.                 first_similarity = response.get("knowledge_base").get("first_similarity")

  318.                 logger.info(f"[LINKAI] knowledge base, search_hit={search_hit}, first_similarity={first_similarity}")

  319.                 plugin_config = pconf("linkai")

  320.                 if plugin_config and plugin_config.get("knowledge_base") and plugin_config.get("knowledge_base").get("search_miss_text_enabled"):

  321.                     search_miss_similarity = plugin_config.get("knowledge_base").get("search_miss_similarity")

  322.                     search_miss_text = plugin_config.get("knowledge_base").get("search_miss_suffix")

  323.                     if not search_hit:

  324.                         return search_miss_text

  325.                     if search_miss_similarity and float(search_miss_similarity) > first_similarity:

  326.                         return search_miss_text

  327.         except Exception as e:

  328.             logger.exception(e)

  329. 

  330. 

  331.     def _fetch_agent_suffix(self, response):

  332.         try:

  333.             plugin_list = []

  334.             logger.debug(f"[LinkAgent] res={response}")

  335.             if response.get("agent") and response.get("agent").get("chain") and response.get("agent").get("need_show_plugin"):

  336.                 chain = response.get("agent").get("chain")

  337.                 suffix = "\n\n- - - - - - - - - - - -"

  338.                 i = 0

  339.                 for turn in chain:

  340.                     plugin_name = turn.get('plugin_name')

  341.                     suffix += "\n"

  342.                     need_show_thought = response.get("agent").get("need_show_thought")

  343.                     if turn.get("thought") and plugin_name and need_show_thought:

  344.                         suffix += f"{turn.get('thought')}\n"

  345.                     if plugin_name:

  346.                         plugin_list.append(turn.get('plugin_name'))

  347.                         suffix += f"{turn.get('plugin_icon')} {turn.get('plugin_name')}"

  348.                         if turn.get('plugin_input'):

  349.                             suffix += f"：{turn.get('plugin_input')}"

  350.                     if i < len(chain) - 1:

  351.                         suffix += "\n"

  352.                     i += 1

  353.                 logger.info(f"[LinkAgent] use plugins: {plugin_list}")

  354.                 return suffix

  355.         except Exception as e:

  356.             logger.exception(e)

  357. 

  358. 

  359.     def _send_image(self, channel, context, image_urls):

  360.         if not image_urls:

  361.             return

  362.         try:

  363.             for url in image_urls:

  364.                 reply = Reply(ReplyType.IMAGE_URL, url)

  365.                 channel.send(reply, context)

  366.         except Exception as e:

  367.             logger.error(e)

  368. 

  369. 

  370. class LinkAISessionManager(SessionManager):

  371.     def session_msg_query(self, query, session_id):

  372.         session = self.build_session(session_id)

  373.         messages = session.messages + [{"role": "user", "content": query}]

  374.         return messages

  375. 

  376.     def session_reply(self, reply, session_id, total_tokens=None, query=None):

  377.         session = self.build_session(session_id)

  378.         if query:

  379.             session.add_query(query)

  380.         session.add_reply(reply)

  381.         try:

  382.             max_tokens = conf().get("conversation_max_tokens", 2500)

  383.             tokens_cnt = session.discard_exceeding(max_tokens, total_tokens)

  384.             logger.debug(f"[LinkAI] chat history, before tokens={total_tokens}, now tokens={tokens_cnt}")

  385.         except Exception as e:

  386.             logger.warning("Exception when counting tokens precisely for session: {}".format(str(e)))

  387.         return session

  388. 

  389. 

  390. class LinkAISession(ChatGPTSession):

  391.     def calc_tokens(self):

  392.         if not self.messages:

  393.             return 0

  394.         return len(str(self.messages))

  395. 

  396.     def discard_exceeding(self, max_tokens, cur_tokens=None):

  397.         cur_tokens = self.calc_tokens()

  398.         if cur_tokens > max_tokens:

  399.             for i in range(0, len(self.messages)):

  400.                 if i > 0 and self.messages[i].get("role") == "assistant" and self.messages[i - 1].get("role") == "user":

  401.                     self.messages.pop(i)

  402.                     self.messages.pop(i - 1)

  403.                     return self.calc_tokens()

  404.         return cur_tokens