Merge pull request #1450 from zhayujie/feat-doc-chat

feat: 文档总结和与内容对话

bot/linkai/link_ai_bot.py

@@ -79,7 +79,10 @@ class LinkAIBot(Bot, OpenAIImage):
                 "frequency_penalty": conf().get("frequency_penalty", 0.0),  # [-2,2]之间，该值越大则更倾向于产生不同的内容
                 "presence_penalty": conf().get("presence_penalty", 0.0),  # [-2,2]之间，该值越大则更倾向于产生不同的内容
             }
-            logger.info(f"[LINKAI] query={query}, app_code={app_code}, mode={body.get('model')}")
+            file_id = context.kwargs.get("file_id")
+            if file_id:
+                body["file_id"] = file_id
+            logger.info(f"[LINKAI] query={query}, app_code={app_code}, mode={body.get('model')}, file_id={file_id}")
             headers = {"Authorization": "Bearer " + linkai_api_key}
 
             # do http request

bridge/bridge.py

@@ -32,6 +32,7 @@ class Bridge(object):
         if model_type in ["claude"]:
             self.btype["chat"] = const.CLAUDEAI
         self.bots = {}
+        self.chat_bots = {}
 
     def get_bot(self, typename):
         if self.bots.get(typename) is None:
@@ -61,6 +62,11 @@ class Bridge(object):
     def fetch_translate(self, text, from_lang="", to_lang="en") -> Reply:
         return self.get_bot("translate").translate(text, from_lang, to_lang)
 
+    def find_chat_bot(self, bot_type: str):
+        if self.chat_bots.get(bot_type) is None:
+            self.chat_bots[bot_type] = create_bot(bot_type)
+        return self.chat_bots.get(bot_type)
+
     def reset_bot(self):
         """
         重置bot路由

bridge/context.py

@@ -9,6 +9,7 @@ class ContextType(Enum):
     IMAGE = 3  # 图片消息
     FILE = 4  # 文件信息
     VIDEO = 5  # 视频信息
+    SHARING = 6  # 分享信息
 
     IMAGE_CREATE = 10  # 创建图片命令
     JOIN_GROUP = 20  # 加入群聊

channel/wechat/wechat_channel.py

@@ -25,7 +25,7 @@ from lib import itchat
 from lib.itchat.content import *
 
 
-@itchat.msg_register([TEXT, VOICE, PICTURE, NOTE])
+@itchat.msg_register([TEXT, VOICE, PICTURE, NOTE, ATTACHMENT, SHARING])
 def handler_single_msg(msg):
     try:
         cmsg = WechatMessage(msg, False)
@@ -36,7 +36,7 @@ def handler_single_msg(msg):
     return None
 
 
-@itchat.msg_register([TEXT, VOICE, PICTURE, NOTE], isGroupChat=True)
+@itchat.msg_register([TEXT, VOICE, PICTURE, NOTE, ATTACHMENT, SHARING], isGroupChat=True)
 def handler_group_msg(msg):
     try:
         cmsg = WechatMessage(msg, True)
@@ -172,6 +172,8 @@ class WechatChannel(ChatChannel):
         elif cmsg.ctype == ContextType.TEXT:
             # logger.debug("[WX]receive group msg: {}, cmsg={}".format(json.dumps(cmsg._rawmsg, ensure_ascii=False), cmsg))
             pass
+        elif cmsg.ctype == ContextType.FILE:
+            logger.debug(f"[WX]receive attachment msg, file_name={cmsg.content}")
         else:
             logger.debug("[WX]receive group msg: {}".format(cmsg.content))
         context = self._compose_context(cmsg.ctype, cmsg.content, isgroup=True, msg=cmsg)

channel/wechat/wechat_message.py

@@ -7,7 +7,6 @@ from common.tmp_dir import TmpDir
 from lib import itchat
 from lib.itchat.content import *
 
-
 class WechatMessage(ChatMessage):
     def __init__(self, itchat_msg, is_group=False):
         super().__init__(itchat_msg)
@@ -42,6 +41,14 @@ class WechatMessage(ChatMessage):
                     self.actual_user_nickname = re.findall(r"\"(.*?)\"", itchat_msg["Content"])[0]
             else:
                 raise NotImplementedError("Unsupported note message: " + itchat_msg["Content"])
+        elif itchat_msg["Type"] == ATTACHMENT:
+            self.ctype = ContextType.FILE
+            self.content = TmpDir().path() + itchat_msg["FileName"]
+            self._prepare_fn = lambda: itchat_msg.download(self.content)
+        elif itchat_msg["Type"] == SHARING:
+            self.ctype = ContextType.SHARING
+            self.content = itchat_msg.get("Url")
+
         else:
             raise NotImplementedError("Unsupported message type: Type:{} MsgType:{}".format(itchat_msg["Type"], itchat_msg["MsgType"]))
 

plugins/godcmd/godcmd.py

@@ -6,6 +6,7 @@ import random
 import string
 from typing import Tuple
 
+import bridge.bridge
 import plugins
 from bridge.bridge import Bridge
 from bridge.context import ContextType
@@ -310,6 +311,8 @@ class Godcmd(Plugin):
                 elif cmd == "reset":
                     if bottype in [const.OPEN_AI, const.CHATGPT, const.CHATGPTONAZURE, const.LINKAI, const.BAIDU, const.XUNFEI]:
                         bot.sessions.clear_session(session_id)
+                        if Bridge().chat_bots.get(bottype):
+                            Bridge().chat_bots.get(bottype).sessions.clear_session(session_id)
                         channel.cancel_session(session_id)
                         ok, result = True, "会话已重置"
                     else:

plugins/linkai/README.md

@@ -1,10 +1,10 @@
 ## 插件说明
 
-基于 LinkAI 提供的知识库、Midjourney绘画等能力对机器人的功能进行增强。平台地址: https://chat.link-ai.tech/console
+基于 LinkAI 提供的知识库、Midjourney绘画、文档对话等能力对机器人的功能进行增强。平台地址: https://chat.link-ai.tech/console
 
 ## 插件配置
 
-将 `plugins/linkai` 目录下的 `config.json.template` 配置模板复制为最终生效的 `config.json`。 (如果未配置则会默认使用`config.json.template`模板中配置，功能默认关闭，可通过指令进行开启)。
+将 `plugins/linkai` 目录下的 `config.json.template` 配置模板复制为最终生效的 `config.json`。 (如果未配置则会默认使用`config.json.template`模板中配置，但功能默认关闭，需要可通过指令进行开启)。
 
 以下是配置项说明：
 
@@ -21,19 +21,25 @@
         "max_tasks": 3,           # 支持同时提交的总任务个数
         "max_tasks_per_user": 1,  # 支持单个用户同时提交的任务个数
         "use_image_create_prefix": true   # 是否使用全局的绘画触发词，如果开启将同时支持由`config.json`中的 image_create_prefix 配置触发
+    },
+    "summary": {
+        "enabled": true,              # 文档总结和对话功能开关
+        "group_enabled": true,        # 是否支持群聊开启
+        "max_summary_words": 50000,   # 文章的最大字数，超过字数则直接忽略
+        "max_file_size": 15000        # 文件的大小限制，单位KB，超过该大小直接忽略
     }
 }
 
 ```
 注意：
 
- - 配置项中 `group_app_map` 部分是用于映射群聊与LinkAI平台上的应用， `midjourney` 部分是 mj 画图的配置，可根据需要进行填写，未填写配置时默认不开启相应功能
+ - 配置项中 `group_app_map` 部分是用于映射群聊与LinkAI平台上的应用， `midjourney` 部分是 mj 画图的配置，`summary` 部分是文档总结及对话功能的配置。三部分的配置相互独立，可按需开启
  - 实际 `config.json` 配置中应保证json格式，不应携带 '#' 及后面的注释
  - 如果是`docker`部署，可通过映射 `plugins/config.json` 到容器中来完成插件配置，参考[文档](https://github.com/zhayujie/chatgpt-on-wechat#3-%E6%8F%92%E4%BB%B6%E4%BD%BF%E7%94%A8)
 
 ## 插件使用
 
-> 使用插件中的知识库管理功能需要首先开启`linkai`对话，依赖全局 `config.json` 中的 `use_linkai` 和 `linkai_api_key` 配置；而midjourney绘画功能则只需填写 `linkai_api_key` 配置，`use_linkai` 无论是否关闭均可使用。具体可参考 [详细文档](https://link-ai.tech/platform/link-app/wechat)。
+> 使用插件中的知识库管理功能需要首先开启`linkai`对话，依赖全局 `config.json` 中的 `use_linkai` 和 `linkai_api_key` 配置；而midjourney绘画和summary文档总结对话功能则只需填写 `linkai_api_key` 配置，`use_linkai` 无论是否关闭均可使用。具体可参考 [详细文档](https://link-ai.tech/platform/link-app/wechat)。
 
 完成配置后运行项目，会自动运行插件，输入 `#help linkai` 可查看插件功能。
 
@@ -77,3 +83,22 @@
 3. 开启 `use_image_create_prefix` 配置后可直接复用全局画图触发词，以"画"开头便可以生成图片。
 4. 提示词内容中包含敏感词或者参数格式错误可能导致绘画失败，生成失败不消耗积分
 5. 若未收到图片可能有两种可能，一种是收到了图片但微信发送失败，可以在后台日志查看有没有获取到图片url，一般原因是受到了wx限制，可以稍后重试或更换账号尝试；另一种情况是图片提示词存在疑似违规，mj不会直接提示错误但会在画图后删掉原图导致程序无法获取，这种情况不消耗积分。
+
+### 3.文档总结对话功能
+
+#### 配置
+
+该功能依赖 LinkAI的知识库及对话功能，需要在项目根目录的config.json中设置 `linkai_api_key`， 同时上述插件配置说明，添加 `summary` 部分的配置，设置 `enabled` 为 true。
+
+如果不想创建 `plugins/linkai/config.json` 配置，可以直接通过 `$linkai sum open` 指令开启该功能。
+
+#### 使用
+
+功能开启后，向机器人发送 **文件** 或 **分享链接卡片** 即可生成摘要，输入 "开启对话" 后，可以与文件或链接的内容进行多轮对话，输入 "退出对话"
+ 可关闭与内容的对话。
+
+#### 限制
+
+ 1. 文件目前 支持 `txt`, `docx`, `pdf`, `md`, `csv`格式，文件大小由 `max_file_size` 限制，最大不超过15M，文件字数由 `max_summary_words` 配置限制，最多可支持百万字的文件，但不建议上传字数过多的文件，一是token消耗过大，而是摘要很难概括到全部内容，但可以通过对话了解细节。
+ 2. 分享链接 目前仅支持 公众号文章，后续会支持更多文章类型及视频链接等
+ 3. 总结及对话的 费用与 LinkAI 3.5-4K 模型的计费方式相同，按文档内容的tokens进行计算

plugins/linkai/config.json.template

@@ -10,5 +10,11 @@
         "max_tasks": 3,
         "max_tasks_per_user": 1,
         "use_image_create_prefix": true
+    },
+    "summary": {
+        "enabled": true,
+        "group_enabled": true,
+        "max_summary_words": 50000,
+        "max_file_size": 15000
     }
 }

plugins/linkai/linkai.py

@@ -4,7 +4,10 @@ from bridge.reply import Reply, ReplyType
 from config import global_config
 from plugins import *
 from .midjourney import MJBot
+from .summary import LinkSummary
 from bridge import bridge
+from common.expired_dict import ExpiredDict
+from common import const
 
 
 @plugins.register(
@@ -12,6 +15,7 @@ from bridge import bridge
     desc="A plugin that supports knowledge base and midjourney drawing.",
     version="0.1.0",
     author="https://link-ai.tech",
+    desire_priority=99
 )
 class LinkAI(Plugin):
     def __init__(self):
@@ -23,8 +27,12 @@ class LinkAI(Plugin):
             self.config = self._load_config_template()
         if self.config:
             self.mj_bot = MJBot(self.config.get("midjourney"))
+        self.sum_config = {}
+        if self.config:
+            self.sum_config = self.config.get("summary")
         logger.info("[LinkAI] inited")
 
+
     def on_handle_context(self, e_context: EventContext):
         """
         消息处理逻辑
@@ -34,10 +42,37 @@ class LinkAI(Plugin):
             return
 
         context = e_context['context']
-        if context.type not in [ContextType.TEXT, ContextType.IMAGE, ContextType.IMAGE_CREATE]:
+        if context.type not in [ContextType.TEXT, ContextType.IMAGE, ContextType.IMAGE_CREATE, ContextType.FILE, ContextType.SHARING]:
             # filter content no need solve
             return
 
+        if context.type == ContextType.FILE and self._is_summary_open(context):
+            # 文件处理
+            context.get("msg").prepare()
+            file_path = context.content
+            if not LinkSummary().check_file(file_path, self.sum_config):
+                return
+            _send_info(e_context, "正在为你加速生成摘要，请稍后")
+            res = LinkSummary().summary_file(file_path)
+            if not res:
+                _set_reply_text("总结出现异常，请稍后再试吧", e_context)
+                return
+            USER_FILE_MAP[_find_user_id(context) + "-sum_id"] = res.get("summary_id")
+            _set_reply_text(res.get("summary") + "\n\n💬 发送 \"开启对话\" 可以开启与文件内容的对话", e_context, level=ReplyType.TEXT)
+            return
+
+        if context.type == ContextType.SHARING and self._is_summary_open(context):
+            if not LinkSummary().check_url(context.content):
+                return
+            _send_info(e_context, "正在为你加速生成摘要，请稍后")
+            res = LinkSummary().summary_url(context.content)
+            if not res:
+                _set_reply_text("总结出现异常，请稍后再试吧", e_context)
+                return
+            _set_reply_text(res.get("summary") + "\n\n💬 发送 \"开启对话\" 可以开启与文章内容的对话", e_context, level=ReplyType.TEXT)
+            USER_FILE_MAP[_find_user_id(context) + "-sum_id"] = res.get("summary_id")
+            return
+
         mj_type = self.mj_bot.judge_mj_task_type(e_context)
         if mj_type:
             # MJ作图任务处理
@@ -49,10 +84,38 @@ class LinkAI(Plugin):
             self._process_admin_cmd(e_context)
             return
 
+        if context.type == ContextType.TEXT and context.content == "开启对话" and _find_sum_id(context):
+            # 文本对话
+            _send_info(e_context, "正在为你开启对话，请稍后")
+            res = LinkSummary().summary_chat(_find_sum_id(context))
+            if not res:
+                _set_reply_text("开启对话失败，请稍后再试吧", e_context)
+                return
+            USER_FILE_MAP[_find_user_id(context) + "-file_id"] = res.get("file_id")
+            _set_reply_text("💡你可以问我关于这篇文章的任何问题，例如：\n\n" + res.get("questions") + "\n\n发送 \"退出对话\" 可以关闭与文章的对话", e_context, level=ReplyType.TEXT)
+            return
+
+        if context.type == ContextType.TEXT and context.content == "退出对话" and _find_file_id(context):
+            del USER_FILE_MAP[_find_user_id(context) + "-file_id"]
+            bot = bridge.Bridge().find_chat_bot(const.LINKAI)
+            bot.sessions.clear_session(context["session_id"])
+            _set_reply_text("对话已退出", e_context, level=ReplyType.TEXT)
+            return
+
+        if context.type == ContextType.TEXT and _find_file_id(context):
+            bot = bridge.Bridge().find_chat_bot(const.LINKAI)
+            context.kwargs["file_id"] = _find_file_id(context)
+            reply = bot.reply(context.content, context)
+            e_context["reply"] = reply
+            e_context.action = EventAction.BREAK_PASS
+            return
+
+
         if self._is_chat_task(e_context):
             # 文本对话任务处理
             self._process_chat_task(e_context)
 
+
     # 插件管理功能
     def _process_admin_cmd(self, e_context: EventContext):
         context = e_context['context']
@@ -94,11 +157,31 @@ class LinkAI(Plugin):
             # 保存插件配置
             super().save_config(self.config)
             _set_reply_text(f"应用设置成功: {app_code}", e_context, level=ReplyType.INFO)
+
+        if len(cmd) == 3 and cmd[1] == "sum" and (cmd[2] == "open" or cmd[2] == "close"):
+            # 知识库开关指令
+            if not _is_admin(e_context):
+                _set_reply_text("需要管理员权限执行", e_context, level=ReplyType.ERROR)
+                return
+            is_open = True
+            tips_text = "开启"
+            if cmd[2] == "close":
+                tips_text = "关闭"
+                is_open = False
+            self.sum_config["enabled"] = is_open
+            _set_reply_text(f"文章总结功能{tips_text}", e_context, level=ReplyType.INFO)
         else:
             _set_reply_text(f"指令错误，请输入{_get_trigger_prefix()}linkai help 获取帮助", e_context,
                             level=ReplyType.INFO)
             return
 
+    def _is_summary_open(self, context) -> bool:
+        if not self.sum_config or not self.sum_config.get("enabled"):
+            return False
+        if not context.kwargs.get("isgroup") and not self.sum_config.get("group_enabled"):
+            return False
+        return True
+
     # LinkAI 对话任务处理
     def _is_chat_task(self, e_context: EventContext):
         context = e_context['context']
@@ -112,7 +195,7 @@ class LinkAI(Plugin):
         """
         context = e_context['context']
         # 群聊应用管理
-        group_name = context.kwargs.get("msg").from_user_nickname
+        group_name = context.get("msg").from_user_nickname
         app_code = self._fetch_group_app_code(group_name)
         if app_code:
             context.kwargs['app_code'] = app_code
@@ -130,7 +213,7 @@ class LinkAI(Plugin):
 
     def get_help_text(self, verbose=False, **kwargs):
         trigger_prefix = _get_trigger_prefix()
-        help_text = "用于集成 LinkAI 提供的知识库、Midjourney绘画等能力。\n\n"
+        help_text = "用于集成 LinkAI 提供的知识库、Midjourney绘画、文档总结对话等能力。\n\n"
         if not verbose:
             return help_text
         help_text += f'📖 知识库\n - 群聊中指定应用: {trigger_prefix}linkai app 应用编码\n'
@@ -140,6 +223,7 @@ class LinkAI(Plugin):
         help_text += f"🎨 绘画\n - 生成: {trigger_prefix}mj 描述词1, 描述词2.. \n - 放大: {trigger_prefix}mju 图片ID 图片序号\n - 变换: {trigger_prefix}mjv 图片ID 图片序号\n - 重置: {trigger_prefix}mjr 图片ID"
         help_text += f"\n\n例如：\n\"{trigger_prefix}mj a little cat, white --ar 9:16\"\n\"{trigger_prefix}mju 11055927171882 2\""
         help_text += f"\n\"{trigger_prefix}mjv 11055927171882 2\"\n\"{trigger_prefix}mjr 11055927171882\""
+        help_text += f"\n\n💡 文档总结和对话\n - 开启: {trigger_prefix}linkai sum open\n - 使用: 发送文件、公众号文章等可生成摘要，并与内容对话"
         return help_text
 
     def _load_config_template(self):
@@ -150,10 +234,17 @@ class LinkAI(Plugin):
                 with open(plugin_config_path, "r", encoding="utf-8") as f:
                     plugin_conf = json.load(f)
                     plugin_conf["midjourney"]["enabled"] = False
+                    plugin_conf["summary"]["enabled"] = False
                     return plugin_conf
         except Exception as e:
             logger.exception(e)
 
+
+def _send_info(e_context: EventContext, content: str):
+    reply = Reply(ReplyType.TEXT, content)
+    channel = e_context["channel"]
+    channel.send(reply, e_context["context"])
+
 # 静态方法
 def _is_admin(e_context: EventContext) -> bool:
     """
@@ -168,11 +259,25 @@ def _is_admin(e_context: EventContext) -> bool:
         return context["receiver"] in global_config["admin_users"]
 
 
+def _find_user_id(context):
+    if context["isgroup"]:
+        return context.kwargs.get("msg").actual_user_id
+    else:
+        return context["receiver"]
+
+
 def _set_reply_text(content: str, e_context: EventContext, level: ReplyType = ReplyType.ERROR):
     reply = Reply(level, content)
     e_context["reply"] = reply
     e_context.action = EventAction.BREAK_PASS
 
-
 def _get_trigger_prefix():
     return conf().get("plugin_trigger_prefix", "$")
+
+def _find_sum_id(context):
+    return USER_FILE_MAP.get(_find_user_id(context) + "-sum_id")
+
+def _find_file_id(context):
+    return USER_FILE_MAP.get(_find_user_id(context) + "-file_id")
+
+USER_FILE_MAP = ExpiredDict(60 * 60)

plugins/linkai/midjourney.py

@@ -5,7 +5,6 @@ import requests
 import threading
 import time
 from bridge.reply import Reply, ReplyType
-import aiohttp
 import asyncio
 from bridge.context import ContextType
 from plugins import EventContext, EventAction

plugins/linkai/summary.py

@@ -0,0 +1,91 @@
+import requests
+from config import conf
+from common.log import logger
+import os
+
+
+class LinkSummary:
+    def __init__(self):
+        pass
+
+    def summary_file(self, file_path: str):
+        file_body = {
+            "file": open(file_path, "rb"),
+            "name": file_path.split("/")[-1],
+        }
+        res = requests.post(url=self.base_url() + "/v1/summary/file", headers=self.headers(), files=file_body, timeout=(5, 180))
+        return self._parse_summary_res(res)
+
+    def summary_url(self, url: str):
+        body = {
+            "url": url
+        }
+        res = requests.post(url=self.base_url() + "/v1/summary/url", headers=self.headers(), json=body, timeout=(5, 180))
+        return self._parse_summary_res(res)
+
+    def summary_chat(self, summary_id: str):
+        body = {
+            "summary_id": summary_id
+        }
+        res = requests.post(url=self.base_url() + "/v1/summary/chat", headers=self.headers(), json=body, timeout=(5, 180))
+        if res.status_code == 200:
+            res = res.json()
+            logger.debug(f"[LinkSum] chat open, res={res}")
+            if res.get("code") == 200:
+                data = res.get("data")
+                return {
+                    "questions": data.get("questions"),
+                    "file_id": data.get("file_id")
+                }
+        else:
+            res_json = res.json()
+            logger.error(f"[LinkSum] summary error, status_code={res.status_code}, msg={res_json.get('message')}")
+            return None
+
+    def _parse_summary_res(self, res):
+        if res.status_code == 200:
+            res = res.json()
+            logger.debug(f"[LinkSum] url summary, res={res}")
+            if res.get("code") == 200:
+                data = res.get("data")
+                return {
+                    "summary": data.get("summary"),
+                    "summary_id": data.get("summary_id")
+                }
+        else:
+            res_json = res.json()
+            logger.error(f"[LinkSum] summary error, status_code={res.status_code}, msg={res_json.get('message')}")
+            return None
+
+    def base_url(self):
+        return conf().get("linkai_api_base", "https://api.link-ai.chat")
+
+    def headers(self):
+        return {"Authorization": "Bearer " + conf().get("linkai_api_key")}
+
+    def check_file(self, file_path: str, sum_config: dict) -> bool:
+        file_size = os.path.getsize(file_path) // 1000
+        with open(file_path, 'r') as f:
+            content = f.read()
+            word_count = len(content)
+
+        if (sum_config.get("max_file_size") and file_size > sum_config.get("max_file_size")) or file_size > 15000\
+                or (sum_config.get("max_summary_words") and word_count > sum_config.get("max_summary_words")):
+            logger.warn(f"[LinkSum] file size exceeds limit, No processing, file_size={file_size}KB, word_count={word_count}")
+            return True
+
+        suffix = file_path.split(".")[-1]
+        support_list = ["txt", "csv", "docx", "pdf", "md"]
+        if suffix not in support_list:
+            logger.warn(f"[LinkSum] unsupported file, suffix={suffix}, support_list={support_list}")
+            return False
+
+        return True
+
+    def check_url(self, url: str):
+        support_list = ["mp.weixin.qq.com"]
+        for support_url in support_list:
+            if support_url in url:
+                return True
+        logger.warn("[LinkSum] unsupported url")
+        return False