Merge pull request #117 from wugren/MVP

Add slack tunnel

.gitignore

@@ -9,6 +9,7 @@ history.txt
 aios_shell_history.txt
 math_school_env.db
 workflows.db
-
+venv-linux
+venv_test
 
 rootfs/test_doc/

src/component/slack_tunnel.py

@@ -0,0 +1,239 @@
+import asyncio
+import datetime
+import json
+import logging
+import os.path
+import re
+import uuid
+import time
+
+import aiofiles
+import aiohttp
+from slack_bolt.adapter.socket_mode.websockets import AsyncSocketModeHandler
+from slack_bolt.app.async_app import AsyncApp
+
+from aios import KnowledgeStore, ObjectType
+from aios.frame.tunnel import AgentTunnel
+from aios.proto.agent_msg import AgentMsg
+from aios.storage.storage import AIStorage
+
+logger = logging.getLogger(__name__)
+
+
+async def download_file(url: str, file_path: str, token: str) -> None:
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url, headers={"Authorization": f"Bearer {token}"}) as resp:
+            if resp.status == 200:
+                f = await aiofiles.open(file_path, mode='wb')
+                while True:
+                    chunk = await resp.content.read(1024)
+                    if not chunk:
+                        break
+                    await f.write(chunk)
+                await f.close()
+
+
+class SlackTunnel(AgentTunnel):
+    type: str
+    token: str
+    app_token: str
+    slack_cache: str
+    app: AsyncSocketModeHandler
+
+    def __init__(self):
+        super().__init__()
+        self.type = "SlackTunnel"
+        self.token = ""
+        self.slack_cache = os.path.join(AIStorage.get_instance().get_myai_dir(), "slack")
+        if not os.path.exists(self.slack_cache):
+            os.makedirs(self.slack_cache)
+
+    @classmethod
+    def register_to_loader(cls):
+        async def load_slack_tunnel(config: dict) -> AgentTunnel:
+            result_tunnel = SlackTunnel()
+            if await result_tunnel.load_from_config(config):
+                return result_tunnel
+            else:
+                return None
+
+        AgentTunnel.register_loader("SlackTunnel", load_slack_tunnel)
+
+    async def load_from_config(self, config: dict) -> bool:
+        self.target_id = config["target"]
+        self.tunnel_id = config["tunnel_id"]
+
+        self.type = "SlackTunnel"
+        self.token = config["token"]
+        self.app_token = config["app_token"]
+
+        return True
+
+    def get_cache_path(self) -> str:
+        today = datetime.datetime.today()
+        path = os.path.join(self.slack_cache, str(today.year), str(today.month))
+        if not os.path.exists(path):
+            os.makedirs(path)
+        return path
+
+    def post_message(self, msg: AgentMsg) -> None:
+        pass
+
+    async def start(self) -> bool:
+        app = AsyncApp(token=self.token)
+
+        bot_info = await app.client.auth_test()
+
+        @app.event("message")
+        async def _handle_message(event):
+            logger.info(json.dumps(event))
+            ty = event["type"]
+            if ty != "message":
+                return
+
+            user = event["user"]
+            user_info = await app.client.users_info(user=user)
+
+            if not user_info["ok"]:
+                return
+
+            user_info = user_info["user"]
+
+            mime_type = None
+            images = []
+            file_type = None
+            video_file = None
+            audio_file = None
+
+
+            if "files" in event:
+                files = event["files"]
+                if files is not None and len(files) > 0:
+                    for file in files:
+                        if file["mode"] == "tombstone":
+                            continue
+
+                        file_path = os.path.join(self.get_cache_path(), file["id"] + "." + file["filetype"])
+                        file_info = await app.client.files_info(file=file["id"])
+                        if not file_info["ok"]:
+                            continue
+                        await download_file(file_info["file"]["url_private_download"], file_path, self.token)
+
+                        if file["mimetype"].startswith("image/"):
+                            if file_type is None:
+                                file_type = "image"
+                            elif file_type != "image":
+                                break
+                            images.append(file_path)
+                        elif file["mimetype"].startswith("video/"):
+                            if file_type is None:
+                                file_type = "video"
+                            video_file = file_path
+                            break
+                        elif file["mimetype"].startswith("audio/"):
+                            if file_type is None:
+                                file_type = "audio"
+                            audio_file = file_path
+                            break
+
+            agent_msg = AgentMsg()
+            agent_msg.topic = "_slack"
+            agent_msg.msg_id = "discord_msg#" + event["client_msg_id"] + "#" + uuid.uuid4().hex
+            agent_msg.target = self.target_id
+            agent_msg.create_time = time.time()
+            agent_msg.sender = user_info["name"]
+            self.ai_bus.register_message_handler(agent_msg.sender, self._process_message)
+
+            content = re.sub("<@.+>", "", event["text"]).strip()
+
+            blocks = event["blocks"]
+            is_metion = False
+            for block in blocks:
+                if block["type"] == "rich_text":
+                    elements = block["elements"]
+                    for element in elements:
+                        if element["type"] == "rich_text_section":
+                            elements = element["elements"]
+
+                            for element in elements:
+                                if element["type"] == "user":
+                                    if element["user_id"] == bot_info.get("user_id"):
+                                        is_metion = True
+                                        break
+
+            if not is_metion:
+                return
+
+            if file_type is None:
+                agent_msg.body = content
+            elif file_type == "image":
+                agent_msg.body = agent_msg.create_image_body(images, content)
+                agent_msg.body_mime = mime_type
+            elif file_type == "video":
+                agent_msg.body = agent_msg.create_video_body(video_file, content)
+                agent_msg.body_mime = mime_type
+            elif file_type == "audio":
+                agent_msg.body = agent_msg.create_audio_body(audio_file, content)
+                agent_msg.body_mime = mime_type
+
+
+            resp_msg: AgentMsg = await self.ai_bus.send_message(agent_msg)
+            if resp_msg is None:
+                await app.client.chat_postMessage(channel=event["channel"], text=f"System Error: Timeout,{self.target_id}  no resopnse! Please check logs/aios.log for more details!")
+            else:
+                if resp_msg.body_mime is None:
+                    if resp_msg.body is None:
+                        return
+
+                    if len(resp_msg.body) < 1:
+                        await app.client.chat_postMessage(channel=event["channel"], text=f"")
+                        return
+
+                    knownledge_object = KnowledgeStore().parse_object_in_message(resp_msg.body)
+                    if knownledge_object is not None:
+                        if knownledge_object.get_object_type() == ObjectType.Image:
+                            image = KnowledgeStore().bytes_from_object(knownledge_object)
+                            try:
+                                async with aiofiles.open("image.jpg", "wb") as f:
+                                    await f.write(image)
+                                    await app.client.files_upload_v2(channel=event["channel"], file="image.jpg")
+                            except Exception as e:
+                                logger.error(f"save image error:{e}")
+                                logger.exception(e)
+                            return
+                    else:
+                        pos = resp_msg.body.find("audio file")
+                        if pos != -1:
+                            audio_file = resp_msg.body[pos+11:].strip()
+                            if audio_file.startswith("\""):
+                                audio_file = audio_file[1:-1]
+                            await app.client.files_upload_v2(channel=event["channel"], file=audio_file)
+                            return
+                    await app.client.chat_postMessage(channel=event["channel"], text=resp_msg.body)
+                else:
+                    if resp_msg.is_image_msg():
+                        text, images = resp_msg.get_image_body()
+                        file_uploads = []
+                        for image in images:
+                            file_uploads.append({"file": image})
+                        await app.client.files_upload_v2(channel=event["channel"], file_uploads=file_uploads, initial_comment=text)
+                    elif resp_msg.is_video_msg():
+                        text, video_file = resp_msg.get_video_body()
+                        await app.client.files_upload_v2(channel=event["channel"], file=video_file, initial_comment=text)
+                    elif resp_msg.is_audio_msg():
+                        text, audio_file = resp_msg.get_audio_body()
+                        await app.client.files_upload_v2(channel=event["channel"], file=audio_file, initial_comment=text)
+                    else:
+                        await app.client.chat_postMessage(channel=event["channel"], text=resp_msg.body)
+
+        handle = AsyncSocketModeHandler(app, self.app_token)
+        asyncio.create_task(handle.start_async())
+        self.app = handle
+        return True
+
+    async def close(self) -> None:
+        await self.app.close_async()
+        self.app = None
+
+    async def _process_message(self, msg: AgentMsg) -> None:
+        logger.warn(f"process message {msg.msg_id} from {msg.sender} to {msg.target}")

src/requirements.txt

@@ -154,3 +154,5 @@ html2text
 docx2txt
 opencv-python
 discord.py
+slack_bolt
+wget

src/service/aios_shell/aios_shell.py

@@ -19,6 +19,7 @@ from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
 from prompt_toolkit.completion import WordCompleter
 from prompt_toolkit.styles import Style
 
+from slack_tunnel import SlackTunnel
 directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../../')
 
@@ -236,6 +237,7 @@ class AIOS_Shell:
         TelegramTunnel.register_to_loader()
         EmailTunnel.register_to_loader()
         DiscordTunnel.register_to_loader()
+        SlackTunnel.register_to_loader()
         user_data_dir = str(AIStorage.get_instance().get_myai_dir())
         tunnels_config_path = os.path.abspath(f"{user_data_dir}/etc/tunnels.cfg.toml")
         tunnel_config = None
@@ -294,6 +296,10 @@ class AIOS_Shell:
             case "discord":
                 tunnel_config["type"] = "DiscordTunnel"
                 input_table["token"] = UserConfigItem("discord bot token\n You can get it from https://discord.com/developers/applications ,read https://discordpy.readthedocs.io/en/stable/discord.html for more details")
+            case "slack":
+                tunnel_config["type"] = "SlackTunnel"
+                input_table["token"] = UserConfigItem("slack bot token\n You can get it from https://api.slack.com/apps")
+                input_table["app_token"] = UserConfigItem("slack app token\n You can get it from https://api.slack.com/apps")
             case _:
                 error_text = FormattedText([("class:error", f"tunnel type {tunnel_type}not support!")])
                 print_formatted_text(error_text,style=shell_style)