Merge branch 'master' into redis
This commit is contained in:
commit
ef399285df
|
@ -7,6 +7,9 @@ from src.handler import *
|
|||
|
||||
|
||||
class Bot:
|
||||
"""
|
||||
Main initializer and dispatcher of messages
|
||||
"""
|
||||
def __init__(self):
|
||||
self.updater = Updater(token=config['bot']['token'])
|
||||
self.dispatcher = self.updater.dispatcher
|
||||
|
|
|
@ -2,13 +2,17 @@ from abc import ABC
|
|||
|
||||
|
||||
class AbstractEntity(ABC):
|
||||
"""
|
||||
Base class for all message entities
|
||||
"""
|
||||
|
||||
def __init__(self, message):
|
||||
self.chat_id = message.chat.id
|
||||
self.chat_type = message.chat.type
|
||||
self.message = message
|
||||
|
||||
def is_private(self):
|
||||
"""Returns True if the message is private.
|
||||
"""Returns True if chat type is private.
|
||||
"""
|
||||
return self.message.chat.type == 'private'
|
||||
|
||||
|
|
|
@ -2,6 +2,9 @@ from .abstract_entity import AbstractEntity
|
|||
|
||||
|
||||
class Command(AbstractEntity):
|
||||
"""
|
||||
Special class for message which contains command
|
||||
"""
|
||||
def __init__(self, message):
|
||||
super(Command, self).__init__(message)
|
||||
self.name = Command.parse_name(message)
|
||||
|
@ -9,8 +12,18 @@ class Command(AbstractEntity):
|
|||
|
||||
@staticmethod
|
||||
def parse_name(message):
|
||||
"""
|
||||
Parses command name from given message
|
||||
:param message: Telegram message object
|
||||
:return: Name of command
|
||||
"""
|
||||
return message.text[1:].split(' ')[0].split('@')[0]
|
||||
|
||||
@staticmethod
|
||||
def parse_args(message):
|
||||
"""
|
||||
Parses command args from given message
|
||||
:param message: Telegram message object
|
||||
:return: List of command args
|
||||
"""
|
||||
return message.text.split()[1:]
|
||||
|
|
|
@ -1,16 +1,19 @@
|
|||
import random
|
||||
import re
|
||||
from .abstract_entity import AbstractEntity
|
||||
from src.utils import deep_get_attr
|
||||
from src.config import config
|
||||
|
||||
|
||||
class Message(AbstractEntity):
|
||||
"""
|
||||
Basic message entity
|
||||
"""
|
||||
def __init__(self, chance, message):
|
||||
super(Message, self).__init__(message)
|
||||
|
||||
self.chance = chance
|
||||
self.entities = message.entities
|
||||
self.anchors = config.getlist('bot', 'anchors')
|
||||
|
||||
if self.has_text():
|
||||
self.text = message.text
|
||||
|
@ -18,39 +21,48 @@ class Message(AbstractEntity):
|
|||
self.text = ''
|
||||
|
||||
def has_text(self):
|
||||
"""Returns True if the message has text.
|
||||
"""
|
||||
Returns True if the message has text.
|
||||
"""
|
||||
return self.message.text.strip() != ''
|
||||
|
||||
def is_sticker(self):
|
||||
"""Returns True if the message is a sticker.
|
||||
"""
|
||||
Returns True if the message is a sticker.
|
||||
"""
|
||||
return self.message.sticker is not None
|
||||
|
||||
def has_entities(self):
|
||||
"""Returns True if the message has entities (attachments).
|
||||
"""
|
||||
Returns True if the message has entities (attachments).
|
||||
"""
|
||||
return self.entities is not None
|
||||
|
||||
def has_anchors(self):
|
||||
"""Returns True if the message contains at least one anchor from anchors config.
|
||||
"""
|
||||
anchors = config.getlist('bot', 'anchors')
|
||||
return self.has_text() and any(a in self.message.text.split(' ') for a in anchors)
|
||||
Returns True if the message contains at least one anchor from anchors config.
|
||||
"""
|
||||
return self.has_text() and any(a in self.message.text.split(' ') for a in self.anchors)
|
||||
|
||||
def is_reply_to_bot(self):
|
||||
"""Returns True if the message is a reply to bot.
|
||||
"""
|
||||
Returns True if the message is a reply to bot.
|
||||
"""
|
||||
user_name = deep_get_attr(self.message, 'reply_to_message.from_user.username')
|
||||
|
||||
return user_name == config['bot']['name']
|
||||
|
||||
def is_random_answer(self):
|
||||
"""Returns True if reply chance for this chat is high enough
|
||||
"""
|
||||
Returns True if reply chance for this chat is high enough
|
||||
"""
|
||||
return random.randint(0, 100) < self.chance
|
||||
|
||||
def should_answer(self):
|
||||
"""
|
||||
Returns True if bot should answer to this message
|
||||
:return: Should answer or not
|
||||
"""
|
||||
return self.has_anchors() \
|
||||
or self.is_private() \
|
||||
or self.is_reply_to_bot() \
|
||||
|
|
|
@ -4,20 +4,24 @@ from src.config import config
|
|||
|
||||
|
||||
class Status(AbstractEntity):
|
||||
bot_name = config['bot']['name']
|
||||
|
||||
"""
|
||||
Special class for message which contains info about status change
|
||||
"""
|
||||
def __init__(self, message):
|
||||
super(Status, self).__init__(message)
|
||||
self.bot_name = config['bot']['name']
|
||||
|
||||
def is_bot_kicked(self):
|
||||
"""Returns True if the bot was kicked from group.
|
||||
"""
|
||||
Returns True if the bot was kicked from group.
|
||||
"""
|
||||
user_name = deep_get_attr(self.message, 'left_chat_member.username')
|
||||
|
||||
return user_name == self.bot_name
|
||||
|
||||
def is_bot_added(self):
|
||||
"""Returns True if the bot was added to group.
|
||||
"""
|
||||
Returns True if the bot was added to group.
|
||||
"""
|
||||
user_name = deep_get_attr(self.message, 'new_chat_member.username')
|
||||
|
||||
|
|
|
@ -1,7 +1,7 @@
|
|||
import logging
|
||||
|
||||
from random import choice
|
||||
from src.config import config, data_learner, reply_generator, media_checker, chance_manager
|
||||
from src.config import config, data_learner, reply_generator, media_checker, chance_repository
|
||||
from telegram.ext import MessageHandler as ParentHandler, Filters
|
||||
from telegram import ChatAction
|
||||
from src.domain.message import Message
|
||||
|
@ -15,12 +15,12 @@ class MessageHandler(ParentHandler):
|
|||
self.data_learner = data_learner
|
||||
self.reply_generator = reply_generator
|
||||
self.media_checker = media_checker
|
||||
self.chance_manager = chance_manager
|
||||
self.chance_repository = chance_repository
|
||||
self.spam_stickers = config.getlist('bot', 'spam_stickers')
|
||||
self.media_checker_messages = config.getlist('media_checker', 'messages')
|
||||
|
||||
def handle(self, bot, update):
|
||||
chance = self.chance_manager.get_chance(update.message.chat.id)
|
||||
chance = self.chance_repository.get(update.message.chat.id)
|
||||
message = Message(chance=chance, message=update.message)
|
||||
|
||||
self.__check_media_uniqueness(bot, message)
|
||||
|
|
|
@ -2,6 +2,9 @@ import redis
|
|||
|
||||
|
||||
class Redis:
|
||||
"""
|
||||
Small redis wrapper, to simplify work with connection pool
|
||||
"""
|
||||
def __init__(self, config):
|
||||
self.pool = redis.ConnectionPool(host=config['redis']['host'],
|
||||
port=config.getint('redis', 'port'),
|
||||
|
|
|
@ -6,6 +6,9 @@ from src.config import config, trigram_repository, job_repository
|
|||
|
||||
|
||||
class ChatPurgeQueue:
|
||||
"""
|
||||
Scheduling and execution of chat purge
|
||||
"""
|
||||
def __init__(self):
|
||||
self.queue = None
|
||||
self.jobs = {}
|
||||
|
@ -21,6 +24,12 @@ class ChatPurgeQueue:
|
|||
return self
|
||||
|
||||
def add(self, chat_id, interval=None, db=True):
|
||||
"""
|
||||
Schedules purge of chat data
|
||||
:param chat_id: ID of chat
|
||||
:param interval: Interval in seconds
|
||||
:param db: Should be added to db or not
|
||||
"""
|
||||
interval = interval if interval is not None else self.default_interval
|
||||
scheduled_at = datetime.now() + timedelta(seconds=interval)
|
||||
|
||||
|
@ -35,6 +44,10 @@ class ChatPurgeQueue:
|
|||
self.job_repository.add(chat_id, scheduled_at)
|
||||
|
||||
def remove(self, chat_id):
|
||||
"""
|
||||
Removes scheduled purge job from queue
|
||||
:param chat_id: ID of chat
|
||||
"""
|
||||
if chat_id not in self.jobs:
|
||||
return
|
||||
|
||||
|
|
|
@ -7,6 +7,10 @@ class DataLearner:
|
|||
self.tokenizer = tokenizer
|
||||
|
||||
def learn(self, message):
|
||||
"""
|
||||
Split message to trigrams and write to DB
|
||||
:param message: Message
|
||||
"""
|
||||
words = self.tokenizer.extract_words(message)
|
||||
trigrams = self.tokenizer.split_to_trigrams(words)
|
||||
|
||||
|
|
|
@ -4,6 +4,9 @@ from urllib.parse import urlparse
|
|||
|
||||
|
||||
class MediaUniquenessChecker:
|
||||
"""
|
||||
Checks message links and photos for uniqueness
|
||||
"""
|
||||
def __init__(self):
|
||||
self.media_repository = media_repository
|
||||
|
||||
|
@ -19,7 +22,7 @@ class MediaUniquenessChecker:
|
|||
def __extract_media(self, message):
|
||||
media = []
|
||||
|
||||
for entity in filter(lambda e: e.type == 'url', message.message.entities):
|
||||
for entity in filter(lambda e: e.type == 'url', message.entities):
|
||||
link = self.__prettify(message.text[entity.offset:entity.length + entity.offset])
|
||||
media.append(link)
|
||||
|
||||
|
|
|
@ -3,6 +3,9 @@ from src.utils import strings_has_equal_letters, capitalize
|
|||
|
||||
|
||||
class ReplyGenerator:
|
||||
"""
|
||||
Handles generation of responses for user message
|
||||
"""
|
||||
def __init__(self):
|
||||
self.redis = redis
|
||||
self.tokenizer = tokenizer
|
||||
|
@ -16,6 +19,11 @@ class ReplyGenerator:
|
|||
self.end_sentence = config['grammar']['end_sentence']
|
||||
|
||||
def generate(self, message):
|
||||
"""
|
||||
Generates response based on message words
|
||||
:param message: Message
|
||||
:return: Response or empty string, if generated response equals to user message
|
||||
"""
|
||||
words = self.tokenizer.extract_words(message)
|
||||
pairs = [trigram[:-1] for trigram in self.tokenizer.split_to_trigrams(words)]
|
||||
messages = [self.__generate_best_message(chat_id=message.chat_id, pair=pair) for pair in pairs]
|
||||
|
|
Loading…
Reference in New Issue