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}")