Добавлена возможность пересылки другого сообщения при обращении к боту.

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

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):

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<bot_id>|<some_name>]
             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):