vanity-bot/vanitybot.py

import re
import logging
import asyncio
from aiogram.contrib.fsm_storage.memory import MemoryStorage
from aiogram import Bot, Dispatcher, executor, types
from aiogram_dialog import DialogManager, DialogRegistry, StartMode
from config import API_TOKEN
from dialog import dialog, MySG
from progress import bg_dialog
from summarize import get_paper_desc

# Initialize bot and dispatcher
bot = Bot(token=API_TOKEN) # type: ignore
dp = Dispatcher(bot, storage=MemoryStorage())
registry = DialogRegistry(dp)
registry.register(dialog)
registry.register(bg_dialog)

help_message = "Hello!\n\n\
Send me a link paper from arxiv.org and \
I'll send you back snippet of paper and arxiv-vanity.com mobile friendly link!\n\
Or add me to chat and I'll be watching the arxiv link and \
reply to them with fancy arxiv-vanity links."

@dp.message_handler(commands=['start'])
async def process_start_command(message: types.Message):
    await message.reply(help_message)


@dp.message_handler(commands=['help'])
async def process_help_command(message: types.Message):
    await message.reply(help_message)
    
@dp.message_handler(commands=['long'])
async def long(message: types.Message):
    import random
    long = "".join(str(random.randint(1,10)) for _ in range(3700))
    await message.reply(long)


@dp.message_handler(regexp=r'arxiv.org\/(?:abs|pdf)\/\d{4}\.\d{5}')
async def vanitify(message: types.Message, dialog_manager: DialogManager):
    papers_ids = re.findall(r'arxiv.org\/(?:abs|pdf)\/(\d{4}\.\d{5}[v]?[\d]?)', message.text)
    
    async def start_dialog(manager=dialog_manager.bg(), state=MySG.main, mode=StartMode.NEW_STACK, data={}):
        await manager.start(state=state, mode=mode, data=data)

    async def get_paper_abs(id_):
        reply_message = f"[Here you can read the paper in mobile friendly way](https://www.arxiv-vanity.com/papers/{id_})"
        data = {
            "id": id_,
            "reply_message": reply_message,
            "url": None,
            "title": None,
            "abs": None
        }
        if paper := await get_paper_desc(id_):
            id_, url, title, abstract, authors = paper.values()
            reply_message = f'{url}\n\n***{title}***\n\n{abstract}\n\n{reply_message}'
            data.update({
                "id": id_,
                "reply_message": reply_message,
                "url": url,
                "title": title,
                "abs": abstract,
                "authors": authors
                })
        else:
            reply_message = f'Something went wrong. Can not reach arxiv.com :('
            data["reply_message"] = reply_message
        return data

    list_data = await asyncio.gather(*[get_paper_abs(id_) for id_ in papers_ids])
    asyncio.gather(*[start_dialog(data=data) for data in list_data])


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    logging.getLogger("asyncio").setLevel(logging.DEBUG)
    logging.getLogger("aiogram_dialog").setLevel(logging.DEBUG)
    executor.start_polling(dp, skip_updates=True)