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