diff --git a/ai_agent.py b/ai_agent.py index b84c1cd..2c2da35 100644 --- a/ai_agent.py +++ b/ai_agent.py @@ -43,11 +43,11 @@ class AiAgent: self.chat_contexts: Dict[int, ChatContext] = {} async def get_reply(self, chat_id: int, chat_prompt: str, - message: AiMessage, forwarded_message: Optional[AiMessage]) -> str: + message: AiMessage, forwarded_messages: List[AiMessage]) -> str: message_text = message.text - if forwarded_message: - message_text += '\n<Цитируемое сообщение от {}>\n'.format(forwarded_message.user_name) - message_text += forwarded_message.text + '\n' + for fwd_message in forwarded_messages: + message_text += '\n<Цитируемое сообщение от {}>\n'.format(fwd_message.user_name) + message_text += fwd_message.text + '\n' message_text += '<Конец цитаты>' context = self._get_chat_context(chat_id, chat_prompt) @@ -81,7 +81,7 @@ class AiAgent: def _get_chat_context(self, chat_id: int, chat_prompt: Optional[str]) -> ChatContext: """Get or create chat context for a specific chat""" if chat_id not in self.chat_contexts: - self.chat_contexts[chat_id] = ChatContext(10) + self.chat_contexts[chat_id] = ChatContext(max_messages=20) prompt = SYSTEM_PROMPT if chat_prompt is not None: prompt += '\n\n' + chat_prompt diff --git a/messages.py b/messages.py index 5d5dfd6..a8f019e 100644 --- a/messages.py +++ b/messages.py @@ -1,6 +1,7 @@ MESSAGE_CHAT_NOT_ACTIVE = 'Извините, но я пока не работаю в этом чате.' MESSAGE_PERMISSION_DENIED = 'Извините, но о таком меня может попросить только администратор чата.' -MESSAGE_NEED_REPLY = 'Извините, но эту команду нужно вызывать в ответном сообщении.' +MESSAGE_NEED_REPLY = 'Извините, но эту команду нужно вызывать в ответ на текстовое сообщение.' +MESSAGE_NEED_REPLY_OR_FORWARD = 'Извините, но эту команду нужно вызывать в ответ на текстовое сообщение или с пересылкой текстовых сообщений.' MESSAGE_DEFAULT_RULES = 'Правила не установлены. Просто ведите себя хорошо.' MESSAGE_DEFAULT_CHECK_RULES = 'Правила чата не установлены. Проверка невозможна.' MESSAGE_DEFAULT_GREETING_JOIN = 'Добро пожаловать, {name}!' diff --git a/tg/handlers/admin.py b/tg/handlers/admin.py index 9fdf47e..95c5f14 100644 --- a/tg/handlers/admin.py +++ b/tg/handlers/admin.py @@ -5,7 +5,7 @@ from aiogram.utils.formatting import Bold import utils from messages import * import tg.tg_database as database -from .default import clear_ai_chat_context, get_ai_reply +from .default import clear_ai_chat_context router = Router() @@ -94,35 +94,6 @@ async def set_ai_prompt_handler(message: Message, bot: Bot): await message.answer('Личность ИИ изменена.') -@router.message(F.text == "!проверка") -async def check_rules_violation_handler(message: Message, bot: Bot): - chat_id = message.chat.id - chat = database.DB.create_chat_if_not_exists(chat_id) - if chat['active'] == 0: - await message.answer(MESSAGE_CHAT_NOT_ACTIVE) - return - - if not await tg_user_is_admin(chat_id, message.from_user.id, bot): - await message.answer(MESSAGE_PERMISSION_DENIED) - return - - if message.reply_to_message is None: - await message.answer(MESSAGE_NEED_REPLY) - return - - chat_rules = chat['rules'] - if chat_rules is None: - await message.answer(MESSAGE_DEFAULT_CHECK_RULES) - return - - prompt = 'В чате действуют следующие правила:\n' - prompt += chat_rules + '\n\n' - prompt += 'Проверь, нарушает ли правила следующее сообщение (используй ссылки на конкретные пункты правил):' - - await message.answer(await get_ai_reply(message.bot, chat_id, message.from_user, prompt, - message.reply_to_message.from_user, message.reply_to_message.text)) - - @router.message(F.text == "!предупреждение") async def warning_handler(message: Message, bot: Bot): chat_id = message.chat.id diff --git a/tg/handlers/default.py b/tg/handlers/default.py index d26081e..74b9e3f 100644 --- a/tg/handlers/default.py +++ b/tg/handlers/default.py @@ -66,7 +66,7 @@ async def any_message_handler(message: Message): bot_user = await message.bot.get_me() ai_message = AiMessage() - ai_fwd_message: Optional[AiMessage] = None + ai_fwd_messages: list[AiMessage] = [] # Ответ на сообщение бота if message.reply_to_message and message.reply_to_message.from_user.id == bot_user.id: @@ -80,9 +80,9 @@ async def any_message_handler(message: Message): return if message.reply_to_message and message.reply_to_message.content_type == ContentType.TEXT: - ai_fwd_message = ( + ai_fwd_messages = [ AiMessage(user_name=await get_user_name_for_ai(message.reply_to_message.from_user), - text=message.reply_to_message.text)) + text=message.reply_to_message.text)] ai_message.user_name = await get_user_name_for_ai(message.from_user) chat_prompt = chat['ai_prompt'] @@ -93,7 +93,7 @@ async def any_message_handler(message: Message): agent = AiAgent(message.bot.config['openrouter_token']) await message.bot.send_chat_action(chat_id, 'typing') - await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_message)) + await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_messages)) def clear_ai_chat_context(chat_id: int): @@ -107,7 +107,7 @@ async def get_ai_reply(bot: Bot, chat_id, user: User, message: str, fwd_user: Us chat_prompt = chat['ai_prompt'] ai_message = AiMessage(user_name=await get_user_name_for_ai(user), text=message) - ai_fwd_message = AiMessage(user_name=await get_user_name_for_ai(fwd_user), text=fwd_message) + ai_fwd_messages = [AiMessage(user_name=await get_user_name_for_ai(fwd_user), text=fwd_message)] global agent if agent is None: @@ -115,4 +115,4 @@ async def get_ai_reply(bot: Bot, chat_id, user: User, message: str, fwd_user: Us agent = AiAgent(bot.config['openrouter_token']) await bot.send_chat_action(chat_id, 'typing') - return await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_message) + return await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_messages) diff --git a/tg/handlers/user.py b/tg/handlers/user.py index 545255c..176ce40 100644 --- a/tg/handlers/user.py +++ b/tg/handlers/user.py @@ -1,12 +1,14 @@ from typing import List, Any from aiogram import Bot, Router, F +from aiogram.enums import ContentType from aiogram.types import Message from aiogram.utils.formatting import Bold, Italic import utils from messages import * import tg.tg_database as database +from .default import get_ai_reply router = Router() @@ -44,6 +46,7 @@ async def help_handler(message: Message): response += '!месяц - статистика сообщений за месяц\n' response += '!молчуны - список молчунов\n' response += '!предупреждения - список участников с предупреждениями\n' + response += '!проверка* - проверить сообщения на нарушение правил\n' response += '\n' response += Bold('Команды для администраторов') + '\n' response += '!старт - начать работу в чате\n' @@ -129,3 +132,28 @@ async def warnings_handler(message: Message, bot: Bot): response = Bold('Участники с предупреждениями') + '\n' response += await format_rating(chat_id, top_users, bot) await message.answer(**response.as_kwargs()) + + +@router.message(F.text == "!проверка") +async def check_rules_violation_handler(message: Message): + chat_id = message.chat.id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + if message.reply_to_message is None or message.reply_to_message.content_type != ContentType.TEXT: + await message.answer(MESSAGE_NEED_REPLY) + return + + chat_rules = chat['rules'] + if chat_rules is None: + await message.answer(MESSAGE_DEFAULT_CHECK_RULES) + return + + prompt = 'В чате действуют следующие правила:\n' + prompt += chat_rules + '\n\n' + prompt += 'Проверь, нарушает ли правила следующее сообщение (используй ссылки на конкретные пункты правил):' + + await message.answer(await get_ai_reply(message.bot, chat_id, message.from_user, prompt, + message.reply_to_message.from_user, message.reply_to_message.text)) diff --git a/vk/handlers/admin.py b/vk/handlers/admin.py index 21ddc48..347aac7 100644 --- a/vk/handlers/admin.py +++ b/vk/handlers/admin.py @@ -6,7 +6,7 @@ from vkbottle_types.codegen.objects import MessagesGetConversationMembers from messages import * import utils import vk.vk_database as database -from .default import clear_ai_chat_context, get_ai_reply +from .default import clear_ai_chat_context labeler = BotLabeler() @@ -149,36 +149,6 @@ async def set_ai_prompt_handler(message: Message): await message.answer('Личность ИИ изменена.') -@labeler.chat_message(text="!проверка") -async def check_rules_violation_handler(message: Message): - chat_id = message.peer_id - chat = database.DB.create_chat_if_not_exists(chat_id) - if chat['active'] == 0: - await message.answer(MESSAGE_CHAT_NOT_ACTIVE) - return - - chat_members = await message.ctx_api.messages.get_conversation_members(peer_id=chat_id, extended=False) - if not vk_user_is_admin(message.from_id, chat_members): - await message.answer(MESSAGE_PERMISSION_DENIED) - return - - if message.reply_message is None: - await message.answer(MESSAGE_NEED_REPLY) - return - - chat_rules = chat['rules'] - if chat_rules is None: - await message.answer(MESSAGE_DEFAULT_CHECK_RULES) - return - - prompt = 'В чате действуют следующие правила:\n' - prompt += chat_rules + '\n\n' - prompt += 'Проверь, нарушает ли правила следующее сообщение (используй ссылки на конкретные пункты правил):' - - await message.answer(await get_ai_reply(message.ctx_api, chat_id, message.from_id, prompt, - message.reply_message.from_id, message.reply_message.text)) - - @labeler.chat_message(text="!предупреждение") async def warning_handler(message: Message): chat_id = message.peer_id diff --git a/vk/handlers/default.py b/vk/handlers/default.py index d41e72d..f60ae78 100644 --- a/vk/handlers/default.py +++ b/vk/handlers/default.py @@ -1,5 +1,5 @@ import re -from typing import Optional +from typing import Optional, Tuple, List from vkbottle import API from vkbottle.bot import Message @@ -50,7 +50,7 @@ async def any_message_handler(message: Message): bot_user = (await message.ctx_api.groups.get_by_id()).groups[0] ai_message = AiMessage() - ai_fwd_message: Optional[AiMessage] = None + ai_fwd_messages: list[AiMessage] = [] # Ответ на сообщение бота if message.reply_message and message.reply_message.from_id == -bot_user.id: @@ -68,10 +68,16 @@ async def any_message_handler(message: Message): else: return - if message.reply_message and message.reply_message.text is not None: - ai_fwd_message = ( + if message.reply_message and len(message.reply_message.text) > 0: + ai_fwd_messages.append( AiMessage(user_name=await get_user_name_for_ai(message.ctx_api, message.reply_message.from_id), text=message.reply_message.text)) + else: + for fwd_message in message.fwd_messages: + if len(fwd_message.text) > 0: + ai_fwd_messages.append( + AiMessage(user_name=await get_user_name_for_ai(message.ctx_api, fwd_message.from_id), + text=fwd_message.text)) ai_message.user_name = await get_user_name_for_ai(message.ctx_api, message.from_id) chat_prompt = chat['ai_prompt'] @@ -82,7 +88,7 @@ async def any_message_handler(message: Message): agent = AiAgent(message.ctx_api.config['openrouter_token']) await message.ctx_api.messages.set_activity(peer_id=chat_id, type='typing') - await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_message)) + await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_messages)) def clear_ai_chat_context(chat_id: int): @@ -91,12 +97,15 @@ def clear_ai_chat_context(chat_id: int): agent.clear_chat_context(chat_id) -async def get_ai_reply(api: API, chat_id, user_id: int, message: str, fwd_user_id: int, fwd_message: str) -> str: +async def get_ai_reply(api: API, chat_id, message: Tuple[int, str], fwd_messages: List[Tuple[int, str]]) -> str: chat = database.DB.create_chat_if_not_exists(chat_id) chat_prompt = chat['ai_prompt'] - ai_message = AiMessage(user_name=await get_user_name_for_ai(api, user_id), text=message) - ai_fwd_message = AiMessage(user_name=await get_user_name_for_ai(api, fwd_user_id), text=fwd_message) + ai_message = AiMessage(user_name=await get_user_name_for_ai(api, message[0]), text=message[1]) + ai_fwd_messages: List[AiMessage] = [] + for fwd_message in fwd_messages: + ai_fwd_messages.append( + AiMessage(user_name=await get_user_name_for_ai(api, fwd_message[0]), text=fwd_message[1])) global agent if agent is None: @@ -104,4 +113,4 @@ async def get_ai_reply(api: API, chat_id, user_id: int, message: str, fwd_user_i agent = AiAgent(api.config['openrouter_token']) await api.messages.set_activity(peer_id=chat_id, type='typing') - return await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_message) + return await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_messages) diff --git a/vk/handlers/user.py b/vk/handlers/user.py index 7ed2d82..323cefe 100644 --- a/vk/handlers/user.py +++ b/vk/handlers/user.py @@ -6,6 +6,7 @@ from vkbottle.framework.labeler import BotLabeler from messages import * import vk.vk_database as database +from .default import get_ai_reply labeler = BotLabeler() @@ -50,6 +51,7 @@ async def rules_handler(message: Message): response += '!молчуны - список молчунов\n' response += '!предупреждения - список участников с предупреждениями\n' response += '!поздравление - запретить/разрешить поздравлять с днем рождения\n' + response += '!проверка* - проверить сообщения на нарушение правил\n' response += '\n' response += bold('Команды для администраторов') + '\n' response += '!старт - начать работу в чате\n' @@ -157,3 +159,35 @@ async def no_birthday_handler(message: Message): await message.answer('Хорошо, я буду поздравлять тебя с днем рождения, если его дата не скрыта.') else: await message.answer('Хорошо, я не буду поздравлять тебя с днем рождения.') + + +@labeler.chat_message(text="!проверка") +async def check_rules_violation_handler(message: Message): + chat_id = message.peer_id + chat = database.DB.create_chat_if_not_exists(chat_id) + if chat['active'] == 0: + await message.answer(MESSAGE_CHAT_NOT_ACTIVE) + return + + chat_rules = chat['rules'] + if chat_rules is None: + await message.answer(MESSAGE_DEFAULT_CHECK_RULES) + return + + prompt = 'В чате действуют следующие правила:\n' + prompt += chat_rules + '\n\n' + prompt += 'Проверь, нарушает ли правила следующее сообщение (используй ссылки на конкретные пункты правил):' + + fwd_messages: list[tuple[int, str]] = [] + if message.reply_message is not None and len(message.reply_message.text) > 0: + fwd_messages.append((message.reply_message.from_id, message.reply_message.text)) + else: + for fwd_message in message.fwd_messages: + if len(fwd_message.text) > 0: + fwd_messages.append((fwd_message.from_id, fwd_message.text)) + + if len(fwd_messages) == 0: + await message.answer(MESSAGE_NEED_REPLY_OR_FORWARD) + return + + await message.answer(await get_ai_reply(message.ctx_api, chat_id, (message.from_id, prompt), fwd_messages))