Add discord tunnel

Dockerfile

@@ -1,4 +1,14 @@
 FROM python:3.11
+
+ENV PYTHON_VERSION=3.11
+RUN apt-get update && apt-get install -y --no-install-recommends \
+		python${PYTHON_VERSION}-venv \
+    	python${PYTHON_VERSION}-dev \
+    	default-libmysqlclient-dev \
+    	build-essential \
+    	pkg-config \
+        && rm -rf /var/lib/apt/lists/* \
+
 WORKDIR /opt/aios
 COPY ./src /opt/aios
 COPY ./rootfs /opt/aios/app

src/aios/agent/agent.py

@@ -311,14 +311,20 @@ class AIAgent(BaseAIAgent):
                     content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames])
                     msg_prompt.messages = [{"role": "user", "content": content}]
             elif msg.is_audio_msg():
-                audio_file = msg.body
+                prompt, audio_file = msg.get_audio_body()
                 resp = await ComputeKernel.get_instance().do_speech_to_text(audio_file, None, prompt=None, response_format="text")
                 if resp.result_code != ComputeTaskResultCode.OK:
                     error_resp = msg.create_error_resp(resp.error_str)
                     return error_resp
                 else:
+                    if prompt is None or prompt == "":
+                        msg.body_mime = "text/plain"
                         msg.body = resp.result_str
                         msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{resp.result_str}"}]
+                    else:
+                        msg.body_mime = "text/plain"
+                        msg.body = f"{msg.sender} prompt:{prompt}\nasr response:{resp.result_str}"
+                        msg_prompt.messages = [{"role": "user", "content": msg.body}]
             else:
                 msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{msg.body}"}]
             session_topic = msg.target + "#" + msg.topic
@@ -352,14 +358,20 @@ class AIAgent(BaseAIAgent):
                     content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames])
                     msg_prompt.messages = [{"role": "user", "content": content}]
             elif msg.is_audio_msg():
-                audio_file = msg.body
+                prompt, audio_file = msg.get_audio_body()
                 resp = await (ComputeKernel.get_instance().do_speech_to_text(audio_file, None, prompt=None, response_format="text"))
                 if resp.result_code != ComputeTaskResultCode.OK:
                     error_resp = msg.create_error_resp(resp.error_str)
                     return error_resp
                 else:
+                    if prompt is None or prompt == "":
+                        msg.body_mime = "text/plain"
                         msg.body = resp.result_str
                         msg_prompt.messages = [{"role":"user","content":resp.result_str}]
+                    else:
+                        msg.body_mime = "text/plain"
+                        msg.body = f"user prompt:{prompt}\nasr response:{resp.result_str}"
+                        msg_prompt.messages = [{"role": "user", "content": msg.body}]
             else:
                 msg_prompt.messages = [{"role":"user","content":msg.body}]
             session_topic = msg.get_sender() + "#" + msg.topic

src/aios/frame/tunnel.py

@@ -69,7 +69,7 @@ class AgentTunnel(ABC):
         self.tunnel_id = None
         self.target_id = None
         self.target_type = None
-        self.ai_bus = None
+        self.ai_bus: AIBus = None
         self.is_connected = False
 
     def connect_to(self, ai_bus:AIBus,target_id: str) -> None:

src/aios/proto/agent_msg.py

@@ -164,6 +164,15 @@ class AgentMsg:
         body = json.loads(video_body)
         return body.get("prompt"), body.get("video")
 
+    @staticmethod
+    def create_audio_body(audio: str, prompt: str = None):
+        return json.dumps({"audio": audio, "prompt": prompt})
+
+    @staticmethod
+    def parse_audio_body(audio_body: str) -> Tuple[str, str]:
+        body = json.loads(audio_body)
+        return body.get("prompt"), body.get("audio")
+
     def set_image(self, sender: str, target: str, image_format: str, images: [str], prompt: str = None, topic: str = None):
         self.sender = sender
         self.target = target
@@ -210,6 +219,22 @@ class AgentMsg:
             return True
         return False
 
+    def set_audio(self, sender: str, target: str, audio_format: str, audio: str, prompt: str = None, topic: str = None):
+        self.sender = sender
+        self.target = target
+        self.create_time = time.time()
+        self.body_mime = f"audio/{audio_format}"
+        self.body = self.create_audio_body(audio, prompt)
+        if topic:
+            self.topic = topic
+
+    def get_audio_body(self) -> Tuple[str, str]:
+        if self.body_mime is None:
+            return None
+        if self.body_mime.startswith("audio/"):
+            return self.parse_audio_body(self.body)
+        return None
+
     def is_audio_msg(self) -> bool:
         if self.body_mime is None:
             return False

src/component/discord_tunnel.py

@@ -0,0 +1,215 @@
+import asyncio
+import datetime
+import time
+import logging
+import os
+import re
+import uuid
+import aiofiles
+from urllib.parse import urlparse
+from typing import Optional
+
+from aios import KnowledgeStore, ObjectType
+from aios.frame.tunnel import AgentTunnel
+from aios.proto.agent_msg import AgentMsg
+import discord
+
+from aios.storage.storage import AIStorage
+
+logger = logging.getLogger(__name__)
+
+
+IMAGE_FORMATS = ["jpg", "jpeg", "png", "gif", "bmp", "webp", "tiff", "tif"]
+VIDEO_FORMATS = ["mp4", "avi", "mov", "wmv", "flv", "mkv", "webm"]
+AUDIO_FORMATS = ["mp3", "wav", "ogg", "flac", "aac", "m4a", "wma", "ape", "alac", "opus", "oga"]
+
+
+class DiscordTunnel(AgentTunnel):
+    target_id: str
+    type: str
+    tunnel_id: str
+
+    def __init__(self, token: Optional[str] = None):
+        super().__init__()
+        self.token = token
+        self.client: discord.Client = None
+
+        self.discord_cache = os.path.join(AIStorage.get_instance().get_myai_dir(), "discord")
+        if not os.path.exists(self.discord_cache):
+            os.makedirs(self.discord_cache)
+
+    @classmethod
+    def register_to_loader(cls):
+        async def load_discord_tunnel(config: dict) -> AgentTunnel:
+            result_tunnel = DiscordTunnel()
+            if await result_tunnel.load_from_config(config):
+                return result_tunnel
+            else:
+                return None
+
+        AgentTunnel.register_loader("DiscordTunnel", load_discord_tunnel)
+
+    async def load_from_config(self, config: dict) -> bool:
+        self.type = "DiscordTunnel"
+        self.token = config["token"]
+        self.target_id = config["target"]
+        self.tunnel_id = config["tunnel_id"]
+        return True
+
+    def get_cache_path(self) -> str:
+        today = datetime.datetime.today()
+        path = os.path.join(self.discord_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:
+        if self.client is None:
+            logger.error("DiscordTunnel is not started")
+            return
+
+        logger.warn("post_message not implemented")
+
+    async def start(self) -> bool:
+        if self.client is not None:
+            logger.warn("DiscordTunnel is already started")
+            return False
+
+        if self.token is None:
+            self.token = os.environ.get("DISCORD_TOKEN")
+            if self.token is None:
+                raise ValueError("Discord token must be provided")
+
+        intents = discord.Intents.default()
+        intents.message_content = True
+        intents.members = True
+        client = discord.Client(intents=intents)
+
+        @client.event
+        async def on_ready():
+            print(f"Logged in as {self.client.user}")
+
+        @client.event
+        async def on_message(message: discord.Message):
+            logger.info(f"Message from {message.author}: {message.content}")
+            if message.author == self.client.user:
+                return
+
+            if len(message.channel.members) > 2:
+                if self.client.user not in message.mentions:
+                    return
+
+            content = re.sub("<@.+>", "", message.content).strip()
+
+            attach_type = None
+            images = []
+            ext = None
+            video_file = None
+            audio_file = None
+            if message.attachments is not None and len(message.attachments) > 0:
+                for attachment in message.attachments:
+                    logger.info(f"Attachment: {attachment.url}")
+                    url = urlparse(attachment.url)
+                    ext = url.path.rsplit(".")[-1]
+
+                    file_path = os.path.join(self.get_cache_path(), attachment.filename)
+                    with open(file_path, "wb") as f:
+                        await attachment.save(f)
+
+                    if ext in IMAGE_FORMATS:
+                        if attach_type is None:
+                            attach_type = "image"
+                        elif attach_type != "image":
+                            break
+                        images.append(file_path)
+                    elif ext in VIDEO_FORMATS:
+                        if attach_type is None:
+                            attach_type = "video"
+                        video_file = file_path
+                        break
+                    elif ext in AUDIO_FORMATS:
+                        if attach_type is None:
+                            attach_type = "audio"
+                        audio_file = file_path
+                        break
+
+            agent_msg = AgentMsg()
+            agent_msg.topic = "_discord"
+            agent_msg.msg_id = "discord_msg#" + str(message.id) + "#" + uuid.uuid4().hex
+            agent_msg.target = self.target_id
+            agent_msg.create_time = time.time()
+            agent_msg.sender = message.author.name
+            self.ai_bus.register_message_handler(agent_msg.sender, self._process_message)
+
+            if attach_type is None:
+                agent_msg.body = content
+            elif attach_type == "image":
+                agent_msg.body = agent_msg.create_image_body(images, content)
+                agent_msg.body_mime = f"image/{ext}"
+            elif attach_type == "video":
+                agent_msg.body = agent_msg.create_video_body(video_file, content)
+                agent_msg.body_mime = f"video/{ext}"
+            elif attach_type == "audio":
+                agent_msg.body = agent_msg.create_audio_body(audio_file, content)
+                agent_msg.body_mime = f"audio/{ext}"
+
+            resp_msg: AgentMsg = await self.ai_bus.send_message(agent_msg)
+            if resp_msg is None:
+                await message.channel.send(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 message.channel.send("")
+                        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 message.channel.send(file=discord.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 message.channel.send(file=discord.File(audio_file))
+                            return
+                    await message.channel.send(resp_msg.body)
+                else:
+                    if resp_msg.is_image_msg():
+                        text, images = resp_msg.get_image_body()
+                        files = []
+                        for image in images:
+                            files.append(discord.File(image))
+                        await message.channel.send(content=text, files=files)
+                    elif resp_msg.is_video_msg():
+                        text, video_file = resp_msg.get_video_body()
+                        await message.channel.send(text, file=discord.File(video_file))
+                    elif resp_msg.is_audio_msg():
+                        text, audio_file = resp_msg.get_audio_body()
+                        await message.channel.send(text, file=discord.File(audio_file))
+                    else:
+                        await message.channel.send(resp_msg.body)
+
+        asyncio.create_task(client.start(self.token))
+        self.client = client
+        print("start finish")
+
+    async def close(self) -> None:
+        await self.client.close()
+        self.client = None
+
+    async def _process_message(self, msg: AgentMsg) -> None:
+        logger.warn(f"process message {msg.msg_id} from {msg.sender} to {msg.target}")

src/component/tg_tunnel.py

@@ -192,14 +192,14 @@ class TelegramTunnel(AgentTunnel):
             ext = audio_file.file_path.rsplit(".")[-1]
             file_path = os.path.join(self.get_cache_path(), audio_file.file_id + f".{ext}")
             await audio_file.download_to_drive(file_path)
-            agent_msg.body = file_path
+            agent_msg.body = agent_msg.create_audio_body(file_path, message.caption)
             agent_msg.body_mime = f"audio/{ext}"
         elif message.voice is not None:
             audio_file = await message.voice.get_file()
             ext = audio_file.file_path.rsplit(".")[-1]
             file_path = os.path.join(self.get_cache_path(), audio_file.file_id + f".{ext}")
             await audio_file.download_to_drive(file_path)
-            agent_msg.body = file_path
+            agent_msg.body = agent_msg.create_audio_body(file_path, message.caption)
             agent_msg.body_mime = f"audio/{ext}"
 
         agent_msg.create_time = time.time()
@@ -304,10 +304,10 @@ class TelegramTunnel(AgentTunnel):
         logger.info(f"process message {agent_msg.msg_id} from {agent_msg.sender} to {agent_msg.target}")
         if agent_msg.msg_type == AgentMsgType.TYPE_GROUPMSG:
             self.ai_bus.register_message_handler(agent_msg.target, self._process_message)
-            resp_msg = await self.ai_bus.send_message(agent_msg,self.target_id,agent_msg.target)
+            resp_msg: AgentMsg = await self.ai_bus.send_message(agent_msg,self.target_id,agent_msg.target)
         else:
             #self.ai_bus.register_message_handler(reomte_user_name, self._process_message)
-            resp_msg = await self.ai_bus.send_message(agent_msg)
+            resp_msg: AgentMsg = await self.ai_bus.send_message(agent_msg)
         #await bot.send_chat_action(chat_id=update.effective_chat.id, action="typing")
 
 
@@ -345,14 +345,32 @@ class TelegramTunnel(AgentTunnel):
                         return
                 await update.message.reply_text(resp_msg.body)
             else:
-                if resp_msg.body_mime.startswith("image"):
+                if resp_msg.is_image_msg():
-                    photo_file = open(resp_msg.body,"rb")
+                    text, images = resp_msg.get_image_body()
-                    if photo_file:
+                    if text is not None:
-                        await update.message.reply_photo(resp_msg.body)
+                        await update.message.reply_text(text)
-                        photo_file.close()
+                    for image in images:
+                        if os.path.exists(image):
+                            await update.message.reply_photo(image)
                         else:
-                        await update.message.reply_text(resp_msg.body)
+                            await update.message.reply_text(image)
+                elif resp_msg.is_video_msg():
+                    text, video_file = resp_msg.get_video_body()
+                    if text is not None:
+                        await update.message.reply_text(text)
+                    if os.path.exists(video_file):
+                        await update.message.reply_video(video_file)
+                    else:
+                        await update.message.reply_text(video_file)
+                elif resp_msg.is_audio_msg():
+                    text, audio_file = resp_msg.get_audio_body()
+                    if text is not None:
+                        await update.message.reply_text(text)
 
+                    if os.path.exists(audio_file):
+                        await update.message.reply_voice(audio_file)
+                    else:
+                        await update.message.reply_text(audio_file)
                 else:
                     await update.message.reply_text(resp_msg.body)
 

src/requirements.txt

@@ -153,4 +153,4 @@ oracledb
 html2text
 docx2txt
 opencv-python
-
+discord.py

src/service/aios_shell/aios_shell.py

@@ -19,6 +19,8 @@ from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
 from prompt_toolkit.completion import WordCompleter
 from prompt_toolkit.styles import Style
 
+from discord_tunnel import DiscordTunnel
+
 directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../../')
 
@@ -230,6 +232,7 @@ class AIOS_Shell:
 
         TelegramTunnel.register_to_loader()
         EmailTunnel.register_to_loader()
+        DiscordTunnel.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
@@ -285,6 +288,9 @@ class AIOS_Shell:
                 input_table["smtp"] = UserConfigItem("smtp server address,like hostname:port")
                 input_table["user"] = UserConfigItem("mail server login user name")
                 input_table["password"] = UserConfigItem("main server login password")
+            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 _:
                 error_text = FormattedText([("class:error", f"tunnel type {tunnel_type}not support!")])
                 print_formatted_text(error_text,style=shell_style)
@@ -535,7 +541,7 @@ class AIOS_Shell:
                     tunnel_type = "telegram"
                 else:
                     tunnel_type = args[1]
-
+                tunnel_type = tunnel_type.lower()
                 tunnel_config = await self.get_tunnel_config_from_input(tunnel_target,tunnel_type)
                 if tunnel_config:
                     if await AgentTunnel.load_tunnel_from_config(tunnel_config):