diff --git a/.gitignore b/.gitignore index 09bdba0..08c0ea6 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,7 @@ history.txt aios_shell_history.txt math_school_env.db workflows.db - +venv-linux +venv_test rootfs/test_doc/ diff --git a/src/component/slack_tunnel.py b/src/component/slack_tunnel.py new file mode 100644 index 0000000..fb9379d --- /dev/null +++ b/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}") diff --git a/src/requirements.txt b/src/requirements.txt index d1aae8f..ed743c4 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -154,3 +154,5 @@ html2text docx2txt opencv-python discord.py +slack_bolt +wget diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index cd32a59..8e68a7f 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/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)