wugren
240 lines
9.5 KiB
Python
240 lines
9.5 KiB
Python
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}")
|