Kirill Kirilenko
e0f521256b
Исправлен выбор размера изображения для Seedream. Обновлена модель генерации аниме.
103 lines
4 KiB
Python
103 lines
4 KiB
Python
import re
|
||
from functools import partial
|
||
from typing import Optional
|
||
|
||
from vkbottle.bot import Message
|
||
from vkbottle.framework.labeler import BotLabeler
|
||
from vkbottle_types.codegen.objects import GroupsGroup
|
||
|
||
import ai_agent
|
||
import utils
|
||
from messages import *
|
||
|
||
import vk.vk_database as database
|
||
from vk.utils import *
|
||
|
||
labeler = BotLabeler()
|
||
|
||
bot_user: Optional[GroupsGroup] = None
|
||
|
||
|
||
# Обычные сообщения (не команды и не действия)
|
||
@labeler.chat_message()
|
||
async def any_message_handler(message: Message):
|
||
bot_id = get_bot_id(message.ctx_api)
|
||
chat_id = message.peer_id
|
||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||
if chat['active'] == 0:
|
||
return
|
||
|
||
# Игнорировать ботов
|
||
if message.from_id < 0:
|
||
return
|
||
|
||
# Не учитывать служебные сообщения
|
||
if message.action is not None:
|
||
return
|
||
|
||
user_id = message.from_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)
|
||
|
||
global bot_user
|
||
if bot_user is None:
|
||
bot_user = (await message.ctx_api.groups.get_by_id()).groups[0]
|
||
|
||
bot_username_mention = '@' + bot_user.screen_name
|
||
pattern = r"\[club" + str(bot_user.id) + r"\|(.+)]"
|
||
bot_mentioned = False
|
||
message_text = message.text
|
||
|
||
if len(message_text) > 0 and re.search(pattern, message_text) is not None:
|
||
# Сообщение содержит [club<bot_id>|<some_name>]
|
||
message_text = re.sub(pattern, r'\1', message_text)
|
||
bot_mentioned = True
|
||
|
||
if len(message_text) > 0 and message.text.find(bot_username_mention) != -1:
|
||
# Сообщение содержит @bot_username
|
||
message_text = message_text.replace(bot_username_mention, bot_user.name)
|
||
bot_mentioned = True
|
||
|
||
ai_fwd_messages: list[ai_agent.Message] = []
|
||
|
||
try:
|
||
if bot_mentioned:
|
||
# Бот упомянут в тексте сообщения
|
||
if message.reply_message:
|
||
# Сообщение также является ответом -> переслать оригинальное сообщение
|
||
ai_fwd_messages = [await create_ai_message(message.reply_message)]
|
||
else:
|
||
# Пересылаем все пересланные сообщения
|
||
for fwd_message in message.fwd_messages:
|
||
ai_fwd_messages.append(await create_ai_message(fwd_message))
|
||
elif message.reply_message and message.reply_message.from_id == -bot_user.id:
|
||
# Ответ на сообщение бота
|
||
last_id = ai_agent.agent.get_last_assistant_message_id(bot_id, chat_id)
|
||
if message.reply_message.message_id != last_id:
|
||
# Оригинального сообщения нет в контексте, или оно не последнее -> переслать его
|
||
ai_fwd_messages = [await create_ai_message(message.reply_message)]
|
||
else:
|
||
return
|
||
except utils.UnsupportedContentType:
|
||
await message.reply(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
|
||
return
|
||
|
||
ai_message = await create_ai_message(message)
|
||
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.ctx_api.messages.set_activity, peer_id=chat_id, type='typing'),
|
||
interval=4)
|
||
|
||
if answer.image is not None:
|
||
photo = await upload_photo(answer.image_hires, chat_id=chat_id, api=message.ctx_api)
|
||
answer_id = (await message.reply(answer.text, attachment=photo)).conversation_message_id
|
||
else:
|
||
answer_id = (await message.reply(answer.text)).conversation_message_id
|
||
|
||
if success:
|
||
ai_agent.agent.set_last_response_id(bot_id, chat_id, answer_id)
|