Реализована отправка изображений в запросе к ИИ для VK.

tg/handlers/default.py

@@ -70,7 +70,7 @@ async def any_message_handler(message: Message, bot: Bot):
                 ai_fwd_messages = [await create_ai_message(message.reply_to_message, bot)]
         else:
             return
-    except UnsupportedContentType:
+    except utils.UnsupportedContentType:
         await message.reply(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
         return
 

tg/handlers/private.py

@@ -48,7 +48,7 @@ async def any_message_handler(message: Message, bot: Bot):
 
     try:
         ai_message = await create_ai_message(message, bot)
-    except UnsupportedContentType:
+    except utils.UnsupportedContentType:
         await message.answer(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
         return
 

tg/utils.py

@@ -6,6 +6,7 @@ from aiogram.enums import ContentType
 from aiogram.types import User, PhotoSize, Message
 
 import ai_agent
+import utils
 
 
 async def get_user_name_for_ai(user: User):
@@ -34,11 +35,6 @@ def get_message_text(message: Message) -> Optional[str]:
         return None
 
 
-class UnsupportedContentType(RuntimeError):
-    def __init__(self):
-        pass
-
-
 async def create_ai_message(message: Message, bot: Bot) -> ai_agent.Message:
     ai_message = ai_agent.Message()
     ai_message.message_id = message.message_id
@@ -50,7 +46,7 @@ async def create_ai_message(message: Message, bot: Bot) -> ai_agent.Message:
             ai_message.text = message.caption
             ai_message.image = await download_photo(message.photo[-1], bot)
         else:
-            raise UnsupportedContentType()
+            raise utils.UnsupportedContentType()
     else:
-        raise UnsupportedContentType()
+        raise utils.UnsupportedContentType()
     return ai_message

utils.py

@@ -6,6 +6,11 @@ from pymorphy3 import MorphAnalyzer
 from time import gmtime
 
 
+class UnsupportedContentType(RuntimeError):
+    def __init__(self):
+        pass
+
+
 _morph = MorphAnalyzer()
 
 

vk/handlers/default.py

@@ -44,64 +44,47 @@ 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 = ai_agent.Message()
-    ai_fwd_messages: list[ai_agent.Message] = []
-
     bot_username_mention = '@' + bot_user.screen_name
     pattern = r"\[club" + str(bot_user.id) + r"\|(.+)]"
     bot_mentioned = False
-    if len(message.text) > 0 and message.text.find(bot_username_mention) != -1:
+    message_text = message.text
-        # Сообщение содержит @bot_username
+
-        ai_message.text = message.text.replace(bot_username_mention, bot_user.name)
+    if len(message_text) > 0 and re.search(pattern, message_text) is not None:
-        bot_mentioned = True
-    elif len(message.text) > 0 and re.search(pattern, message.text) is not None:
         # Сообщение содержит [club<bot_id>|<some_name>]
-        ai_message.text = re.sub(pattern, r'\1', message.text)
+        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:
             # Ответ на сообщение бота
-        if len(message.text) > 0:
-            ai_message.text = message.text
-        else:
-            await message.reply(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
-            return
-
             last_id = ai_agent.agent.get_last_assistant_message_id(bot_id, chat_id)
             if message.reply_message.message_id != last_id:
-            # Оригинального сообщения нет в контексте, или оно не последнее
+                # Оригинального сообщения нет в контексте, или оно не последнее -> переслать его
-            if len(message.reply_message.text) > 0:
+                ai_fwd_messages = [await create_ai_message(message.reply_message)]
-                ai_fwd_messages = [
-                    ai_agent.Message(user_name=await get_user_name_for_ai(message.ctx_api,
-                                                                          message.reply_message.from_id),
-                                     text=message.reply_message.text)]
-            else:
-                await message.reply(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
-                return
         else:
             return
-
+    except utils.UnsupportedContentType:
-    if bot_mentioned:
-        if message.reply_message:
-            if len(message.reply_message.text) > 0:
-                ai_fwd_messages.append(
-                    ai_agent.Message(
-                        user_name=await get_user_name_for_ai(message.ctx_api, message.reply_message.from_id),
-                        text=message.reply_message.text))
-            else:
-                await message.reply(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
-                return
-        else:
-            for fwd_message in message.fwd_messages:
-                if len(fwd_message.text) > 0:
-                    ai_fwd_messages.append(
-                        ai_agent.Message(user_name=await get_user_name_for_ai(message.ctx_api, fwd_message.from_id),
-                                         text=fwd_message.text))
-                else:
         await message.reply(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
         return
 
-    ai_message.user_name = await get_user_name_for_ai(message.ctx_api, message.from_id)
+    ai_message = await create_ai_message(message)
-    ai_message.message_id = message.conversation_message_id
+    ai_message.text = message_text
 
     answer, success = await utils.run_with_progress(
             partial(ai_agent.agent.get_group_chat_reply, bot_id, chat_id, ai_message, ai_fwd_messages),

vk/handlers/private.py

@@ -52,9 +52,11 @@ async def any_message_handler(message: Message):
         await message.answer(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
         return
 
-    ai_message = ai_agent.Message()
+    try:
-    ai_message.text = message.text
+        ai_message = await create_ai_message(message)
-    ai_message.message_id = message.message_id
+    except utils.UnsupportedContentType:
+        await message.answer(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
+        return
 
     answer, success = await utils.run_with_progress(
         partial(ai_agent.agent.get_private_chat_reply, bot_id, chat_id, ai_message),

vk/utils.py

@@ -1,4 +1,14 @@
+from typing import List
+
+import aiohttp
+
 from vkbottle import API
+from vkbottle.bot import Message
+from vkbottle_types.codegen.objects import PhotosPhotoSizes
+from vkbottle_types.objects import MessagesMessageAttachmentType
+
+import ai_agent
+import utils
 
 
 class MyAPI(API):
@@ -18,3 +28,38 @@ async def get_user_name_for_ai(api: API, user_id: int):
         return "{} {}".format(user[0].first_name, user[0].last_name)
     else:
         return '@id' + str(user_id)
+
+
+async def download_photo(photos: List[PhotosPhotoSizes]) -> bytes:
+    max_photo_size = 16*1024*1024
+    async with aiohttp.ClientSession() as session:
+        for size_type in ['w', 'z', 'y', 'x', 'm', 's']:
+            for photo in photos:
+                if photo.type != size_type:
+                    continue
+                async with session.head(photo.url) as response:
+                    if response.status != 200 or response.content_length > max_photo_size:
+                        break
+                async with session.get(photo.url) as response:
+                    if response.status == 200:
+                        return await response.read()
+                break
+
+    raise RuntimeError(f"Failed to download photo. Status code: {response.status}")
+
+
+async def create_ai_message(message: Message) -> ai_agent.Message:
+    ai_message = ai_agent.Message()
+    ai_message.message_id = message.conversation_message_id
+    ai_message.user_name = await get_user_name_for_ai(message.ctx_api, message.from_id)
+    if message.text is not None:
+        ai_message.text = message.text
+    for attachment in message.attachments:
+        if attachment.type == MessagesMessageAttachmentType.PHOTO:
+            ai_message.image = await download_photo(attachment.photo.sizes)
+            break
+        else:
+            continue
+    if ai_message.text is None and ai_message.image is None:
+        raise utils.UnsupportedContentType()
+    return ai_message