From 2296176b1cb99d8214f9ec65a542c94b2b2e607d Mon Sep 17 00:00:00 2001 From: Kirill Kirilenko Date: Mon, 29 Dec 2025 23:41:37 +0300 Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D0=B1=D0=B0=D0=B2=D0=BB=D0=B5?= =?UTF-8?q?=D0=BD=D0=B0=20=D0=B2=D0=BE=D0=B7=D0=BC=D0=BE=D0=B6=D0=BD=D0=BE?= =?UTF-8?q?=D1=81=D1=82=D1=8C=20=D0=BF=D0=B5=D1=80=D0=B5=D1=81=D1=8B=D0=BB?= =?UTF-8?q?=D0=BA=D0=B8=20=D0=B4=D1=80=D1=83=D0=B3=D0=BE=D0=B3=D0=BE=20?= =?UTF-8?q?=D1=81=D0=BE=D0=BE=D0=B1=D1=89=D0=B5=D0=BD=D0=B8=D1=8F=20=D0=BF?= =?UTF-8?q?=D1=80=D0=B8=20=D0=BE=D0=B1=D1=80=D0=B0=D1=89=D0=B5=D0=BD=D0=B8?= =?UTF-8?q?=D0=B8=20=D0=BA=20=D0=B1=D0=BE=D1=82=D1=83.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ai_agent.py | 17 +++++++++++++++-- tg/handlers/default.py | 35 +++++++++++++++++++++++------------ vk/handlers/default.py | 32 ++++++++++++++++++++++---------- 3 files changed, 60 insertions(+), 24 deletions(-) diff --git a/ai_agent.py b/ai_agent.py index ec888c4..b84c1cd 100644 --- a/ai_agent.py +++ b/ai_agent.py @@ -31,14 +31,27 @@ class ChatContext: self.messages.pop() +@dataclass() +class AiMessage: + user_name: str = None + text: str = None + + class AiAgent: def __init__(self, api_token: str): self.client = OpenRouter(api_key=api_token) self.chat_contexts: Dict[int, ChatContext] = {} - async def get_reply(self, chat_id: int, chat_prompt: str, user_name: str, message: str) -> str: + async def get_reply(self, chat_id: int, chat_prompt: str, + message: AiMessage, forwarded_message: Optional[AiMessage]) -> str: + message_text = message.text + if forwarded_message: + message_text += '\n<Цитируемое сообщение от {}>\n'.format(forwarded_message.user_name) + message_text += forwarded_message.text + '\n' + message_text += '<Конец цитаты>' + context = self._get_chat_context(chat_id, chat_prompt) - context.add_message(role="user", content=f"[{user_name}]: {message}") + context.add_message(role="user", content=f"[{message.user_name}]: {message_text}") try: # Get response from OpenRouter diff --git a/tg/handlers/default.py b/tg/handlers/default.py index af89806..b92bd59 100644 --- a/tg/handlers/default.py +++ b/tg/handlers/default.py @@ -6,7 +6,7 @@ from aiogram.types.user import User from aiogram.enums.content_type import ContentType import utils -from ai_agent import AiAgent +from ai_agent import AiAgent, AiMessage import tg.tg_database as database router = Router() @@ -34,6 +34,17 @@ ACCEPTED_CONTENT_TYPES: list[ContentType] = [ ] +async def get_user_name_for_ai(user: User): + if user.first_name and user.last_name: + return "{} {}".format(user.first_name, user.last_name) + elif user.first_name: + return user.first_name + elif user.username: + return user.username + else: + return str(user.id) + + @router.message(F.content_type.in_(ACCEPTED_CONTENT_TYPES)) async def any_message_handler(message: Message): chat_id = message.chat.id @@ -54,26 +65,26 @@ async def any_message_handler(message: Message): if bot_user is None: bot_user = await message.bot.get_me() + ai_message = AiMessage() + ai_fwd_message: Optional[AiMessage] = None + # Ответ на сообщение бота if message.reply_to_message and message.reply_to_message.from_user.id == bot_user.id: - message_text = message.text + ai_message.text = message.text else: # Сообщение содержит @bot_username bot_username_mention = '@' + bot_user.username if message.content_type == ContentType.TEXT and message.text.find(bot_username_mention) != -1: - message_text = message.text.replace(bot_username_mention, bot_user.first_name) + ai_message.text = message.text.replace(bot_username_mention, bot_user.first_name) else: return - if message.from_user.first_name and message.from_user.last_name: - user_name = "{} {}".format(message.from_user.first_name, message.from_user.last_name) - elif message.from_user.first_name: - user_name = message.from_user.first_name - elif message.from_user.username: - user_name = message.from_user.username - else: - user_name = str(message.from_user.id) + if message.reply_to_message and message.reply_to_message.content_type == ContentType.TEXT: + ai_fwd_message = ( + AiMessage(user_name=await get_user_name_for_ai(message.reply_to_message.from_user), + text=message.reply_to_message.text)) + ai_message.user_name = await get_user_name_for_ai(message.from_user) chat_prompt = chat['ai_prompt'] global agent @@ -82,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, user_name, message_text)) + await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_message)) def clear_ai_chat_context(chat_id: int): diff --git a/vk/handlers/default.py b/vk/handlers/default.py index d059af1..2bb8708 100644 --- a/vk/handlers/default.py +++ b/vk/handlers/default.py @@ -1,12 +1,13 @@ import re from typing import Optional +from vkbottle import API from vkbottle.bot import Message from vkbottle.framework.labeler import BotLabeler from vkbottle_types.codegen.objects import GroupsGroup import utils -from ai_agent import AiAgent +from ai_agent import AiAgent, AiMessage import vk.vk_database as database labeler = BotLabeler() @@ -15,6 +16,14 @@ agent: Optional[AiAgent] = None bot_user: Optional[GroupsGroup] = None +async def get_user_name_for_ai(api: API, user_id: int): + user = await api.users.get(user_ids=[user_id]) + if len(user) == 1: + return "{} {}".format(user[0].first_name, user[0].last_name) + else: + return '@id' + str(user_id) + + # Обычные сообщения (не команды и не действия) @labeler.chat_message() async def any_message_handler(message: Message): @@ -40,28 +49,31 @@ async def any_message_handler(message: Message): if bot_user is None: bot_user = (await message.ctx_api.groups.get_by_id()).groups[0] + ai_message = AiMessage() + ai_fwd_message: Optional[AiMessage] = None + # Ответ на сообщение бота if message.reply_message and message.reply_message.from_id == -bot_user.id: - message_text = message.text + ai_message.text = message.text else: # Сообщение содержит @bot_username bot_username_mention = '@' + bot_user.screen_name if message.text is not None and message.text.find(bot_username_mention) != -1: - message_text = message.text.replace(bot_username_mention, bot_user.name) + ai_message.text = message.text.replace(bot_username_mention, bot_user.name) else: # Сообщение содержит [club|] pattern = r"\[club" + str(bot_user.id) + r"\|(.+)]" if message.text is not None and re.search(pattern, message.text) is not None: - message_text = re.sub(pattern, r'\1', message.text) + ai_message.text = re.sub(pattern, r'\1', message.text) else: return - user = await message.ctx_api.users.get(user_ids=[message.from_id]) - if len(user) == 1: - user_name = "{} {}".format(user[0].first_name, user[0].last_name) - else: - user_name = '@id' + str(message.from_id) + if message.reply_message and message.reply_message.text is not None: + ai_fwd_message = ( + AiMessage(user_name=await get_user_name_for_ai(message.ctx_api, message.reply_message.from_id), + text=message.reply_message.text)) + ai_message.user_name = await get_user_name_for_ai(message.ctx_api, message.from_id) chat_prompt = chat['ai_prompt'] global agent @@ -70,7 +82,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, user_name, message_text)) + await message.reply(await agent.get_reply(chat_id, chat_prompt, ai_message, ai_fwd_message)) def clear_ai_chat_context(chat_id: int):