opendan/src/aios_kernel/tg_tunnel.py

import logging
import threading
import asyncio
import uuid
import time
import aiofiles

from telegram import Update
from telegram import Bot
from telegram.ext import Updater
from telegram.error import Forbidden, NetworkError

from knowledge.object.object_id import ObjectType

from .knowledge_base import KnowledgeBase

from .tunnel import AgentTunnel
from .storage import AIStorage
from .contact_manager import ContactManager,Contact,FamilyMember
from .agent_message import AgentMsg


logger = logging.getLogger(__name__)

class TelegramTunnel(AgentTunnel):

    @classmethod
    def register_to_loader(cls):
        async def load_tg_tunnel(config:dict) -> AgentTunnel:
            result_tunnel = TelegramTunnel("")
            if await result_tunnel.load_from_config(config):
                return result_tunnel
            else:
                return None

        AgentTunnel.register_loader("TelegramTunnel",load_tg_tunnel)


    async def load_from_config(self,config:dict)->bool:
        self.type = "TelegramTunnel"
        self.tg_token = config["token"]
        self.target_id = config["target"]
        self.tunnel_id = config["tunnel_id"]
        if config.get("allow") is not None:
            self.allow_group = config["allow"]

        return True

    def dump_to_config(self) -> dict:
        pass

    def __init__(self,tg_token:str) -> None:
        super().__init__()
        self.is_start = False
        self.tg_token = tg_token
        self.bot:Bot = None
        self.update_queue = None
        self.allow_group = "contact"
        self.in_process_tg_msg = {}

    async def _do_process_raw_message(self,bot: Bot, update_id: int) -> int:
        # Request updates after the last update_id
        updates = await bot.get_updates(offset=update_id, timeout=10, allowed_updates=Update.ALL_TYPES)
        for update in updates:
            next_update_id = update.update_id + 1

            if update.message and update.message.text:

                await self.on_message(bot,update)
            return next_update_id

        return update_id

    async def start(self) -> bool:
        if self.is_start:
            logger.warning(f"tunnel {self.tunnel_id} is already started")
            return False
        self.is_start = True
        logger.info(f"tunnel {self.tunnel_id} is starting...")

        self.bot = Bot(self.tg_token)
        self.update_queue = asyncio.Queue()
        self.bot_updater = Updater(self.bot,update_queue=self.update_queue)


        async def _run_app():
            try:
                update_id = (await self.bot.get_updates())[0].update_id
            except IndexError:
                update_id = None

            #logger.info("listening for new messages...")
            while True:
                try:
                    update_id = await self._do_process_raw_message(self.bot, update_id)
                except NetworkError:
                    await asyncio.sleep(1)
                except Forbidden:
                    # The user has removed or blocked the bot.
                    update_id += 1
                except  Exception as e:
                    logger.error(f"tg_tunnel error:{e}")
                    await asyncio.sleep(1)

                    

        asyncio.create_task(_run_app())
        logger.info(f"tunnel {self.tunnel_id} started.")
        return True

    async def close(self) -> None:
        pass

    async def _process_message(self, msg: AgentMsg) -> None:
        logger.warn(f"process message {msg.msg_id} from {msg.sender} to {msg.target}")

    async def conver_tg_msg_to_agent_msg(self,update:Update) -> AgentMsg:
        agent_msg = AgentMsg()
        agent_msg.topic = "_telegram"
        agent_msg.msg_id = "tg_msg#" + str(update.message.message_id) + "#" + uuid.uuid4().hex
        agent_msg.target = self.target_id
        agent_msg.body = update.message.text
        agent_msg.create_time = time.time()
        #if update.message.photo is not None:
        #    agent_msg.body_mime = "image"
        #    agent_msg.body = update.message.photo[-1].get_file().download()
        return agent_msg



    async def on_message(self, bot:Bot, update: Update) -> None:
        if update.effective_user.is_bot:
            logger.warning(f"ignore message from telegram bot {update.effective_user.id}")
            return None
        
        if self.in_process_tg_msg.get(update.message.message_id) is not None:
            logger.warning(f"ignore message from telegram bot {update.effective_user.id}")
            return None
        
        self.in_process_tg_msg[update.message.message_id] = True


        cm : ContactManager = ContactManager.get_instance()
        reomte_user_name = f"{update.effective_user.id}@telegram"
        #owner_tg_username = AIStorage.get_instance().get_user_config().get_value("telegram")
        #owner_name = AIStorage.get_instance().get_user_config().get_value("username")


        contact : Contact = cm.find_contact_by_telegram(update.effective_user.username)
        if contact is None:
            contact = cm.find_contact_by_telegram(str(update.effective_user.id))

        if contact is not None:
            reomte_user_name = contact.name
            if not contact.is_family_member:
                if self.allow_group != "contact" and self.allow_group !="guest":
                    await update.message.reply_text(f"You're not supposed to talk to me! Please contact my father~")
                    return
            
        else:
            if self.allow_group != "guest":
                await update.message.reply_text(f"You're not supposed to talk to me! Please contact my father~")
                return

            if cm.is_auto_create_contact_from_telegram:
                contact_name = update.effective_user.first_name
                if update.effective_user.last_name is not None:
                    contact_name += " " + update.effective_user.last_name

                contact = Contact(contact_name)
                contact.telegram = update.effective_user.username if update.effective_user.username is not None else str(update.effective_user.id)
                contact.added_by = self.target_id
                cm.add_contact(contact.name, contact)
                reomte_user_name = contact.name
             
        agent_msg = await self.conver_tg_msg_to_agent_msg(update)
        agent_msg.sender = reomte_user_name
        self.ai_bus.register_message_handler(reomte_user_name, self._process_message)
        #await bot.send_chat_action(chat_id=update.effective_chat.id, action="typing")
        resp_msg = await self.ai_bus.send_message(agent_msg)
        logger.info(f"process message {agent_msg.msg_id} from {agent_msg.sender} to {agent_msg.target}")
        if resp_msg is None:
            await update.message.reply_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 not None:
                    knowledge_object = KnowledgeBase().parse_object_in_message(resp_msg.body)
                    if knowledge_object is not None:
                        if knowledge_object.get_object_type() == ObjectType.Image:
                            image = KnowledgeBase().bytes_from_object(knowledge_object)
                            try:
                                async with aiofiles.open("tg_send_temp.png", mode='wb') as local_file:
                                    if local_file:
                                        await local_file.write(image)
                                        await update.message.reply_photo("tg_send_temp.png")
                            except Exception as e:
                                logger.error(f"save image error: {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 update.message.reply_voice(audio_file)
                            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()
                    else:
                        await update.message.reply_text(resp_msg.body)

                else:
                    await update.message.reply_text(resp_msg.body)