vanitybot.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. import re
  2. import logging
  3. from aiogram.contrib.fsm_storage.memory import MemoryStorage
  4. from aiogram import Bot, Dispatcher, executor, types
  5. from aiogram_dialog import DialogManager, DialogRegistry, StartMode
  6. from config import API_TOKEN
  7. from dialog import dialog, MySG
  8. from summarize import get_paper_desc
  9. # Initialize bot and dispatcher
  10. bot = Bot(token=API_TOKEN)
  11. dp = Dispatcher(bot, storage=MemoryStorage())
  12. registry = DialogRegistry(dp)
  13. registry.register(dialog)
  14. help_message = "Hello!\n\n\
  15. Send me a link paper from arxiv.org and \
  16. I'll send you back snippet of paper and arxiv-vanity.com mobile friendly link!\n\
  17. Or add me to chat and I'll be watching the arxiv link and \
  18. reply to them with fancy arxiv-vanity links."
  19. @dp.message_handler(commands=['start'])
  20. async def process_start_command(message: types.Message):
  21. await message.reply(help_message)
  22. @dp.message_handler(commands=['help'])
  23. async def process_help_command(message: types.Message):
  24. await message.reply(help_message)
  25. @dp.message_handler(regexp=r'arxiv.org\/(?:abs|pdf)\/\d{4}\.\d{5}')
  26. async def vanitify(message: types.Message, dialog_manager: DialogManager):
  27. papers_ids = re.findall(r'arxiv.org\/(?:abs|pdf)\/(\d{4}\.\d{5})', message.text)
  28. for id_ in papers_ids:
  29. reply_message = f"[Here you can read the paper in mobile friendly way](https://www.arxiv-vanity.com/papers/{id_})"
  30. data = {
  31. "id": id_,
  32. "reply_message": reply_message,
  33. "url": None,
  34. "title": None,
  35. "abs": None
  36. }
  37. if desc := get_paper_desc(id_):
  38. url, title, description = desc
  39. reply_message = f'{url}\n\n***{title}***\n\n{description}\n\n{reply_message}'
  40. data.update({
  41. "reply_message": reply_message,
  42. "url": url,
  43. "title": title,
  44. "abs": description
  45. })
  46. else:
  47. reply_message = f'Something went wrong. Can not reach arxiv.com :('
  48. data["reply_message"] = reply_message
  49. await dialog_manager.start(MySG.main, mode=StartMode.NEW_STACK, data=data)
  50. if __name__ == "__main__":
  51. logging.basicConfig(level=logging.DEBUG)
  52. executor.start_polling(dp, skip_updates=True)