J1147
/
ai-on-wechat
Beobachten 1
Favorisieren 0
Fork 0
Code Issues 0 Pull-Requests 0 Releases 0 Wiki Aktivität
Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
1008 Commits
2 Branches
152MB
Struktur: f71820bf4e
Branches Tags
${ item.name }
Erstelle Branch ${ searchTerm }
von „f71820bf4e“
${ noResults }
ai-on-wechat/bot/linkai/link_ai_bot.py

469 Zeilen
21KB
Originalformat Blame Verlauf 

  1. # access LinkAI knowledge base platform

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

  3. 

  4. import re

  5. import time

  6. import requests

  7. import config

  8. from bot.bot import Bot

  9. from bot.chatgpt.chat_gpt_session import ChatGPTSession

  10. from bot.session_manager import SessionManager

  11. from bridge.context import Context, ContextType

  12. from bridge.reply import Reply, ReplyType

  13. from common.log import logger

  14. from config import conf, pconf

  15. import threading

  16. from common import memory, utils

  17. import base64

  18. import os

  19. 

  20. class LinkAIBot(Bot):

  21.     # authentication failed

  22.     AUTH_FAILED_CODE = 401

  23.     NO_QUOTA_CODE = 406

  24. 

  25.     def __init__(self):

  26.         super().__init__()

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

  28.         self.args = {}

  29. 

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

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

  32.             return self._chat(query, context)

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

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

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

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

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

  38.             if ok:

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

  40.             else:

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

  42.             return reply

  43.         else:

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

  45.             return reply

  46. 

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

  48.         """

  49.         发起对话请求

  50.         :param query: 请求提示词

  51.         :param context: 对话上下文

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

  53.         :return: 回复

  54.         """

  55.         if retry_count > 2:

  56.             # exit from retry 2 times

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

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

  59. 

  60.         try:

  61.             # load config

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

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

  64.                 app_code = None

  65.             else:

  66.                 plugin_app_code = self._find_group_mapping_code(context)

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

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

  69. 

  70.             session_id = context["session_id"]

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

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

  73. 

  74.             # image process

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

  76.             if img_cache:

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

  78.                 if messages:

  79.                     session_message = messages

  80. 

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

  82.             # remove system message

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

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

  85.                     session_message.pop(0)

  86.             body = {

  87.                 "app_code": app_code,

  88.                 "messages": session_message,

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

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

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

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

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

  94.                 "session_id": session_id,

  95.                 "sender_id": session_id,

  96.                 "channel_type": conf().get("channel_type", "wx")

  97.             }

  98.             try:

  99.                 from linkai import LinkAIClient

  100.                 client_id = LinkAIClient.fetch_client_id()

  101.                 if client_id:

  102.                     body["client_id"] = client_id

  103.                     # start: client info deliver

  104.                     if context.kwargs.get("msg"):

  105.                         body["session_id"] = context.kwargs.get("msg").from_user_id

  106.                         if context.kwargs.get("msg").is_group:

  107.                             body["is_group"] = True

  108.                             body["group_name"] = context.kwargs.get("msg").from_user_nickname

  109.                             body["sender_name"] = context.kwargs.get("msg").actual_user_nickname

  110.                         else:

  111.                             if body.get("channel_type") in ["wechatcom_app"]:

  112.                                 body["sender_name"] = context.kwargs.get("msg").from_user_id

  113.                             else:

  114.                                 body["sender_name"] = context.kwargs.get("msg").from_user_nickname

  115. 

  116.             except Exception as e:

  117.                 pass

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

  119.             if file_id:

  120.                 body["file_id"] = file_id

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

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

  123. 

  124.             # do http request

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

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

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

  128.             if res.status_code == 200:

  129.                 # execute success

  130.                 response = res.json()

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

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

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

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

  135.     

  136.                 agent_suffix = self._fetch_agent_suffix(response)

  137.                 if agent_suffix:

  138.                     reply_content += agent_suffix

  139.                 if not agent_suffix:

  140.                     knowledge_suffix = self._fetch_knowledge_search_suffix(response)

  141.                     if knowledge_suffix:

  142.                         reply_content += knowledge_suffix

  143.                 # image process

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

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

  146.                     thread.start()

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

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

  149.                 reply_content = self._process_url(reply_content)

  150.                 return Reply(ReplyType.TEXT, reply_content)

  151. 

  152.             else:

  153.                 response = res.json()

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

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

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

  157. 

  158.                 if res.status_code >= 500:

  159.                     # server error, need 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. 

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

  165. 

  166.         except Exception as e:

  167.             logger.exception(e)

  168.             # retry

  169.             time.sleep(2)

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

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

  172. 

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

  174.         try:

  175.             enable_image_input = False

  176.             app_info = self._fetch_app_info(app_code)

  177.             if not app_info:

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

  179.                 return None

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

  181.             for plugin in plugins:

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

  183.                     enable_image_input = True

  184.             if not enable_image_input:

  185.                 return

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

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

  188.             msg.prepare()

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

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

  191.             memory.USER_IMAGE_CACHE[session_id] = None

  192.             return messages

  193.         except Exception as e:

  194.             logger.exception(e)

  195. 

  196.     def _find_group_mapping_code(self, context):

  197.         try:

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

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

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

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

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

  203.                     if group_mapping and group_name:

  204.                         return group_mapping.get(group_name)

  205.         except Exception as e:

  206.             logger.exception(e)

  207.             return None

  208. 

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

  210.         try:

  211.             suffix = utils.get_path_suffix(path)

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

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

  214.                 messages = [{

  215.                     "role": "user",

  216.                     "content": [

  217.                         {

  218.                             "type": "text",

  219.                             "text": query

  220.                         },

  221.                         {

  222.                             "type": "image_url",

  223.                             "image_url": {

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

  225.                             }

  226.                         }

  227.                     ]

  228.                 }]

  229.                 return messages

  230.         except Exception as e:

  231.             logger.exception(e)

  232. 

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

  234.         if retry_count >= 2:

  235.             # exit from retry 2 times

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

  237.             return {

  238.                 "total_tokens": 0,

  239.                 "completion_tokens": 0,

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

  241.             }

  242. 

  243.         try:

  244.             body = {

  245.                 "app_code": app_code,

  246.                 "messages": session.messages,

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

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

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

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

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

  252.             }

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

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

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

  256. 

  257.             # do http request

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

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

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

  261.             if res.status_code == 200:

  262.                 # execute success

  263.                 response = res.json()

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

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

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

  267.                 return {

  268.                     "total_tokens": total_tokens,

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

  270.                     "content": reply_content,

  271.                 }

  272. 

  273.             else:

  274.                 response = res.json()

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

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

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

  278. 

  279.                 if res.status_code >= 500:

  280.                     # server error, need retry

  281.                     time.sleep(2)

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

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

  284. 

  285.                 return {

  286.                     "total_tokens": 0,

  287.                     "completion_tokens": 0,

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

  289.                 }

  290. 

  291.         except Exception as e:

  292.             logger.exception(e)

  293.             # retry

  294.             time.sleep(2)

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

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

  297. 

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

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

  300.         # do http request

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

  302.         params = {"app_code": app_code}

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

  304.         if res.status_code == 200:

  305.             return res.json()

  306.         else:

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

  308. 

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

  310.         try:

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

  312.             headers = {

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

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

  315.             }

  316.             data = {

  317.                 "prompt": query,

  318.                 "n": 1,

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

  320.                 "response_format": "url",

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

  322.             }

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

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

  325.             t2 = time.time()

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

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

  328.             return True, image_url

  329. 

  330.         except Exception as e:

  331.             logger.error(format(e))

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

  333. 

  334. 

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

  336.         try:

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

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

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

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

  341.                 plugin_config = pconf("linkai")

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

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

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

  345.                     if not search_hit:

  346.                         return search_miss_text

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

  348.                         return search_miss_text

  349.         except Exception as e:

  350.             logger.exception(e)

  351. 

  352. 

  353.     def _fetch_agent_suffix(self, response):

  354.         try:

  355.             plugin_list = []

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

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

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

  359.                 suffix = "\n\n- - - - - - - - - - - -"

  360.                 i = 0

  361.                 for turn in chain:

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

  363.                     suffix += "\n"

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

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

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

  367.                     if plugin_name:

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

  369.                         if turn.get('plugin_icon'):

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

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

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

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

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

  375.                         suffix += "\n"

  376.                     i += 1

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

  378.                 return suffix

  379.         except Exception as e:

  380.             logger.exception(e)

  381. 

  382.     def _process_url(self, text):

  383.         try:

  384.             url_pattern = re.compile(r'\[(.*?)\]\((http[s]?://.*?)\)')

  385.             def replace_markdown_url(match):

  386.                 return f"{match.group(2)}"

  387.             return url_pattern.sub(replace_markdown_url, text)

  388.         except Exception as e:

  389.             logger.error(e)

  390. 

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

  392.         if not image_urls:

  393.             return

  394.         max_send_num = conf().get("max_media_send_count")

  395.         send_interval = conf().get("media_send_interval")

  396.         try:

  397.             i = 0

  398.             for url in image_urls:

  399.                 if max_send_num and i >= max_send_num:

  400.                     continue

  401.                 i += 1

  402.                 if url.endswith(".mp4"):

  403.                     reply_type = ReplyType.VIDEO_URL

  404.                 elif url.endswith(".pdf") or url.endswith(".doc") or url.endswith(".docx") or url.endswith(".csv"):

  405.                     reply_type = ReplyType.FILE

  406.                     url = _download_file(url)

  407.                     if not url:

  408.                         continue

  409.                 else:

  410.                     reply_type = ReplyType.IMAGE_URL

  411.                 reply = Reply(reply_type, url)

  412.                 channel.send(reply, context)

  413.                 if send_interval:

  414.                     time.sleep(send_interval)

  415.         except Exception as e:

  416.             logger.error(e)

  417. 

  418. 

  419. def _download_file(url: str):

  420.     try:

  421.         file_path = "tmp"

  422.         if not os.path.exists(file_path):

  423.             os.makedirs(file_path)

  424.         file_name = url.split("/")[-1]  # 获取文件名

  425.         file_path = os.path.join(file_path, file_name)

  426.         response = requests.get(url)

  427.         with open(file_path, "wb") as f:

  428.             f.write(response.content)

  429.         return file_path

  430.     except Exception as e:

  431.         logger.warn(e)

  432. 

  433. 

  434. class LinkAISessionManager(SessionManager):

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

  436.         session = self.build_session(session_id)

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

  438.         return messages

  439. 

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

  441.         session = self.build_session(session_id)

  442.         if query:

  443.             session.add_query(query)

  444.         session.add_reply(reply)

  445.         try:

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

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

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

  449.         except Exception as e:

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

  451.         return session

  452. 

  453. 

  454. class LinkAISession(ChatGPTSession):

  455.     def calc_tokens(self):

  456.         if not self.messages:

  457.             return 0

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

  459. 

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

  461.         cur_tokens = self.calc_tokens()

  462.         if cur_tokens > max_tokens:

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

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

  465.                     self.messages.pop(i)

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

  467.                     return self.calc_tokens()

  468.         return cur_tokens