diff --git a/df_utils.py b/df_utils.py index 173069e..51b3759 100644 --- a/df_utils.py +++ b/df_utils.py @@ -2,7 +2,8 @@ from google.cloud import dialogflow_v2beta1 as dialogflow from dotenv import load_dotenv -load_dotenv('.env') + +load_dotenv() async def detect_intent_text(project_id, session_id, text, language_code): @@ -16,5 +17,5 @@ async def detect_intent_text(project_id, session_id, text, language_code): if response.query_result.intent.is_fallback: return None - else: - return response.query_result.fulfillment_text + + return response.query_result.fulfillment_text diff --git a/docker-compose.yml b/docker-compose.yml index 0be6ab7..522b645 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -12,7 +12,6 @@ services: - SESSION_NAME=NoVagonBot - GOOGLE_PROJECT_ID=novagon-rg9f - GOOGLE_APPLICATION_CREDENTIALS=df_key.json - - GOOGLE_MODEL_LANGUAGE=ru - POSTGRES_HOST=${POSTGRES_HOST} - POSTGRES_DATABASE=postgres - POSTGRES_USER=${POSTGRES_USER} diff --git a/main.py b/main.py index e2ec543..539cb99 100644 --- a/main.py +++ b/main.py @@ -1,9 +1,13 @@ from telethon import TelegramClient, events from telethon import types + from datetime import timedelta +from PIL import Image +import io +import pytesseract + from db import DBAction, db_action, create_table -import aiosqlite from df_utils import detect_intent_text @@ -15,6 +19,8 @@ from dotenv import load_dotenv load_dotenv() +languages = ['rus', 'eng'] + bot = TelegramClient(os.environ['SESSION_NAME'], int(os.environ['API_ID']), os.environ['API_HASH']).start(bot_token=os.environ['BOT_TOKEN']) @@ -176,35 +182,50 @@ async def manage_stickerpack(msg): raise events.StopPropagation +async def on_mute(msg): + if random.randint(0, 6) == 1: + await msg.respond( + 'Бот работает в режиме обучения, напишите /report в ответ на сообщение, если оно не было распознано или распознано по ошибке') + await bot.edit_permissions(msg.chat_id, msg.sender_id, timedelta(minutes=10), send_messages=False) + + @bot.on(events.NewMessage()) async def process_message(msg): """if await check_admin(msg) or await check_user_in_exceptions(msg.sender_id, msg.chat_id): - return""" + raise events.StopPropagation""" if msg.text: - resp = await detect_intent_text(os.environ['GOOGLE_PROJECT_ID'], msg.chat_id, msg.raw_text, - os.environ['GOOGLE_MODEL_LANGUAGE']) - if resp: - await msg.reply(random.choice(['Харе уже свои вагоны слать', - 'Вагончиками балуемся?', - 'Вагоны тут не приветствуются', - 'Надоел уже со своими вагонами'])) - if random.randint(0, 6) == 1: - await msg.respond( - 'Бот работает в режиме обучения, напишите /report в ответ на сообщение, если оно не было распознано или распознано по ошибке') - await bot.edit_permissions(msg.chat_id, msg.sender_id, timedelta(minutes=10), send_messages=False) - + for lang in languages: + resp = await detect_intent_text(os.environ['GOOGLE_PROJECT_ID'], msg.chat_id, msg.text, lang) + if resp: + await msg.reply(random.choice(['Харе уже свои вагоны слать', + 'Вагончиками балуемся?', + 'Вагоны тут не приветствуются', + 'Надоел уже со своими вагонами'])) + await on_mute(msg) raise events.StopPropagation if msg.sticker: banstickerpacks = [stickerpack[0] for stickerpack in await db_action('SELECT pack_id FROM banned_stickerpacks WHERE chat_id = ?', (msg.chat_id,), DBAction.fetchall)] if any([x.stickerset.id in banstickerpacks for x in msg.sticker.attributes if isinstance(x, types.DocumentAttributeSticker)]): - await bot.edit_permissions(msg.chat_id, msg.sender_id, timedelta(minutes=10), send_messages=False) await msg.reply(random.choice(['Ай-ай-ай, стикеры с вагончиками, не хорошо', 'Стикерами с вагонами балуемся? Негоже так делать', 'Надоел уже со своими вагонами', 'Щас высажу с поезда, будете тут про номерные свои кричать!'])) + await on_mute(msg) + raise events.StopPropagation + + if msg.photo: + for lang in languages: + image_text = pytesseract.image_to_string(Image.open(io.BytesIO(await bot.download_media(msg.photo, file=bytes))), lang=lang) + resp = await detect_intent_text(os.environ['GOOGLE_PROJECT_ID'], msg.chat_id, image_text, lang) + if resp: + await msg.reply(random.choice(['Картинки с текстом про вагончики тоже нельзя 🙃', + 'Мемами с вагонами балуемся? Не надо так делать', + 'Вагоны нельзя ни в каком виде, даже в в мемах'])) + await on_mute(msg) + raise events.StopPropagation async def main(): diff --git a/requirements.txt b/requirements.txt index f5055c8..8db1356 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,6 @@ telethon aiosqlite google-cloud-dialogflow -python-dotenv \ No newline at end of file +python-dotenv +pytesseract +Pillow \ No newline at end of file