Kirill Kirilenko
e0f521256b
Исправлен выбор размера изображения для Seedream. Обновлена модель генерации аниме.
93 lines
3.5 KiB
Python
93 lines
3.5 KiB
Python
from functools import partial
|
||
|
||
from aiogram import Router, F, Bot
|
||
from aiogram.types import Message
|
||
from aiogram.enums.content_type import ContentType
|
||
|
||
import ai_agent
|
||
import utils
|
||
from messages import *
|
||
|
||
import tg.tg_database as database
|
||
from tg.utils import *
|
||
|
||
router = Router()
|
||
|
||
ACCEPTED_CONTENT_TYPES: list[ContentType] = [
|
||
ContentType.TEXT,
|
||
ContentType.ANIMATION,
|
||
ContentType.AUDIO,
|
||
ContentType.DOCUMENT,
|
||
ContentType.PAID_MEDIA,
|
||
ContentType.PHOTO,
|
||
ContentType.STICKER,
|
||
ContentType.STORY,
|
||
ContentType.VIDEO,
|
||
ContentType.VIDEO_NOTE,
|
||
ContentType.VOICE,
|
||
ContentType.CHECKLIST,
|
||
ContentType.CONTACT,
|
||
ContentType.POLL,
|
||
ContentType.VENUE,
|
||
ContentType.LOCATION
|
||
]
|
||
|
||
|
||
@router.message(F.content_type.in_(ACCEPTED_CONTENT_TYPES))
|
||
async def any_message_handler(message: Message, bot: Bot):
|
||
chat_id = message.chat.id
|
||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||
if chat['active'] == 0:
|
||
return
|
||
|
||
# Игнорировать ботов
|
||
if message.from_user.is_bot:
|
||
return
|
||
|
||
user_id = message.from_user.id
|
||
database.DB.create_user_if_not_exists(bot.id, chat_id, user_id)
|
||
database.DB.user_set_last_message(bot.id, chat_id, user_id, utils.posix_time())
|
||
database.DB.user_increment_messages(bot.id, chat_id, user_id)
|
||
|
||
bot_user = await bot.me()
|
||
|
||
ai_fwd_messages: list[ai_agent.Message] = []
|
||
|
||
try:
|
||
message_text = get_message_text(message)
|
||
bot_username_mention = '@' + bot_user.username
|
||
if message_text is not None and message_text.find(bot_username_mention) != -1:
|
||
# Сообщение содержит @bot_username
|
||
message_text = message_text.replace(bot_username_mention, bot_user.first_name)
|
||
if message.reply_to_message:
|
||
# Сообщение также является ответом -> переслать оригинальное сообщение
|
||
ai_fwd_messages = [await create_ai_message(message.reply_to_message, bot)]
|
||
elif message.reply_to_message and message.reply_to_message.from_user.id == bot_user.id:
|
||
# Ответ на сообщение бота
|
||
last_id = ai_agent.agent.get_last_assistant_message_id(bot.id, chat_id)
|
||
if message.reply_to_message.message_id != last_id:
|
||
# Оригинального сообщения нет в контексте, или оно не последнее -> переслать его
|
||
ai_fwd_messages = [await create_ai_message(message.reply_to_message, bot)]
|
||
else:
|
||
return
|
||
except utils.UnsupportedContentType:
|
||
await message.reply(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
|
||
return
|
||
|
||
ai_message = await create_ai_message(message, bot)
|
||
ai_message.text = message_text
|
||
|
||
answer: ai_agent.Message
|
||
success: bool
|
||
answer, success = await utils.run_with_progress(
|
||
partial(ai_agent.agent.get_group_chat_reply, bot.id, chat_id, ai_message, ai_fwd_messages),
|
||
partial(message.bot.send_chat_action, chat_id, 'typing'),
|
||
interval=4)
|
||
|
||
if answer.image is not None:
|
||
answer_id = (await message.reply_photo(photo=wrap_photo(answer.image), caption=answer.text)).message_id
|
||
await message.answer_document(document=wrap_document(answer.image_hires, 'image', 'png'))
|
||
else:
|
||
answer_id = (await message.reply(answer.text)).message_id
|
||
if success:
|
||
ai_agent.agent.set_last_response_id(bot.id, chat_id, answer_id)
|