from typing import List import aiohttp from vkbottle import API, PhotoMessageUploader 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): def __init__(self, bot_id: int, api_token: str): super().__init__(api_token) self.bot_id = bot_id def get_bot_id(api: API) -> int: my_api: MyAPI = api return my_api.bot_id 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) 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 len(message.text) > 0: 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 async def upload_photo(image: bytes, chat_id: int, api: API) -> str: return await PhotoMessageUploader(api).upload(file_source=image, peer_id=chat_id)