[WIP]
This commit is contained in:
parent
af916a0f09
commit
03c8ea48c8
|
@ -11,6 +11,7 @@ from src.service.chat_purge_queue import ChatPurgeQueue
|
|||
from src.service.data_learner import DataLearner
|
||||
from src.service.reply_generator import ReplyGenerator
|
||||
from src.service.links_checker import LinksChecker
|
||||
from src.service.chance_manager import ChanceManager
|
||||
|
||||
|
||||
class Bot:
|
||||
|
@ -22,11 +23,13 @@ class Bot:
|
|||
logging.info("Bot started")
|
||||
|
||||
redis = Redis(config)
|
||||
chance_manager = ChanceManager(redis)
|
||||
|
||||
self.dispatcher.add_handler(MessageHandler(data_learner=DataLearner(),
|
||||
reply_generator=ReplyGenerator(),
|
||||
links_checker=LinksChecker(redis)))
|
||||
self.dispatcher.add_handler(CommandHandler())
|
||||
links_checker=LinksChecker(redis),
|
||||
chance_manager=chance_manager))
|
||||
self.dispatcher.add_handler(CommandHandler(chance_manager=chance_manager))
|
||||
self.dispatcher.add_handler(StatusHandler(chat_purge_queue=ChatPurgeQueue(self.updater.job_queue, redis)))
|
||||
|
||||
if config['updates']['mode'] == 'polling':
|
||||
|
|
|
@ -2,6 +2,7 @@ from abc import ABC
|
|||
|
||||
|
||||
class AbstractEntity(ABC):
|
||||
def __init__(self, chat, message):
|
||||
self.chat = chat
|
||||
def __init__(self, message):
|
||||
self.chat_id = message.chat.id
|
||||
self.chat_type = message.chat.type
|
||||
self.message = message
|
||||
|
|
|
@ -2,8 +2,8 @@ from .abstract_entity import AbstractEntity
|
|||
|
||||
|
||||
class Command(AbstractEntity):
|
||||
def __init__(self, chat, message):
|
||||
super(Command, self).__init__(chat=chat, message=message)
|
||||
def __init__(self, message):
|
||||
super(Command, self).__init__(message)
|
||||
self.name = Command.parse_name(message)
|
||||
self.args = Command.parse_args(message)
|
||||
|
||||
|
|
|
@ -7,8 +7,10 @@ from src.config import config
|
|||
|
||||
|
||||
class Message(AbstractEntity):
|
||||
def __init__(self, chat, message):
|
||||
super(Message, self).__init__(chat=chat, message=message)
|
||||
def __init__(self, chance, message):
|
||||
super(Message, self).__init__(message)
|
||||
|
||||
self.chance = chance
|
||||
|
||||
if self.has_text():
|
||||
self.text = message.text
|
||||
|
@ -63,7 +65,7 @@ class Message(AbstractEntity):
|
|||
def is_random_answer(self):
|
||||
"""Returns True if reply chance for this chat is high enough
|
||||
"""
|
||||
return random.randint(0, 100) < getattr(self.chat, 'random_chance', config['bot']['default_chance'])
|
||||
return random.randint(0, 100) < self.chance
|
||||
|
||||
def should_answer(self):
|
||||
return self.has_anchors() \
|
||||
|
|
|
@ -4,8 +4,8 @@ from src.config import config
|
|||
|
||||
|
||||
class Status(AbstractEntity):
|
||||
def __init__(self, chat, message):
|
||||
super(Status, self).__init__(chat=chat, message=message)
|
||||
def __init__(self, message):
|
||||
super(Status, self).__init__(message)
|
||||
|
||||
def is_bot_kicked(self):
|
||||
"""Returns True if the bot was kicked from group.
|
||||
|
|
|
@ -1,31 +0,0 @@
|
|||
import logging
|
||||
|
||||
from orator.orm import Model
|
||||
from orator.orm import has_many
|
||||
|
||||
import src.entity.pair
|
||||
|
||||
|
||||
class Chat(Model):
|
||||
__fillable__ = ['telegram_id', 'chat_type', 'random_chance']
|
||||
|
||||
@has_many
|
||||
def pairs(self):
|
||||
return src.entity.pair.Pair
|
||||
|
||||
# def migrate_to_chat_id(self, new_id):
|
||||
# logging.info("[Chat %s %s] Migrating ID to %s" % (self.chat_type, self.telegram_id, new_id))
|
||||
# self.telegram_id = new_id
|
||||
# self.save()
|
||||
|
||||
@staticmethod
|
||||
def get_chat(message):
|
||||
return Chat.first_or_create(telegram_id=message.chat.id, chat_type=message.chat.type)
|
||||
|
||||
# Events
|
||||
Chat.created(lambda chat: logging.info("[Chat %s %s] Created with internal ID #%s" %
|
||||
(chat.chat_type, chat.telegram_id, chat.id)))
|
||||
|
||||
#TODO Fix
|
||||
#Chat.updating(lambda chat: logging.info("[Chat %s %s New gab level is set to %s" %
|
||||
# (chat.chat_type, chat.telegram_id, chat.random_chance)))
|
|
@ -5,12 +5,12 @@ from telegram import Update
|
|||
from telegram.ext import Handler
|
||||
|
||||
from src.domain.command import Command
|
||||
from src.entity.chat import Chat
|
||||
from src.entity.pair import Pair
|
||||
from .moderate_command import ModerateCommand
|
||||
|
||||
|
||||
class CommandHandler(Handler):
|
||||
def __init__(self):
|
||||
def __init__(self, chance_manager):
|
||||
super(CommandHandler, self).__init__(self.handle)
|
||||
self.commands = {
|
||||
'start': self.__start_command,
|
||||
|
@ -21,6 +21,7 @@ class CommandHandler(Handler):
|
|||
'get_stats': self.__get_stats_command,
|
||||
'moderate': self.__moderate_command
|
||||
}
|
||||
self.chance_manager = chance_manager
|
||||
|
||||
def check_update(self, update):
|
||||
if isinstance(update, Update) and update.message:
|
||||
|
@ -36,8 +37,7 @@ class CommandHandler(Handler):
|
|||
return self.callback(dispatcher.bot, update, **optional_args)
|
||||
|
||||
def handle(self, bot, update):
|
||||
chat = Chat.get_chat(update.message)
|
||||
command = Command(chat=chat, message=update.message)
|
||||
command = Command(update.message)
|
||||
|
||||
try:
|
||||
self.commands[command.name](bot, command)
|
||||
|
@ -46,12 +46,12 @@ class CommandHandler(Handler):
|
|||
|
||||
def __reply(self, bot, command, message):
|
||||
logging.debug("[Chat %s %s command] %s: %s" %
|
||||
(command.chat.chat_type,
|
||||
command.chat.telegram_id,
|
||||
(command.chat_type,
|
||||
command.chat_id,
|
||||
command.name,
|
||||
message))
|
||||
|
||||
bot.send_message(chat_id=command.chat.telegram_id,
|
||||
bot.send_message(chat_id=command.chat_id,
|
||||
reply_to_message_id=command.message.message_id,
|
||||
text=message)
|
||||
|
||||
|
@ -89,22 +89,24 @@ In 12 hours, I'll forget everything that have been learned in your chat, so you
|
|||
|
||||
def __set_chance_command(self, bot, command):
|
||||
try:
|
||||
random_chance = int(command.args[0])
|
||||
chance = int(command.args[0])
|
||||
|
||||
if random_chance < 1 or random_chance > 50:
|
||||
if chance < 1 or chance > 50:
|
||||
raise ValueError
|
||||
|
||||
command.chat.update(random_chance=random_chance)
|
||||
self.chance_manager.set_chance(chat_id=command.chat_id, chance=chance)
|
||||
|
||||
self.__reply(bot, command, 'Set chance to: {}'.format(random_chance))
|
||||
self.__reply(bot, command, 'Set chance to: {}'.format(chance))
|
||||
except (IndexError, ValueError):
|
||||
self.__reply(bot, command, 'Usage: /set_chance 1-50.')
|
||||
|
||||
def __get_chance_command(self, bot, command):
|
||||
self.__reply(bot, command, 'Current chance: {}'.format(command.chat.random_chance))
|
||||
self.__reply(bot, command, 'Current chance: {}'
|
||||
.format(self.chance_manager.get_chance(command.chat_id)))
|
||||
|
||||
def __get_stats_command(self, bot, command):
|
||||
self.__reply(bot, command, 'Pairs: {}'.format(command.chat.pairs().count()))
|
||||
self.__reply(bot, command, 'Pairs: {}'
|
||||
.format(Pair.where('chat_id', command.chat_id).count()))
|
||||
|
||||
def __moderate_command(self, bot, command):
|
||||
try:
|
||||
|
|
|
@ -4,13 +4,11 @@ from random import choice
|
|||
from src.config import config
|
||||
from telegram.ext import MessageHandler as ParentHandler, Filters
|
||||
from telegram import ChatAction
|
||||
|
||||
from src.domain.message import Message
|
||||
from src.entity.chat import Chat
|
||||
|
||||
|
||||
class MessageHandler(ParentHandler):
|
||||
def __init__(self, data_learner, reply_generator, links_checker):
|
||||
def __init__(self, data_learner, reply_generator, links_checker, chance_manager):
|
||||
super(MessageHandler, self).__init__(
|
||||
Filters.text | Filters.sticker,
|
||||
self.handle)
|
||||
|
@ -18,15 +16,16 @@ class MessageHandler(ParentHandler):
|
|||
self.data_learner = data_learner
|
||||
self.reply_generator = reply_generator
|
||||
self.links_checker = links_checker
|
||||
self.chance_manager = chance_manager
|
||||
|
||||
def handle(self, bot, update):
|
||||
chat = Chat.get_chat(update.message)
|
||||
message = Message(chat=chat, message=update.message)
|
||||
chance = self.chance_manager.get_chance(update.message.chat.id)
|
||||
message = Message(chance=chance, message=update.message)
|
||||
|
||||
if message.has_text():
|
||||
logging.debug("[Chat %s %s bare_text] %s" %
|
||||
(message.chat.chat_type,
|
||||
message.chat.telegram_id,
|
||||
(message.chat_type,
|
||||
message.chat_id,
|
||||
message.text))
|
||||
|
||||
if message.has_text() and not message.is_editing():
|
||||
|
@ -38,12 +37,12 @@ class MessageHandler(ParentHandler):
|
|||
should_answer = message.should_answer()
|
||||
|
||||
if should_answer:
|
||||
bot.send_chat_action(chat_id=message.chat.telegram_id, action=ChatAction.TYPING)
|
||||
bot.send_chat_action(chat_id=message.chat_id, action=ChatAction.TYPING)
|
||||
|
||||
self.data_learner.learn(message)
|
||||
|
||||
if message.has_links() and self.links_checker.check(message.chat.telegram_id, message.links):
|
||||
bot.send_sticker(chat_id=message.chat.telegram_id,
|
||||
if message.has_links() and self.links_checker.check(message.chat_id, message.links):
|
||||
bot.send_sticker(chat_id=message.chat_id,
|
||||
reply_to_message_id=message.message.message_id,
|
||||
sticker=choice(config.getlist('links', 'stickers')))
|
||||
|
||||
|
@ -52,20 +51,20 @@ class MessageHandler(ParentHandler):
|
|||
reply_id = None if not message.is_reply_to_bot() else message.message.message_id
|
||||
|
||||
logging.debug("[Chat %s %s answer] %s" %
|
||||
(message.chat.chat_type,
|
||||
message.chat.telegram_id,
|
||||
(message.chat_type,
|
||||
message.chat_id,
|
||||
text))
|
||||
|
||||
bot.send_message(chat_id=message.chat.telegram_id,
|
||||
bot.send_message(chat_id=message.chat_id,
|
||||
reply_to_message_id=reply_id,
|
||||
text=text)
|
||||
|
||||
def __process_sticker(self, bot, message):
|
||||
if message.should_answer():
|
||||
logging.debug("[Chat %s %s spam_sticker]" %
|
||||
(message.chat.chat_type,
|
||||
message.chat.telegram_id))
|
||||
(message.chat_type,
|
||||
message.chat_id))
|
||||
|
||||
bot.send_sticker(chat_id=message.chat.telegram_id,
|
||||
bot.send_sticker(chat_id=message.chat_id,
|
||||
reply_to_message_id=message.message.message_id,
|
||||
sticker=choice(config.getlist('bot', 'spam_stickers')))
|
||||
|
|
|
@ -33,7 +33,7 @@ class ModerateCommand:
|
|||
return []
|
||||
|
||||
to_keep = []
|
||||
for pair in self.__find_pairs(self.entity.chat.id, list(found_words.keys())):
|
||||
for pair in self.__find_pairs(self.entity.chat_id, list(found_words.keys())):
|
||||
if pair.first_id in found_words:
|
||||
to_keep.append(pair.first_id)
|
||||
if pair.second_id in found_words:
|
||||
|
@ -47,13 +47,13 @@ class ModerateCommand:
|
|||
user_id = self.entity.message.from_user.id
|
||||
admin_ids = list(map(
|
||||
lambda member: member.user.id,
|
||||
self.bot.get_chat_administrators(chat_id=self.entity.chat.telegram_id)
|
||||
self.bot.get_chat_administrators(chat_id=self.entity.chat_id)
|
||||
))
|
||||
|
||||
return user_id in admin_ids
|
||||
|
||||
def remove_word(self, word_id):
|
||||
pairs_ids = self.__find_pairs(self.entity.chat.id, [word_id]).lists('id')
|
||||
pairs_ids = self.__find_pairs(self.entity.chat_id, [word_id]).lists('id')
|
||||
|
||||
Pair.where_in('id', pairs_ids).delete()
|
||||
Reply.where_in('pair_id', pairs_ids).delete()
|
||||
|
|
Loading…
Reference in New Issue