upo-app-agents/src/app/interface/telegram_app.py

import io
import os
import json
import httpx
import logging
import warnings
from enum import Enum
from markdown_pdf import MarkdownPdf, Section
from telegram import CallbackQuery, InlineKeyboardButton, InlineKeyboardMarkup, Message, Update, User
from telegram.constants import ChatAction
from telegram.ext import Application, CallbackQueryHandler, CommandHandler, ContextTypes, ConversationHandler, MessageHandler, filters
from app.agents.pipeline import Pipeline, PipelineInputs

# per per_message di ConversationHandler che rompe sempre qualunque input tu metta
warnings.filterwarnings("ignore")
logging = logging.getLogger("telegram")


# Lo stato cambia in base al valore di ritorno delle funzioni async
# END state è già definito in telegram.ext.ConversationHandler
# Un semplice schema delle interazioni:
#      /start
#         ║
#         V
#  ╔══ CONFIGS <═════╗
#  ║      ║ ╚══> SELECT_CONFIG
#  ║      V
#  ║  start_team (polling for updates)
#  ║      ║
#  ║      V
#  ╚═══> END
CONFIGS, SELECT_CONFIG = range(2)

# Usato per separare la query arrivata da Telegram
QUERY_SEP = "|==|"

class ConfigsChat(Enum):
    MODEL_TEAM = "Team Model"
    MODEL_OUTPUT = "Output Model"
    STRATEGY = "Strategy"

class TelegramApp:
    def __init__(self):
        token = os.getenv("TELEGRAM_BOT_TOKEN")
        assert token, "TELEGRAM_BOT_TOKEN environment variable not set"

        self.user_requests: dict[User, PipelineInputs] = {}
        self.token = token
        self.create_bot()

    def add_miniapp_url(self, url: str) -> None:
        try:
            endpoint = f"https://api.telegram.org/bot{self.token}/setChatMenuButton"
            payload = {"menu_button": json.dumps({
                "type": "web_app",
                "text": "MiniApp",
                "web_app": { "url": url }
            })}
            httpx.post(endpoint, data=payload)
        except httpx.HTTPError as e:
            logging.warning(f"Failed to update mini app URL: {e}")

    def create_bot(self) -> None:
        """
        Initialize the Telegram bot and set up the conversation handler.
        """
        app = Application.builder().token(self.token).build()

        app.add_error_handler(self.__error_handler)
        app.add_handler(ConversationHandler(
            per_message=False, # capire a cosa serve perchè da un warning quando parte il server
            entry_points=[CommandHandler('start', self.__start)],
            states={
                CONFIGS: [
                    CallbackQueryHandler(self.__model_team, pattern=ConfigsChat.MODEL_TEAM.name),
                    CallbackQueryHandler(self.__model_output, pattern=ConfigsChat.MODEL_OUTPUT.name),
                    CallbackQueryHandler(self.__strategy, pattern=ConfigsChat.STRATEGY.name),
                    CallbackQueryHandler(self.__cancel, pattern='^cancel$'),
                    MessageHandler(filters.TEXT, self.__start_team)  # Any text message
                ],
                SELECT_CONFIG: [
                    CallbackQueryHandler(self.__select_config, pattern=f"^__select_config{QUERY_SEP}.*$"),
                ]
            },
            fallbacks=[CommandHandler('start', self.__start)],
        ))
        self.app = app

    def run(self) -> None:
        self.app.run_polling()

    ########################################
    # Funzioni di utilità
    ########################################
    async def start_message(self, user: User, query: CallbackQuery | Message) -> None:
        confs = self.user_requests.setdefault(user, PipelineInputs())

        str_model_team = f"{ConfigsChat.MODEL_TEAM.value}:   {confs.team_model.label}"
        str_model_output = f"{ConfigsChat.MODEL_OUTPUT.value}:   {confs.team_leader_model.label}"
        str_strategy = f"{ConfigsChat.STRATEGY.value}:   {confs.strategy.label}"

        msg, keyboard = (
            "Please choose an option or write your query",
            InlineKeyboardMarkup([
                [InlineKeyboardButton(str_model_team, callback_data=ConfigsChat.MODEL_TEAM.name)],
                [InlineKeyboardButton(str_model_output, callback_data=ConfigsChat.MODEL_OUTPUT.name)],
                [InlineKeyboardButton(str_strategy, callback_data=ConfigsChat.STRATEGY.name)],
                [InlineKeyboardButton("Cancel", callback_data='cancel')]
            ])
        )

        if isinstance(query, CallbackQuery):
            await query.edit_message_text(msg, reply_markup=keyboard, parse_mode='MarkdownV2')
        else:
            await query.reply_text(msg, reply_markup=keyboard, parse_mode='MarkdownV2')

    async def handle_callbackquery(self, update: Update) -> tuple[CallbackQuery, User]:
        assert update.callback_query and update.callback_query.from_user, "Update callback_query or user is None"
        query = update.callback_query
        await query.answer()  # Acknowledge the callback query
        return query, query.from_user

    async def handle_message(self, update: Update) -> tuple[Message, User]:
        assert update.message and update.message.from_user, "Update message or user is None"
        return update.message, update.message.from_user

    def build_callback_data(self, callback: str, config: ConfigsChat, labels: list[str]) -> list[tuple[str, str]]:
        return [(label, QUERY_SEP.join((callback, config.value, str(i)))) for i, label in enumerate(labels)]

    async def __error_handler(self, update: object, context: ContextTypes.DEFAULT_TYPE) -> None:
        try:
            logging.error(f"Unhandled exception in Telegram handler: {context.error}")

            # Try to notify the user in chat if possible
            if isinstance(update, Update) and update.effective_chat:
                chat_id = update.effective_chat.id
                msg = "An error occurred while processing your request."
                await context.bot.send_message(chat_id=chat_id, text=msg)

        except Exception:
            # Ensure we never raise from the error handler itself
            logging.exception("Exception in the error handler")

    #########################################
    # Funzioni async per i comandi e messaggi
    #########################################
    async def __start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
        message, user = await self.handle_message(update)
        logging.info(f"@{user.username} started the conversation.")
        await self.start_message(user, message)
        return CONFIGS

    async def __model_team(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
        return await self._model_select(update, ConfigsChat.MODEL_TEAM)

    async def __model_output(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
        return await self._model_select(update, ConfigsChat.MODEL_OUTPUT)

    async def _model_select(self, update: Update, state: ConfigsChat, msg: str | None = None) -> int:
        query, user = await self.handle_callbackquery(update)

        req = self.user_requests[user]
        models = self.build_callback_data("__select_config", state, req.list_models_names())
        inline_btns = [[InlineKeyboardButton(name, callback_data=callback_data)] for name, callback_data in models]

        await query.edit_message_text(msg or state.value, reply_markup=InlineKeyboardMarkup(inline_btns))
        return SELECT_CONFIG

    async def __strategy(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
        query, user = await self.handle_callbackquery(update)

        req = self.user_requests[user]
        strategies = self.build_callback_data("__select_config", ConfigsChat.STRATEGY, req.list_strategies_names())
        inline_btns = [[InlineKeyboardButton(name, callback_data=callback_data)] for name, callback_data in strategies]

        await query.edit_message_text("Select a strategy", reply_markup=InlineKeyboardMarkup(inline_btns))
        return SELECT_CONFIG

    async def __select_config(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
        query, user = await self.handle_callbackquery(update)
        logging.debug(f"@{user.username} --> {query.data}")

        req = self.user_requests[user]
        _, state, index = str(query.data).split(QUERY_SEP)
        if state == str(ConfigsChat.MODEL_TEAM):
            req.choose_team(int(index))
        if state == str(ConfigsChat.MODEL_OUTPUT):
            req.choose_team_leader(int(index))
        if state == str(ConfigsChat.STRATEGY):
            req.choose_strategy(int(index))

        await self.start_message(user, query)
        return CONFIGS

    async def __start_team(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
        message, user = await self.handle_message(update)

        confs = self.user_requests[user]
        confs.user_query = message.text or ""

        logging.info(f"@{user.username} started the team with [{confs.team_model.label}, {confs.team_leader_model.label}, {confs.strategy.label}]")
        await self.__run_team(update, confs)

        logging.info(f"@{user.username} team finished.")
        return ConversationHandler.END

    async def __cancel(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
        query, user = await self.handle_callbackquery(update)
        logging.info(f"@{user.username} canceled the conversation.")
        if user in self.user_requests:
            del self.user_requests[user]
        await query.edit_message_text("Conversation canceled. Use /start to begin again.")
        return ConversationHandler.END

    async def __run_team(self, update: Update, inputs: PipelineInputs) -> None:
        if not update.message: return

        bot = update.get_bot()
        msg_id = update.message.message_id - 1
        chat_id = update.message.chat_id

        configs_str = [
            'Running with configurations:   ',
            f'Team:      {inputs.team_model.label}',
            f'Output:    {inputs.team_leader_model.label}',
            f'Strategy:  {inputs.strategy.label}',
            f'Query:     "{inputs.user_query}"'
        ]
        full_message = f"""```\n{'\n'.join(configs_str)}\n```\n\n"""
        first_message = full_message + "Generating report, please wait"
        msg = await bot.edit_message_text(chat_id=chat_id, message_id=msg_id, text=first_message, parse_mode='MarkdownV2')
        if isinstance(msg, bool): return

        # Remove user query and bot message
        await bot.delete_message(chat_id=chat_id, message_id=update.message.id)

        # TODO migliorare messaggi di attesa
        await bot.send_chat_action(chat_id=chat_id, action=ChatAction.TYPING)
        pipeline = Pipeline(inputs)
        report_content = await pipeline.interact_async()
        await msg.delete()

        # attach report file to the message
        pdf = MarkdownPdf(toc_level=2, optimize=True)
        pdf.add_section(Section(report_content, toc=False))

        # TODO vedere se ha senso dare il pdf o solo il messaggio
        document = io.BytesIO()
        pdf.save_bytes(document)
        document.seek(0)
        await bot.send_document(chat_id=chat_id, document=document, filename="report.pdf", parse_mode='MarkdownV2', caption=full_message)