Add discord tunnel

This commit is contained in:
wugren
2023-12-07 14:55:43 +08:00
parent 8739bf6a76
commit 3494aa8f6b
8 changed files with 313 additions and 27 deletions
+11 -1
View File
@@ -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
@@ -12,4 +22,4 @@ RUN pip install --no-cache-dir -r /opt/aios/requirements.txt
ENV PYTHONDONTWRITEBYTECODE 1
ENV PYTHONUNBUFFERED 1
CMD ["python3","./service/aios_shell/aios_shell.py"]
CMD ["python3","./service/aios_shell/aios_shell.py"]
+18 -6
View File
@@ -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:
msg.body = resp.result_str
msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{resp.result_str}"}]
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:
msg.body = resp.result_str
msg_prompt.messages = [{"role":"user","content":resp.result_str}]
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
+6 -6
View File
@@ -34,7 +34,7 @@ class AgentTunnel(ABC):
tunnel.connect_to(AIBus.get_default_bus(),tunnel.target_id)
await tunnel.start()
else:
logger.error(f"load tunnel {tunnel_id} failed")
logger.error(f"load tunnel {tunnel_id} failed")
else:
logger.error(f"load tunnel {tunnel_id} failed,loader not found")
@@ -49,12 +49,12 @@ class AgentTunnel(ABC):
await tunnel.start()
return True
else:
logger.error(f"load tunnel {tunnel_config['tunnel_id']} failed")
logger.error(f"load tunnel {tunnel_config['tunnel_id']} failed")
else:
logger.error(f"load tunnel {tunnel_config['type']} failed,loader not found")
return False
@classmethod
async def get_tunnel_by_agentid(cls,agent_id:str):
result = []
@@ -62,14 +62,14 @@ class AgentTunnel(ABC):
if tunnel.target_id == agent_id:
result.append(tunnel)
return result
def __init__(self) -> None:
super().__init__()
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:
@@ -87,7 +87,7 @@ class AgentTunnel(ABC):
@abstractmethod
def post_message(self, msg: AgentMsg) -> None:
pass
@abstractmethod
async def start(self) -> bool:
+25
View File
@@ -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
+215
View File
@@ -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}")
+28 -10
View File
@@ -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"):
photo_file = open(resp_msg.body,"rb")
if photo_file:
await update.message.reply_photo(resp_msg.body)
photo_file.close()
if resp_msg.is_image_msg():
text, images = resp_msg.get_image_body()
if text is not None:
await update.message.reply_text(text)
for image in images:
if os.path.exists(image):
await update.message.reply_photo(image)
else:
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(resp_msg.body)
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)
+1 -1
View File
@@ -153,4 +153,4 @@ oracledb
html2text
docx2txt
opencv-python
discord.py
+9 -3
View File
@@ -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)
@@ -423,7 +429,7 @@ class AIOS_Shell:
elif sub_cmd == "add":
if len(args) < 3:
return show_text
model_name = args[1]
url = args[2]
ComputeNodeConfig.get_instance().add_node("llama", url, model_name)
@@ -434,7 +440,7 @@ class AIOS_Shell:
elif sub_cmd == "rm":
if len(args) < 3:
return show_text
model_name = args[1]
url = args[2]
ComputeNodeConfig.get_instance().remove_node("llama", url, model_name)
@@ -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):