import asyncio import io import logging import os from contextlib import asynccontextmanager, nullcontext import aiohttp from aiogram import Bot, Dispatcher, F, Router from aiogram.client.session.aiohttp import AiohttpSession from aiogram.client.telegram import TEST from aiogram.filters import CommandStart from aiogram.types import Message from apinex_client import APINEX_STT_URL, classify_intent, extract_tone_from_images, extract_transcript, generate_posts, transcribe_audio from user_store import get_tone_of_voice, init_db, save_tone_of_voice logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) router = Router() MAX_AUDIO_BYTES = 20 * 1024 * 1024 MAX_IMAGE_BYTES = 10 * 1024 * 1024 TYPING_REFRESH_SECONDS = 4 MEDIA_GROUP_COLLECT_SECONDS = 0.8 _media_groups: dict[tuple[int, str], dict[str, object]] = {} _media_groups_lock = asyncio.Lock() def split_text(text: str, limit: int = 4000) -> list[str]: chunks: list[str] = [] remaining = text.strip() while len(remaining) > limit: split_at = remaining.rfind(" ", 0, limit + 1) if split_at < limit // 2: split_at = limit chunks.append(remaining[:split_at].strip()) remaining = remaining[split_at:].strip() if remaining: chunks.append(remaining) return chunks @router.message(CommandStart()) async def start(message: Message) -> None: await message.answer( "Я помогаю писать посты для соцсетей. Пришли скриншот своих постов, чтобы настроить стиль, " "или попроси написать пост, например: «Придумай пост про запуск нового продукта»." ) async def _send_three_posts(message: Message, posts: list[str]) -> None: for post in posts: await message.answer(post) @asynccontextmanager async def typing_status(bot: Bot, chat_id: int): try: await bot.send_chat_action(chat_id, "typing") except Exception as exc: logger.debug("Typing indicator unavailable: %s", type(exc).__name__) async def refresh_typing() -> None: while True: await asyncio.sleep(TYPING_REFRESH_SECONDS) try: await bot.send_chat_action(chat_id, "typing") except Exception as exc: logger.debug("Typing indicator refresh failed: %s", type(exc).__name__) task = asyncio.create_task(refresh_typing()) try: yield finally: task.cancel() try: await task except asyncio.CancelledError: pass async def _classify_and_respond(message: Message, text: str, show_typing: bool = True) -> None: if message.from_user is None: await message.answer("Не удалось определить пользователя.") return reply_message = message.reply_to_message original_post = "" if reply_message and reply_message.from_user and reply_message.from_user.is_bot: original_post = (reply_message.text or reply_message.caption or "").strip() try: async with (typing_status(message.bot, message.chat.id) if show_typing else nullcontext()): intent = await classify_intent(text, original_post or None) if intent["intent"] == "other": await message.answer("Я помогаю только писать и редактировать посты для социальных сетей.") return if intent["intent"] == "revise_post": if not original_post: await message.answer("Чтобы исправить пост, ответь реплаем на сообщение с его вариантом.") return tone = await asyncio.to_thread(get_tone_of_voice, message.from_user.id) posts = await generate_posts( intent, tone, original_post=original_post, edit_request=text, ) else: tone = await asyncio.to_thread(get_tone_of_voice, message.from_user.id) posts = await generate_posts(intent, tone) await _send_three_posts(message, posts) except (aiohttp.ClientError, asyncio.TimeoutError, RuntimeError, ValueError) as exc: logger.warning("Post request failed: %s", type(exc).__name__) await message.answer("Не получилось подготовить посты. Попробуй ещё раз позже.") except Exception: logger.exception("Unexpected post generation error") await message.answer("Произошла ошибка при подготовке постов.") async def _download_message_image(message: Message, bot: Bot) -> tuple[bytes, str] | None: document = message.document if document: mime_type = (document.mime_type or "").lower() if not mime_type.startswith("image/"): return None file_obj = document elif message.photo: file_obj = message.photo[-1] mime_type = "image/jpeg" else: return None if file_obj.file_size and file_obj.file_size > MAX_IMAGE_BYTES: await message.answer("Скриншот слишком большой. Пришли изображение до 10 МБ.") return None output = io.BytesIO() await bot.download(file_obj, destination=output) image = output.getvalue() if not image: await message.answer("Не получилось скачать изображение. Попробуй ещё раз.") return None if len(image) > MAX_IMAGE_BYTES: await message.answer("Скриншот слишком большой. Пришли изображение до 10 МБ.") return None return image, mime_type async def _process_image_group(messages: list[Message], bot: Bot) -> None: first = messages[0] if first.from_user is None: await first.answer("Не удалось определить пользователя.") return if len(messages) > 10: await first.answer("В одной галерее можно обработать не более 10 изображений.") return images: list[tuple[bytes, str]] = [] for message in messages: image_payload = await _download_message_image(message, bot) if image_payload: images.append(image_payload) if not images: return tone = await extract_tone_from_images(images) if not tone: await first.answer("Не смогла прочитать посты на изображениях. Пришли более чёткие скриншоты.") return await asyncio.to_thread(save_tone_of_voice, first.from_user.id, tone) await first.answer(f"Настроила стиль постов:\n\n{tone}") async def _collect_media_group(message: Message, bot: Bot) -> list[Message]: media_group_id = message.media_group_id if not media_group_id: return [message] key = (message.chat.id, media_group_id) future: asyncio.Future[list[Message]] = asyncio.get_running_loop().create_future() async with _media_groups_lock: group = _media_groups.get(key) if group is None: group = {"messages": [message], "future": future} _media_groups[key] = group asyncio.create_task(_finish_media_group(key)) else: group["messages"].append(message) future = group["future"] try: return await future except Exception: async with _media_groups_lock: group = _media_groups.get(key) if group: try: group["messages"].remove(message) except ValueError: pass raise async def _finish_media_group(key: tuple[int, str]) -> None: await asyncio.sleep(MEDIA_GROUP_COLLECT_SECONDS) async with _media_groups_lock: group = _media_groups.pop(key, None) if group is None: return group["future"].set_result(group["messages"]) async def _handle_image_message(message: Message, bot: Bot) -> None: try: if message.media_group_id: messages = await _collect_media_group(message, bot) if messages[0].message_id != message.message_id: return else: messages = [message] async with typing_status(bot, message.chat.id): await _process_image_group(messages, bot) except (aiohttp.ClientError, asyncio.TimeoutError, RuntimeError, ValueError) as exc: logger.warning("Tone of voice update failed: %s", type(exc).__name__) await message.answer("Не получилось прочитать скриншот. Попробуй отправить его ещё раз.") except Exception: logger.exception("Unexpected image tone extraction error") await message.answer("Произошла ошибка при чтении скриншота.") @router.message(F.photo | (F.document & F.document.mime_type.startswith("image/"))) async def update_tone_from_image(message: Message, bot: Bot) -> None: await _handle_image_message(message, bot) async def _process_voice(message: Message, bot: Bot) -> None: if not os.environ.get("APINEX_API_KEY", "").strip(): await message.answer("Расшифровка голосовых пока не настроена.") return audio_file = message.voice or message.audio if audio_file is None: return if audio_file.file_size and audio_file.file_size > MAX_AUDIO_BYTES: await message.answer("Аудиофайл слишком большой. Максимальный размер — 20 МБ.") return audio_buffer = io.BytesIO() await bot.download(audio_file, destination=audio_buffer) audio = audio_buffer.getvalue() if not audio: await message.answer("Не удалось получить аудиофайл.") return if len(audio) > MAX_AUDIO_BYTES: await message.answer("Аудиофайл слишком большой. Максимальный размер — 20 МБ.") return content_type = "audio/ogg" if message.voice else (message.audio.mime_type or "audio/mpeg") transcript = await transcribe_audio(audio, content_type) if not transcript: await message.answer("Не удалось распознать речь в аудио.") return await _classify_and_respond(message, transcript, show_typing=False) @router.message(F.voice | F.audio) async def transcribe_voice(message: Message, bot: Bot) -> None: try: async with typing_status(bot, message.chat.id): await _process_voice(message, bot) except (aiohttp.ClientError, asyncio.TimeoutError, RuntimeError, ValueError) as exc: logger.warning("Voice post request failed: %s", type(exc).__name__) await message.answer("Не удалось обработать голосовое. Попробуй ещё раз или напиши текстом.") except Exception: logger.exception("Unexpected voice post processing failure") await message.answer("Произошла ошибка при обработке голосового.") @router.message(F.text & ~F.text.startswith("/")) async def handle_text(message: Message) -> None: if message.text: await _classify_and_respond(message, message.text) async def main() -> None: init_db() token = os.environ.get("BOT_TOKEN", "").strip() if not token: raise RuntimeError("BOT_TOKEN is required") test_env = os.environ.get("BOT_ENV", "test").strip().lower() == "test" session = AiohttpSession(api=TEST) if test_env else None bot = Bot(token=token, session=session) dp = Dispatcher() dp.include_router(router) try: await dp.start_polling(bot) finally: await bot.session.close() if __name__ == "__main__": asyncio.run(main())