diff --git a/Dockerfile b/Dockerfile index e4ab46d..66dffbd 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM python:3.13-slim +FROM python:3.11-slim RUN apt-get update && apt-get install -y --no-install-recommends ffmpeg && rm -rf /var/lib/apt/lists/* @@ -10,5 +10,6 @@ RUN pip install --no-cache-dir -r requirements.txt COPY . . VOLUME /app/bot/data +VOLUME /app/models CMD ["python", "main.py"] diff --git a/README.md b/README.md index e4de6b2..cd9388d 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,11 @@ - **Модерация стикеров и GIF** — если пользователь отправляет 25+ стикеров/GIF за 60 секунд, бот удаляет все последующие стикеры и GIF от этого пользователя в течение 5 минут. Текстовые сообщения не затрагиваются. - **Восстановление раскладки (`/res`)** — ответьте командой `/res` на сообщение с текстом в неправильной раскладке, и бот переведёт его (например, `gbdj` → `пиво`). - **Погода (`/weather`)** — покажет подробную погоду в указанном городе через OpenWeatherMap API. -- **AI-ассистент "Астра" (`/ai`)** — перешлите сообщение или ответьте на него командой `/ai`, и Astra проанализирует его с учётом контекста последних 10 сообщений. +- **AI-ассистент "Астра" (`/ai`)** — задайте вопрос, ответьте командой `/ai` на сообщение или перешлите сообщение, и Astra проанализирует его с учётом контекста и ранее сохранённых выжимок диалогов. +- **Диалог с Астрой** — после `/ai` без вопроса начинается диалог. Отвечайте на сообщения бота, чтобы продолжать. + - Фаза 1: до 50 сообщений. + - Фаза 2: дополнительно до 20 сообщений. + - После исчерпания лимита — перерыв 1 час. - **Управление AI (`/aino`, `/aiyes`)** — владелец чата может заблокировать/разблокировать пользователя от AI. - **Поддержка SOCKS5 прокси** — настраивается через `.env`. @@ -20,24 +24,33 @@ umb/ ├── requirements.txt ├── main.py # точка входа ├── config.py # конфигурация +├── Dockerfile +├── docker-compose.yml └── bot/ ├── __init__.py ├── bot.py # инициализация бота и сессии + ├── setup_commands.py # меню команд ├── routers/ │ ├── __init__.py │ ├── moderation.py # роутер модерации стикеров/GIF │ ├── layout.py # роутер /start, /res │ ├── weather.py # роутер /weather - │ └── ai.py # роутер /ai, /aino, /aiyes - ├── utils/ - │ ├── __init__.py - │ ├── layout_converter.py # конвертер раскладки en→ru - │ ├── weather.py # функция получения погоды - │ ├── ai_client.py # OpenRouter AI клиент - │ ├── database.py # SQLAlchemy модели и функции БД - │ └── moderation.py # логика модерации (устарел, теперь в database.py) - └── data/ - └── umb.db # SQLite база данных + │ ├── ai.py # роутер /ai, /aino, /aiyes, /aiclear, /aiuser + │ ├── dialogue.py # диалоговый режим с Астрой + │ ├── voice.py # распознавание голосовых сообщений + │ └── yadisk.py # скачивание с Яндекс.Диска + └── utils/ + ├── __init__.py + ├── proxy.py # общий SOCKS5-коннектор + ├── layout_converter.py # конвертер раскладки en→ru + ├── weather.py # функция получения погоды + ├── ai_client.py # OpenRouter/RouterAI клиент + ├── memory.py # саммари и эмбеддинги + ├── database.py # SQLAlchemy модели и функции БД + ├── voice.py # загрузка и транскрибация голоса + ├── yadisk_download.py # загрузка файлов с Яндекс.Диска + ├── s3_client.py # загрузка больших файлов в S3 + └── logging_config.py # настройка логирования ``` ## Установка @@ -61,6 +74,11 @@ API_WEATHER=your_openweather_api_key OPENROUTER_API_KEY=your_openrouter_api_key ROUTERAI_API_KEY=your_routerai_api_key ROUTERAI_BASE_URL=https://routerai.ru/api/v1 + +# Опционально: периодическая проверка работоспособности бесплатных моделей. +# Внимание: проверка расходует токены API. По умолчанию отключена. +AI_HEALTH_CHECK_ENABLED=false +AI_HEALTH_CHECK_INTERVAL=600 ``` Для включения прокси установите `PROXY_ENABLED=true` и укажите корректный `PROXY_URL`. @@ -72,6 +90,14 @@ source .venv/bin/activate python main.py ``` +## Запуск в Docker + +```bash +docker compose up -d --build +``` + +Модели Whisper сохраняются в `./models`, база данных и логи — в `./bot/data`. + ## Команды | Команда | Описание | @@ -79,9 +105,15 @@ python main.py | `/start` | Приветствие и информация о боте | | `/res` | Восстановить раскладку (ответить на сообщение) | | `/weather <город>` | Погода в указанном городе | -| `/ai` | AI-анализ пересланного/ответного сообщения | +| `/ai <вопрос>` | Задать вопрос Астре | +| `/ai` (ответом/пересылкой) | Проанализировать сообщение | +| `/ai` (без вопроса) | Начать диалоговый режим | +| `/aiclear` | Очистить диалог и сохранить выжимку | +| `/aiuser` | Список пользователей чата (только creator) | | `/aino @username` | Заблокировать пользователя от AI (только creator, 24ч) | | `/aiyes @username` | Разблокировать пользователя от AI (только creator) | +| `/ydf <ссылка>` | Скачать файл с Яндекс.Диска | +| `/ydf (new) <ссылка>` | Скачать файл, игнорируя кеш Telegram | ## Настройки модерации @@ -93,12 +125,9 @@ python main.py ## AI-ассистент "Астра" -- **Бесплатные модели** (OpenRouter): пробует по очереди все, при rate-limit переключается на следующую - - `deepseek/deepseek-v4-flash:free` - - `google/gemma-4-26b-a4b-it:free` - - `minimax/minimax-m2.5:free` - - `qwen/qwen3-next-80b-a3b-instruct:free` -- **Платный fallback** (RouterAI): если все бесплатные модели вернули 429, используется `deepseek/deepseek-v4-flash` через `routerai.ru`. В ответе добавляется уведомление `⚡ Обработано через платный API`. -- **Контекст**: хранит последние 10 сообщений каждого пользователя в SQLite -- **Стиль**: отвечает кратко и по делу, если не просят развёрнуто +- **Бесплатные модели** (OpenRouter): пробует по очереди доступные модели с суффиксом `:free`. При rate-limit автоматически переключается на следующую. +- **Платный fallback** (RouterAI): если все бесплатные модели недоступны, используется `deepseek/deepseek-v4-flash` через `routerai.ru`. В ответе добавляется уведомление `⚡ Обработано через платный API`. +- **Контекст**: хранит последние 10 сообщений каждого пользователя в SQLite. +- **Долгосрочная память**: при переходе между фазами диалога и при `/aiclear` создаётся краткая выжимка, которая затем подбирается по смыслу к новым вопросам. +- **Стиль**: отвечает кратко и по делу, если не просят развёрнуто. - **Потеря контекста**: говорит что-то милое ("я потеряла мысль", "смотри, какая птичка!") diff --git a/bot/routers/ai.py b/bot/routers/ai.py index dc2ac56..3f75cbe 100644 --- a/bot/routers/ai.py +++ b/bot/routers/ai.py @@ -1,4 +1,5 @@ import logging +import time from aiogram import Router, F from aiogram.types import Message @@ -15,11 +16,19 @@ from bot.utils.database import ( clear_dialogue, clear_user_context, get_or_create_dialogue, - increment_dialogue_count, save_chat_user, ) -from bot.utils.memory import save_summary_with_embedding, find_relevant_summaries, generate_summary -from config import AI_BLOCK_DEFAULT_DURATION, AI_CONTEXT_LIMIT, AI_DIALOGUE_LIMIT, AI_PHASE2_LIMIT +from bot.utils.memory import ( + save_summary_with_embedding, + find_relevant_summaries, + generate_summary, +) +from config import ( + AI_BLOCK_DEFAULT_DURATION, + AI_CONTEXT_LIMIT, + AI_DIALOGUE_LIMIT, + AI_PHASE2_LIMIT, +) logger = logging.getLogger(__name__) @@ -48,8 +57,8 @@ async def get_user_id_by_username(message: Message, username: str) -> int | None if username_clean.isdigit(): return int(username_clean) - except Exception as e: - logger.error("Error resolving username %s: %s", username, e) + except Exception as exc: + logger.error("Error resolving username %s: %s", username, exc) return None @@ -62,8 +71,8 @@ async def _is_creator(message: Message) -> bool: for admin in admins: if admin.status == "creator" and admin.user.id == message.from_user.id: return True - except Exception as e: - logger.error("Error checking creator: %s", e) + except Exception as exc: + logger.error("Error checking creator: %s", exc) return False @@ -76,10 +85,18 @@ async def cmd_ai(message: Message): user_id = message.from_user.id chat_id = message.chat.id - await save_chat_user(user_id, chat_id, message.from_user.username, message.from_user.full_name or "") + await save_chat_user( + user_id, + chat_id, + message.from_user.username, + message.from_user.full_name or "", + ) if await is_ai_blocked(user_id, chat_id): - await message.answer("Тебе временно недоступен AI. Обратись к владельцу чата.", parse_mode=None) + await message.answer( + "Тебе временно недоступен AI. Обратись к владельцу чата.", + parse_mode=None, + ) return text_parts = message.text.split(maxsplit=1) @@ -88,12 +105,21 @@ async def cmd_ai(message: Message): if message.reply_to_message: target_text = message.reply_to_message.text or message.reply_to_message.caption if not target_text: - await message.answer("Могу работать только с текстовыми сообщениями.", parse_mode=None) + await message.answer( + "Могу работать только с текстовыми сообщениями.", + parse_mode=None, + ) return prompt = f"Проанализируй это сообщение:\n\n{target_text}" context_text = target_text - elif message.forward_from and (message.forward_from.text or message.forward_from.caption): - target_text = message.forward_from.text or message.forward_from.caption + elif message.forward_origin: + target_text = message.text or message.caption + if not target_text: + await message.answer( + "Могу работать только с текстовыми сообщениями.", + parse_mode=None, + ) + return prompt = f"Проанализируй это сообщение:\n\n{target_text}" context_text = target_text elif has_direct_question: @@ -107,21 +133,26 @@ async def cmd_ai(message: Message): context = await get_user_context(user_id, chat_id, AI_CONTEXT_LIMIT) + extra_system_content = None try: relevant = await find_relevant_summaries(user_id, chat_id, prompt, top_k=2) if relevant: - summary_text = "Из прошлых диалогов:\n" + "\n---\n".join(relevant) - context.insert(0, {"role": "system", "content": summary_text}) - except Exception as e: - logger.error("Error fetching relevant summaries: %s", e) + extra_system_content = "Из прошлых диалогов:\n" + "\n---\n".join(relevant) + except Exception as exc: + logger.error("Error fetching relevant summaries: %s", exc) async def update_status(text: str): try: await status_msg.edit_text(text, parse_mode=None) - except Exception: - pass + except Exception as exc: + logger.debug("Status update failed: %s", exc) - response = await ask_ai(prompt, context, status_callback=update_status) + response = await ask_ai( + prompt, + context, + status_callback=update_status, + extra_system_content=extra_system_content, + ) if not response: response = "Не удалось получить ответ от AI." @@ -130,9 +161,14 @@ async def cmd_ai(message: Message): await add_context_message(user_id, chat_id, response, "assistant") try: - await status_msg.edit_text(response, parse_mode="HTML", disable_web_page_preview=True) + await status_msg.edit_text( + response, parse_mode="HTML", disable_web_page_preview=True + ) except Exception: - await status_msg.edit_text(response, parse_mode=None, disable_web_page_preview=True) + try: + await status_msg.edit_text(response, parse_mode=None, disable_web_page_preview=True) + except Exception as exc: + logger.warning("Failed to edit status message: %s", exc) async def _start_dialogue(message: Message): @@ -142,8 +178,10 @@ async def _start_dialogue(message: Message): dialogue = await get_or_create_dialogue(user_id, chat_id) if not dialogue["is_active"]: - remaining = int(dialogue["blocked_until"] - __import__("time").time()) if dialogue["blocked_until"] else 0 - mins = remaining // 60 + remaining = ( + int(dialogue["blocked_until"] - time.time()) if dialogue["blocked_until"] else 0 + ) + mins = max(0, remaining // 60) await message.answer( f"Твой лимит диалога исчерпан. Попробуй через {mins} мин.", parse_mode=None, @@ -182,10 +220,13 @@ async def cmd_aiuser(message: Message): parse_mode=None, ) - msg = await message.answer("✅ Список отправлен в ЛС.", parse_mode=None) - except Exception as e: - logger.error("aiuser error: %s", e) - await message.answer("Не удалось получить список пользователей. Возможно, у бота нет доступа.", parse_mode=None) + await message.answer("✅ Список отправлен в ЛС.", parse_mode=None) + except Exception as exc: + logger.error("aiuser error: %s", exc) + await message.answer( + "Не удалось получить список пользователей. Возможно, у бота нет доступа.", + parse_mode=None, + ) @router.message(Command("aiclear")) @@ -195,28 +236,49 @@ async def cmd_aiclear(message: Message): try: context_messages = await get_user_context(user_id, chat_id, limit=50) - cnt = len(context_messages) if context_messages else 0 - if context_messages and cnt >= 2: - status_msg = await message.answer("Сохраняю выжимку диалога...", parse_mode=None) + cnt = len(context_messages) + if cnt >= 2: + status_msg = await message.answer( + "Сохраняю выжимку диалога...", parse_mode=None + ) summary = await generate_summary(context_messages) if summary: await save_summary_with_embedding(user_id, chat_id, summary) await clear_user_context(user_id, chat_id) - await status_msg.edit_text("✅ Выжимка сохранена. Диалог очищен.", parse_mode=None) - logger.info("aiclear | user=%d chat=%d context=%d summary=%d emb=saved", user_id, chat_id, cnt, len(summary)) + await status_msg.edit_text( + "✅ Выжимка сохранена. Диалог очищен.", parse_mode=None + ) + logger.info( + "aiclear | user=%d chat=%d context=%d summary=%d", + user_id, + chat_id, + cnt, + len(summary), + ) else: - logger.warning("aiclear | user=%d chat=%d context=%d summary=None", user_id, chat_id, cnt) + logger.warning( + "aiclear | user=%d chat=%d context=%d summary=None", + user_id, + chat_id, + cnt, + ) else: - logger.warning("aiclear | user=%d chat=%d not enough context (%d < 2)", user_id, chat_id, cnt) - except Exception as e: - logger.error("aiclear summary error: %s", e) + logger.warning( + "aiclear | user=%d chat=%d not enough context (%d < 2)", + user_id, + chat_id, + cnt, + ) + except Exception as exc: + logger.error("aiclear summary error: %s", exc) await clear_dialogue(user_id, chat_id) - - from bot.utils.database import unblock_user_from_ai await unblock_user_from_ai(user_id, chat_id) - await message.answer("✅ Диалог очищен. Можешь начать новый.", parse_mode=None) + await message.answer( + "✅ Диалог очищен. Можешь начать новый.", + parse_mode=None, + ) @router.message(Command("aino")) @@ -226,16 +288,30 @@ async def cmd_aino(message: Message): username = extract_username(message.text) if not username: - await message.answer("Укажи пользователя: /aino @username или /aino id", parse_mode=None) + await message.answer( + "Укажи пользователя: /aino @username или /aino id", + parse_mode=None, + ) return target_user_id = await get_user_id_by_username(message, username) if not target_user_id: - await message.answer(f"Не удалось найти пользователя {username}.", parse_mode=None) + await message.answer( + f"Не удалось найти пользователя {username}.", + parse_mode=None, + ) return - await block_user_from_ai(target_user_id, message.chat.id, message.from_user.id, AI_BLOCK_DEFAULT_DURATION) - await message.answer(f"Пользователь {username} заблокирован от AI на 24 часа.", parse_mode=None) + await block_user_from_ai( + target_user_id, + message.chat.id, + message.from_user.id, + AI_BLOCK_DEFAULT_DURATION, + ) + await message.answer( + f"Пользователь {username} заблокирован от AI на 24 часа.", + parse_mode=None, + ) @router.message(Command("aiyes")) @@ -245,16 +321,28 @@ async def cmd_aiyes(message: Message): username = extract_username(message.text) if not username: - await message.answer("Укажи пользователя: /aiyes @username или /aiyes id", parse_mode=None) + await message.answer( + "Укажи пользователя: /aiyes @username или /aiyes id", + parse_mode=None, + ) return target_user_id = await get_user_id_by_username(message, username) if not target_user_id: - await message.answer(f"Не удалось найти пользователя {username}.", parse_mode=None) + await message.answer( + f"Не удалось найти пользователя {username}.", + parse_mode=None, + ) return unblocked = await unblock_user_from_ai(target_user_id, message.chat.id) if unblocked: - await message.answer(f"Пользователь {username} разблокирован для AI.", parse_mode=None) + await message.answer( + f"Пользователь {username} разблокирован для AI.", + parse_mode=None, + ) else: - await message.answer(f"Пользователь {username} не был заблокирован от AI.", parse_mode=None) + await message.answer( + f"Пользователь {username} не был заблокирован от AI.", + parse_mode=None, + ) diff --git a/bot/routers/dialogue.py b/bot/routers/dialogue.py index 30b2df8..fe12cbe 100644 --- a/bot/routers/dialogue.py +++ b/bot/routers/dialogue.py @@ -4,6 +4,7 @@ import time from aiogram import Router, F from aiogram.types import Message +from bot.bot import bot from bot.utils.ai_client import ask_ai from bot.utils.database import ( add_context_message, @@ -16,7 +17,11 @@ from bot.utils.database import ( save_chat_user, clear_user_context, ) -from bot.utils.memory import find_relevant_summaries, save_summary_with_embedding, generate_summary +from bot.utils.memory import ( + find_relevant_summaries, + save_summary_with_embedding, + generate_summary, +) from config import AI_CONTEXT_LIMIT, AI_DIALOGUE_LIMIT, AI_PHASE2_LIMIT, AI_COOLDOWN logger = logging.getLogger(__name__) @@ -24,19 +29,34 @@ logger = logging.getLogger(__name__) router = Router() +async def _is_astra_message(replied: Message) -> bool: + """Check whether the replied message was sent by this bot (Astra).""" + if not replied.from_user or not replied.from_user.is_bot: + return False + try: + me = await bot.me() + return replied.from_user.id == me.id + except Exception as exc: + logger.warning("Could not verify bot identity: %s", exc) + return False + + @router.message(F.text, F.reply_to_message.as_("replied")) async def handle_dialogue_reply(message: Message, replied: Message): user_id = message.from_user.id chat_id = message.chat.id - await save_chat_user(user_id, chat_id, message.from_user.username, message.from_user.full_name or "") + await save_chat_user( + user_id, + chat_id, + message.from_user.username, + message.from_user.full_name or "", + ) if await is_ai_blocked(user_id, chat_id): return - replied_from_bot = replied.from_user and replied.from_user.is_bot - - if not replied_from_bot: + if not await _is_astra_message(replied): return dialogue = await get_or_create_dialogue(user_id, chat_id) @@ -48,64 +68,77 @@ async def handle_dialogue_reply(message: Message, replied: Message): current_count = count_result["msg_count"] phase = dialogue["phase"] - logger.info("Dialogue msg | user=%d chat=%d phase=%d count=%d", user_id, chat_id, phase, current_count) + logger.info( + "Dialogue msg | user=%d chat=%d phase=%d count=%d", + user_id, + chat_id, + phase, + current_count, + ) - if phase == 1: - limit = AI_DIALOGUE_LIMIT - else: - limit = AI_PHASE2_LIMIT + limit = AI_DIALOGUE_LIMIT if phase == 1 else AI_PHASE2_LIMIT - if current_count > limit: + # We allow answering on the exact limit message; the next one triggers phase/cooldown. + should_transition = current_count > limit + + if should_transition: if phase == 1: - logger.info("Dialogue phase1→2 | user=%d chat=%d msg_count=%d", user_id, chat_id, current_count) + logger.info( + "Dialogue phase1→2 | user=%d chat=%d msg_count=%d", + user_id, + chat_id, + current_count, + ) await _transition_to_phase2(message, user_id, chat_id) else: - logger.info("Dialogue ended | user=%d chat=%d msg_count=%d", user_id, chat_id, current_count) + logger.info( + "Dialogue ended | user=%d chat=%d msg_count=%d", + user_id, + chat_id, + current_count, + ) await _end_dialogue(message, user_id, chat_id) return - if phase == 1: - context_limit = AI_DIALOGUE_LIMIT - else: - context_limit = AI_PHASE2_LIMIT - + context_limit = AI_CONTEXT_LIMIT context_messages = await get_user_context(user_id, chat_id, limit=context_limit) extra_context = "" try: - relevant = await find_relevant_summaries(user_id, chat_id, message.text or "", top_k=2) + relevant = await find_relevant_summaries( + user_id, chat_id, message.text or "", top_k=2 + ) if relevant: extra_context = "\n\nИз прошлых диалогов:\n" + "\n---\n".join(relevant[:2]) - except Exception as e: - logger.error("Error fetching relevant summaries: %s", e) + except Exception as exc: + logger.error("Error fetching relevant summaries: %s", exc) remaining = limit - current_count warning = "" - if remaining <= 5: + if 0 < remaining <= 5: warning = f"\n\n⚠️ Осталось {remaining} сообщений в этом диалоге." - system_prompt_extra = "" + system_extras = [] if extra_context: - system_prompt_extra += extra_context + system_extras.append(extra_context) if warning: - system_prompt_extra += warning + system_extras.append(warning) - message_text = message.text or "" - if system_prompt_extra: - message_text += "\n\n(Контекст)" + system_prompt_extra + extra_system_content = "\n".join(system_extras) if system_extras else None status_msg = await message.answer("✍️", parse_mode=None) async def update_status(text: str): try: await status_msg.edit_text(text, parse_mode=None) - except Exception: - pass + except Exception as exc: + logger.debug("Status update failed: %s", exc) response = await ask_ai( - message_text, + message.text or "", context_messages, status_callback=update_status, + extra_system_content=extra_system_content, ) if not response: @@ -121,21 +154,23 @@ async def handle_dialogue_reply(message: Message, replied: Message): try: await status_msg.delete() - except Exception: - pass + except Exception as exc: + logger.debug("Status delete failed: %s", exc) - if remaining <= 3 and remaining > 0: + if 0 < remaining <= 3: try: - warn_msg = await message.reply( + await message.reply( f"⚠️ Осталось {remaining} сообщений. Память почти заполнена.", parse_mode=None, ) - except Exception: - pass + except Exception as exc: + logger.debug("Warning message failed: %s", exc) async def _transition_to_phase2(message: Message, user_id: int, chat_id: int): - status_msg = await message.answer("Сохраняю выжимку диалога...", parse_mode=None) + status_msg = await message.answer( + "Сохраняю выжимку диалога...", parse_mode=None + ) try: context = await get_user_context(user_id, chat_id, limit=AI_DIALOGUE_LIMIT) @@ -144,15 +179,17 @@ async def _transition_to_phase2(message: Message, user_id: int, chat_id: int): if summary: await save_summary_with_embedding(user_id, chat_id, summary) await clear_user_context(user_id, chat_id) - except Exception as e: - logger.error("Error saving summary: %s", e) + except Exception as exc: + logger.error("Error saving summary: %s", exc) await reset_dialogue_to_phase2(user_id, chat_id) try: - await status_msg.edit_text("✅ Начинаю новую сессию (осталось 20 сообщений).", parse_mode=None) - except Exception: - pass + await status_msg.edit_text( + "✅ Начинаю новую сессию (осталось 20 сообщений).", parse_mode=None + ) + except Exception as exc: + logger.debug("Status edit failed: %s", exc) async def _end_dialogue(message: Message, user_id: int, chat_id: int): @@ -162,9 +199,12 @@ async def _end_dialogue(message: Message, user_id: int, chat_id: int): summary = await generate_summary(context) if summary: await save_summary_with_embedding(user_id, chat_id, summary) - except Exception as e: - logger.error("Error saving final summary: %s", e) + except Exception as exc: + logger.error("Error saving final summary: %s", exc) await block_dialogue(user_id, chat_id, AI_COOLDOWN) - await message.reply("Твой лимит исчерпан. Возвращайся через час.", parse_mode=None) + await message.reply( + "Твой лимит исчерпан. Возвращайся через час.", + parse_mode=None, + ) diff --git a/bot/routers/moderation.py b/bot/routers/moderation.py index 627feb1..66785ec 100644 --- a/bot/routers/moderation.py +++ b/bot/routers/moderation.py @@ -1,9 +1,17 @@ +import logging + from aiogram import Router, F from aiogram.types import Message -from bot.utils.database import add_sticker_message, ban_user_stickers, is_user_sticker_banned +from bot.utils.database import ( + add_sticker_message, + ban_user_stickers, + is_user_sticker_banned, +) from config import MODERATION_LIMIT +logger = logging.getLogger(__name__) + router = Router() @@ -13,13 +21,28 @@ async def handle_sticker_or_gif(message: Message): chat_id = message.chat.id if await is_user_sticker_banned(user_id, chat_id): - await message.delete() + try: + await message.delete() + logger.info( + "Deleted sticker/gif from banned user | user=%d chat=%d", + user_id, + chat_id, + ) + except Exception as exc: + logger.warning("Could not delete sticker/gif: %s", exc) return count = await add_sticker_message(user_id, chat_id) + logger.debug("Sticker/gif count | user=%d chat=%d count=%d", user_id, chat_id, count) if count >= MODERATION_LIMIT: await ban_user_stickers(user_id, chat_id) + logger.info( + "Sticker/gif ban triggered | user=%d chat=%d count=%d", + user_id, + chat_id, + count, + ) await message.answer( f"⚠️ {message.from_user.full_name}, вы превысили лимит стикеров/GIF!\n" "Отправка стикеров и GIF ограничена на 5 минут.", diff --git a/bot/routers/yadisk.py b/bot/routers/yadisk.py index 7eb407c..7e4f984 100644 --- a/bot/routers/yadisk.py +++ b/bot/routers/yadisk.py @@ -14,6 +14,7 @@ router = Router() logger = logging.getLogger("yadisk") MAX_TELEGRAM_FILE_SIZE = 50 * 1024 * 1024 +MAX_DOWNLOAD_SIZE = 2 * 1024 * 1024 * 1024 # 2 GB safety limit def is_valid_yandex_public_link(value: str) -> bool: @@ -35,19 +36,21 @@ async def yandex_download_handler(message: Message): await message.answer("Укажите ссылку: /ydf https://yadi.sk/...") return - url = parts[1].strip() + raw_arg = parts[1].strip() + force_refresh = raw_arg.startswith("(new)") + url = raw_arg.replace("(new)", "", 1).strip() if force_refresh else raw_arg - if url.startswith("(new)"): - url = url.replace("(new)", "", 1).strip() - if not url: - await message.answer("После (new) укажите ссылку.") - return - await save_file_id(url, None) + if not url: + await message.answer("После (new) укажите ссылку.") + return if not is_valid_yandex_public_link(url): await message.answer("Это не похоже на публичную ссылку Яндекс.Диска.") return + if force_refresh: + await save_file_id(url, None) + cached_id = await get_file_id(url) if cached_id: await message.answer_document( @@ -63,6 +66,8 @@ async def yandex_download_handler(message: Message): nonlocal last_text if total <= 0: return + if total > MAX_DOWNLOAD_SIZE: + raise RuntimeError("Файл слишком большой для загрузки.") pct = int(downloaded / total * 100) new_text = f"Загрузка: {pct}%" if new_text != last_text and pct % 5 == 0: @@ -72,15 +77,20 @@ async def yandex_download_handler(message: Message): except Exception: pass + file_path = None try: file_path = await download_yandex_file(url, progress_callback=progress) except Exception as exc: logger.exception("Ошибка скачивания с Яндекс.Диска") - await status_msg.edit_text("Не удалось скачать файл. Проверьте ссылку.") + await status_msg.edit_text(f"Не удалось скачать файл: {exc}") return try: - file_size = os.path.getsize(file_path) + file_size = await asyncio.to_thread(os.path.getsize, file_path) + + if file_size > MAX_DOWNLOAD_SIZE: + await status_msg.edit_text("Файл превышает максимально допустимый размер.") + return if file_size < MAX_TELEGRAM_FILE_SIZE: doc = FSInputFile(file_path) @@ -97,7 +107,8 @@ async def yandex_download_handler(message: Message): else: await status_msg.edit_text("Не удалось загрузить файл в облако.") finally: - try: - os.remove(file_path) - except OSError: - pass + if file_path: + try: + await asyncio.to_thread(os.remove, file_path) + except OSError: + pass diff --git a/bot/setup_commands.py b/bot/setup_commands.py index 572f9ca..8179772 100644 --- a/bot/setup_commands.py +++ b/bot/setup_commands.py @@ -10,8 +10,10 @@ COMMANDS = [ BotCommand(command="res", description="Исправить раскладку"), BotCommand(command="weather", description="Погода в городе"), BotCommand(command="ai", description="Спросить Астру"), - BotCommand(command="aino", description="Заблокировать (админ)"), - BotCommand(command="aiyes", description="Разблокировать (админ)"), + BotCommand(command="aiclear", description="Очистить диалог с Астрой"), + BotCommand(command="aiuser", description="Список пользователей чата (админ)"), + BotCommand(command="aino", description="Заблокировать от AI (админ)"), + BotCommand(command="aiyes", description="Разблокировать для AI (админ)"), BotCommand(command="ydf", description="Скачать с Яндекс.Диска"), ] diff --git a/bot/utils/ai_client.py b/bot/utils/ai_client.py index fdcdc78..b91ec7f 100644 --- a/bot/utils/ai_client.py +++ b/bot/utils/ai_client.py @@ -1,19 +1,20 @@ import asyncio +import html import logging import re import time import aiohttp -from aiohttp_socks import ProxyConnector, ProxyType +from bot.utils.proxy import get_proxy_connector from config import ( OPENROUTER_API_KEY, AI_SYSTEM_PROMPT, - PROXY_ENABLED, - PROXY_URL, ROUTERAI_API_KEY, ROUTERAI_BASE_URL, ROUTERAI_MODEL, + AI_HEALTH_CHECK_ENABLED, + AI_HEALTH_CHECK_INTERVAL, ) logger = logging.getLogger(__name__) @@ -23,54 +24,35 @@ ROUTERAI_URL = f"{ROUTERAI_BASE_URL}/chat/completions" PAID_NOTICE = "\n\n⚡ Обработано через платный API" -_request_timeout = aiohttp.ClientTimeout(total=60, sock_connect=15, sock_read=30) +_free_models_cache: list[str] = [] +_free_models_cache_time = 0.0 +_free_models_cache_ttl = 3600 + +_working_models_cache: list[str] = [] +_working_models_cache_time = 0.0 +_working_models_cache_ttl = 600 def get_client_timeout(total: int = 60) -> aiohttp.ClientTimeout: return aiohttp.ClientTimeout(total=total, sock_connect=15, sock_read=30) -_free_models_cache = [] -_free_models_cache_time = 0 -_free_models_cache_ttl = 3600 - -_working_models_cache = [] -_working_models_cache_time = 0 -_working_models_cache_ttl = 600 - - -def _get_connector(): - if PROXY_ENABLED and PROXY_URL: - parsed = PROXY_URL.replace("socks5://", "").replace("socks5h://", "") - if "@" in parsed: - auth, host_port = parsed.split("@", 1) - username, password = auth.split(":", 1) - else: - username = None - password = None - host_port = parsed - - host, port = host_port.rsplit(":", 1) - port = int(port) - - return ProxyConnector( - proxy_type=ProxyType.SOCKS5, - host=host, - port=port, - username=username, - password=password, - ) - return None - - def _md_to_html(text: str) -> str: - text = re.sub(r"```(\w*)\n(.*?)```", r"
\2
", text, flags=re.DOTALL) - text = re.sub(r"`(.*?)`", r"\1", text) - text = re.sub(r"\*\*(.*?)\*\*", r"\1", text) - text = re.sub(r"\*(.*?)\*", r"\1", text) - text = re.sub(r"__(.*?)__", r"\1", text) - text = re.sub(r"~~(.*?)~~", r"\1", text) - text = re.sub(r"\[(.*?)\]\((.*?)\)", r'\1', text) + """Convert a small subset of Markdown to Telegram HTML, escaping raw HTML first.""" + text = html.escape(text) + + def _pre_repl(match: re.Match) -> str: + lang = match.group(1) + code = html.unescape(match.group(2)) + return f'
{html.escape(code)}
' if lang else f"
{html.escape(code)}
" + + text = re.sub(r"```(\w*)\n(.*?)```", _pre_repl, text, flags=re.DOTALL) + text = re.sub(r"`([^`]+)`", r"\1", text) + text = re.sub(r"\*\*(.+?)\*\*", r"\1", text) + text = re.sub(r"\*(.+?)\*", r"\1", text) + text = re.sub(r"__(.+?)__", r"\1", text) + text = re.sub(r"~~(.+?)~~", r"\1", text) + text = re.sub(r"\[(.+?)\]\((.+?)\)", r'\1', text) return text @@ -88,41 +70,44 @@ async def _fetch_free_models() -> list[str]: "Content-Type": "application/json", } - try: - connector = _get_connector() - async with aiohttp.ClientSession(connector=connector) as session: - async with session.get("https://openrouter.ai/api/v1/models", headers=headers) as response: - if response.status == 200: - data = await response.json() - models = data.get("data", []) - - free_models = [] - for model in models: - model_id = model.get("id", "") - if model_id.endswith(":free"): - free_models.append(model_id) - - if free_models: - _free_models_cache = free_models - _free_models_cache_time = now - logger.info(f"Fetched {len(free_models)} free models") - return free_models - except Exception as e: - logger.warning(f"Failed to fetch free models: {e}") - - if _free_models_cache: - return _free_models_cache - fallback = [ "deepseek/deepseek-v4-flash:free", "google/gemma-4-26b-a4b-it:free", "minimax/minimax-m2.5:free", "qwen/qwen3-next-80b-a3b-instruct:free", ] + + connector = get_proxy_connector() + try: + async with aiohttp.ClientSession(connector=connector) as session: + async with session.get( + "https://openrouter.ai/api/v1/models", + headers=headers, + timeout=get_client_timeout(30), + ) as response: + if response.status == 200: + data = await response.json() + models = data.get("data", []) + free_models = [m.get("id", "") for m in models if m.get("id", "").endswith(":free")] + + if free_models: + _free_models_cache = free_models + _free_models_cache_time = now + logger.info("Fetched %d free models", len(free_models)) + return free_models + except Exception as exc: + logger.warning("Failed to fetch free models: %s", exc) + finally: + if connector: + await connector.close() + + if _free_models_cache: + return _free_models_cache + return fallback -async def _test_model(session, model: str) -> bool: +async def _test_model(session: aiohttp.ClientSession, model: str) -> bool: headers = { "Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json", @@ -137,94 +122,122 @@ async def _test_model(session, model: str) -> bool: if response.status == 200: data = await response.json() choices = data.get("choices", []) - return bool(choices and choices[0]["message"].get("content")) + return bool(choices and choices[0].get("message", {}).get("content")) except Exception: pass return False -async def _update_working_models(): +async def _update_working_models() -> None: global _working_models_cache, _working_models_cache_time free_models = await _fetch_free_models() if not free_models: return - connector = _get_connector() - working = [] + connector = get_proxy_connector() + working: list[str] = [] - async with aiohttp.ClientSession(connector=connector) as session: - for model in free_models: - if await _test_model(session, model): - working.append(model) - await asyncio.sleep(0.3) + try: + async with aiohttp.ClientSession(connector=connector) as session: + for model in free_models: + if await _test_model(session, model): + working.append(model) + await asyncio.sleep(0.3) + finally: + if connector: + await connector.close() _working_models_cache = working _working_models_cache_time = time.time() - logger.info(f"Health check: {len(working)}/{len(free_models)} models working") + logger.info("Health check: %d/%d models working", len(working), len(free_models)) -async def start_model_health_check(): +async def start_model_health_check() -> None: + if not AI_HEALTH_CHECK_ENABLED: + logger.info("Model health check is disabled") + return + await asyncio.sleep(30) try: await _update_working_models() - except Exception as e: - logger.error(f"Initial health check error: {e}") + except Exception as exc: + logger.error("Initial health check error: %s", exc) while True: - await asyncio.sleep(600) + await asyncio.sleep(AI_HEALTH_CHECK_INTERVAL) try: await _update_working_models() - except Exception as e: - logger.error(f"Health check error: {e}") + except Exception as exc: + logger.error("Health check error: %s", exc) -def _log_usage(source: str, model: str, data: dict, latency: float): +def _log_usage(source: str, model: str, data: dict, latency: float) -> None: usage = data.get("usage") if usage: logger.info( "AI %s | model=%s in_tok=%s out_tok=%s total_tok=%s latency=%.1fs", - source, model, + source, + model, usage.get("prompt_tokens", "?"), usage.get("completion_tokens", "?"), usage.get("total_tokens", "?"), latency, ) else: - logger.info( - "AI %s | model=%s latency=%.1fs", - source, model, latency, - ) + logger.info("AI %s | model=%s latency=%.1fs", source, model, latency) -async def _try_openrouter(session, model: str, messages: list[dict], headers: dict, payload: dict) -> str | None: - payload["model"] = model +async def _try_openrouter( + session: aiohttp.ClientSession, + model: str, + headers: dict, + base_payload: dict, + raw: bool = False, +) -> tuple[str | None, bool]: + """Returns (content, should_retry_later).""" + payload = {**base_payload, "model": model} start = time.monotonic() async with session.post(OPENROUTER_URL, json=payload, headers=headers) as response: latency = time.monotonic() - start + if response.status == 429: + retry_after = response.headers.get("Retry-After") + logger.warning( + "OpenRouter rate limited | model=%s retry_after=%s", + model, + retry_after, + ) + return None, True + if response.status != 200: error_body = await response.text() - logger.warning("OpenRouter error | model=%s status=%s latency=%.1fs error=%s", model, response.status, latency, error_body[:200]) - return None + logger.warning( + "OpenRouter error | model=%s status=%s latency=%.1fs error=%s", + model, + response.status, + latency, + error_body[:200], + ) + return None, False data = await response.json() choices = data.get("choices", []) if not choices: logger.warning("OpenRouter empty choices | model=%s latency=%.1fs", model, latency) - return None + return None, False - content = choices[0]["message"].get("content") + content = choices[0].get("message", {}).get("content") if not content: logger.warning("OpenRouter empty content | model=%s latency=%.1fs", model, latency) - return None + return None, False _log_usage("OpenRouter", model, data, latency) - return _md_to_html(content) + return content if raw else _md_to_html(content), False -async def _try_routerai(session, messages: list[dict]) -> str | None: +async def _try_routerai(session: aiohttp.ClientSession, messages: list[dict]) -> str | None: if not ROUTERAI_API_KEY: logger.warning("RouterAI skipped | key not set") return None @@ -247,7 +260,12 @@ async def _try_routerai(session, messages: list[dict]) -> str | None: if response.status != 200: error_body = await response.text() - logger.warning("RouterAI error | status=%s latency=%.1fs error=%s", response.status, latency, error_body[:200]) + logger.warning( + "RouterAI error | status=%s latency=%.1fs error=%s", + response.status, + latency, + error_body[:200], + ) return None data = await response.json() @@ -256,7 +274,7 @@ async def _try_routerai(session, messages: list[dict]) -> str | None: logger.warning("RouterAI empty choices | latency=%.1fs", latency) return None - content = choices[0]["message"].get("content") + content = choices[0].get("message", {}).get("content") if not content: logger.warning("RouterAI empty content | latency=%.1fs", latency) return None @@ -271,29 +289,24 @@ async def ask_ai_simple(prompt: str) -> str | None: "Content-Type": "application/json", } - payload = { + base_payload = { "messages": [{"role": "user", "content": prompt}], "max_tokens": 512, } - connector = _get_connector() + connector = get_proxy_connector() try: - async with aiohttp.ClientSession(timeout=_request_timeout, connector=connector) as session: + async with aiohttp.ClientSession( + timeout=get_client_timeout(), connector=connector + ) as session: free_models = await _fetch_free_models() for model in free_models[:5]: - payload["model"] = model - start = time.monotonic() - async with session.post(OPENROUTER_URL, json=payload, headers=headers) as response: - latency = time.monotonic() - start - if response.status == 200: - data = await response.json() - choices = data.get("choices", []) - if choices and choices[0]["message"].get("content"): - _log_usage("ask_ai_simple", model, data, latency) - return choices[0]["message"]["content"] - else: - logger.warning("ask_ai_simple fallback fail | model=%s status=%s latency=%.1fs", model, response.status, latency) + result, _ = await _try_openrouter(session, model, headers, base_payload, raw=True) + if result: + # ask_ai_simple returns raw text for internal use (summaries) + return result + logger.warning("ask_ai_simple fallback fail | model=%s", model) logger.warning("ask_ai_simple | all free models failed, trying RouterAI") routerai_payload = { @@ -311,17 +324,29 @@ async def ask_ai_simple(prompt: str) -> str | None: if response.status == 200: data = await response.json() choices = data.get("choices", []) - if choices and choices[0]["message"].get("content"): + if choices and choices[0].get("message", {}).get("content"): _log_usage("ask_ai_simple (RouterAI)", ROUTERAI_MODEL, data, latency) return choices[0]["message"]["content"] - except Exception as e: - logger.error("ask_ai_simple error: %s", e) + except Exception as exc: + logger.error("ask_ai_simple error: %s", exc) + finally: + if connector: + await connector.close() return None -async def ask_ai(prompt: str, context_messages: list[dict] | None = None, status_callback=None) -> str: - messages = [{"role": "system", "content": AI_SYSTEM_PROMPT}] +async def ask_ai( + prompt: str, + context_messages: list[dict] | None = None, + status_callback=None, + extra_system_content: str | None = None, +) -> str: + system_content = AI_SYSTEM_PROMPT + if extra_system_content: + system_content += "\n\n" + extra_system_content + + messages: list[dict] = [{"role": "system", "content": system_content}] if context_messages: messages.extend(context_messages) @@ -335,13 +360,12 @@ async def ask_ai(prompt: str, context_messages: list[dict] | None = None, status "X-Title": "UMB Bot", } - payload = { + base_payload = { "messages": messages, "max_tokens": 512, } - timeout = aiohttp.ClientTimeout(total=60, sock_connect=15, sock_read=30) - connector = _get_connector() + connector = get_proxy_connector() waiting_messages = [ "Думаю...", @@ -354,10 +378,12 @@ async def ask_ai(prompt: str, context_messages: list[dict] | None = None, status ] try: - async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: + async with aiohttp.ClientSession( + timeout=get_client_timeout(), connector=connector + ) as session: free_models = await _fetch_free_models() - if _working_models_cache: + if _working_models_cache and (time.time() - _working_models_cache_time) < _working_models_cache_ttl: models_to_try = [m for m in _working_models_cache if m in free_models] if not models_to_try: models_to_try = free_models @@ -371,7 +397,7 @@ async def ask_ai(prompt: str, context_messages: list[dict] | None = None, status wait_idx = min(i, len(waiting_messages) - 1) await status_callback(waiting_messages[wait_idx]) - result = await _try_openrouter(session, model, messages, or_headers, payload) + result, _ = await _try_openrouter(session, model, or_headers, base_payload) if result: return result @@ -390,8 +416,8 @@ async def ask_ai(prompt: str, context_messages: list[dict] | None = None, status logger.error("ask_ai | all models failed") return "Извини, ни одна модель не смогла ответить. Попробуй позже." - except aiohttp.ClientError as e: - logger.error("AI request network error: %s", e) + except aiohttp.ClientError as exc: + logger.error("AI request network error: %s", exc) return "Не удалось связаться с AI сервисом. Проверь соединение." except asyncio.TimeoutError: logger.error("ask_ai | timeout after all models exhausted") @@ -399,3 +425,6 @@ async def ask_ai(prompt: str, context_messages: list[dict] | None = None, status except Exception: logger.exception("ask_ai | unexpected error") return "Произошла ошибка при обработке запроса." + finally: + if connector: + await connector.close() diff --git a/bot/utils/database.py b/bot/utils/database.py index 1811ae3..b4c9aaa 100644 --- a/bot/utils/database.py +++ b/bot/utils/database.py @@ -1,7 +1,18 @@ import time -from datetime import datetime -from sqlalchemy import Column, Integer, String, Float, BigInteger, Text +from sqlalchemy import ( + Column, + Integer, + String, + Float, + BigInteger, + Text, + Boolean, + Index, + select, + delete, +) +from sqlalchemy.dialects.sqlite import insert from sqlalchemy.orm import declarative_base from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker @@ -20,6 +31,10 @@ class UserContext(Base): role = Column(String(16), nullable=False) timestamp = Column(Float, nullable=False) + __table_args__ = ( + Index("ix_user_context_user_chat", "user_id", "chat_id"), + ) + class AiBlockedUser(Base): __tablename__ = "ai_blocked_users" @@ -31,6 +46,10 @@ class AiBlockedUser(Base): blocked_at = Column(Float, nullable=False) expires_at = Column(Float, nullable=False) + __table_args__ = ( + Index("ix_ai_blocked_user_chat", "user_id", "chat_id"), + ) + class FileIdCache(Base): __tablename__ = "file_ids" @@ -50,6 +69,10 @@ class StickerBan(Base): ban_until = Column(Float, nullable=True) ban_trigger = Column(Integer, default=0) + __table_args__ = ( + Index("ix_sticker_ban_user_chat", "user_id", "chat_id"), + ) + class ChatUser(Base): __tablename__ = "chat_users" @@ -61,6 +84,10 @@ class ChatUser(Base): full_name = Column(String, nullable=False) last_seen = Column(Float, nullable=False) + __table_args__ = ( + Index("ix_chat_user_chat_user", "chat_id", "user_id"), + ) + class DialogueSession(Base): __tablename__ = "dialogue_sessions" @@ -70,10 +97,14 @@ class DialogueSession(Base): chat_id = Column(BigInteger, nullable=False, index=True) phase = Column(Integer, default=1) msg_count = Column(Integer, default=0) - is_active = Column(Integer, default=0) + is_active = Column(Boolean, default=False) blocked_until = Column(Float, nullable=True) last_activity = Column(Float, nullable=False) + __table_args__ = ( + Index("ix_dialogue_session_user_chat", "user_id", "chat_id"), + ) + class ConversationSummary(Base): __tablename__ = "conversation_summaries" @@ -85,6 +116,10 @@ class ConversationSummary(Base): embedding = Column(Text, nullable=True) created_at = Column(Float, nullable=False) + __table_args__ = ( + Index("ix_conversation_summary_user_chat", "user_id", "chat_id"), + ) + engine = create_async_engine(DATABASE_URL, echo=False) async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) @@ -96,13 +131,11 @@ async def init_db(): async def add_context_message(user_id: int, chat_id: int, text: str, role: str) -> None: - if text is None: - text = "" async with async_session() as session: msg = UserContext( user_id=user_id, chat_id=chat_id, - text=text, + text=text or "", role=role, timestamp=time.time(), ) @@ -112,8 +145,6 @@ async def add_context_message(user_id: int, chat_id: int, text: str, role: str) async def get_user_context(user_id: int, chat_id: int, limit: int = 10) -> list[dict]: async with async_session() as session: - from sqlalchemy import select - stmt = ( select(UserContext) .where(UserContext.user_id == user_id, UserContext.chat_id == chat_id) @@ -125,10 +156,28 @@ async def get_user_context(user_id: int, chat_id: int, limit: int = 10) -> list[ return [{"role": m.role, "content": m.text} for m in reversed(messages)] +async def clear_user_context(user_id: int, chat_id: int) -> None: + async with async_session() as session: + stmt = delete(UserContext).where( + UserContext.user_id == user_id, + UserContext.chat_id == chat_id, + ) + await session.execute(stmt) + await session.commit() + + +async def cleanup_old_context_messages(older_than_days: int = 30) -> int: + """Remove user context messages older than the given number of days.""" + cutoff = time.time() - (older_than_days * 86400) + async with async_session() as session: + stmt = delete(UserContext).where(UserContext.timestamp < cutoff) + result = await session.execute(stmt) + await session.commit() + return result.rowcount + + async def is_ai_blocked(user_id: int, chat_id: int) -> bool: async with async_session() as session: - from sqlalchemy import select - stmt = ( select(AiBlockedUser) .where(AiBlockedUser.user_id == user_id, AiBlockedUser.chat_id == chat_id) @@ -138,15 +187,16 @@ async def is_ai_blocked(user_id: int, chat_id: int) -> bool: return result.scalar_one_or_none() is not None -async def block_user_from_ai(user_id: int, chat_id: int, blocked_by: int, duration: int = 86400) -> None: +async def block_user_from_ai( + user_id: int, chat_id: int, blocked_by: int, duration: int = 86400 +) -> None: async with async_session() as session: - from sqlalchemy import select, delete - - stmt = delete(AiBlockedUser).where( - AiBlockedUser.user_id == user_id, - AiBlockedUser.chat_id == chat_id, + await session.execute( + delete(AiBlockedUser).where( + AiBlockedUser.user_id == user_id, + AiBlockedUser.chat_id == chat_id, + ) ) - await session.execute(stmt) now = time.time() entry = AiBlockedUser( @@ -162,8 +212,6 @@ async def block_user_from_ai(user_id: int, chat_id: int, blocked_by: int, durati async def unblock_user_from_ai(user_id: int, chat_id: int) -> bool: async with async_session() as session: - from sqlalchemy import select, delete - stmt = delete(AiBlockedUser).where( AiBlockedUser.user_id == user_id, AiBlockedUser.chat_id == chat_id, @@ -175,8 +223,6 @@ async def unblock_user_from_ai(user_id: int, chat_id: int) -> bool: async def get_sticker_ban(user_id: int, chat_id: int) -> dict | None: async with async_session() as session: - from sqlalchemy import select - stmt = select(StickerBan).where( StickerBan.user_id == user_id, StickerBan.chat_id == chat_id, @@ -195,8 +241,6 @@ async def get_sticker_ban(user_id: int, chat_id: int) -> dict | None: async def add_sticker_message(user_id: int, chat_id: int) -> int: async with async_session() as session: - from sqlalchemy import select - stmt = select(StickerBan).where( StickerBan.user_id == user_id, StickerBan.chat_id == chat_id, @@ -229,8 +273,6 @@ async def add_sticker_message(user_id: int, chat_id: int) -> int: async def ban_user_stickers(user_id: int, chat_id: int, duration: int = 300) -> None: async with async_session() as session: - from sqlalchemy import select - stmt = select(StickerBan).where( StickerBan.user_id == user_id, StickerBan.chat_id == chat_id, @@ -246,8 +288,6 @@ async def ban_user_stickers(user_id: int, chat_id: int, duration: int = 300) -> async def is_user_sticker_banned(user_id: int, chat_id: int) -> bool: async with async_session() as session: - from sqlalchemy import select - stmt = select(StickerBan).where( StickerBan.user_id == user_id, StickerBan.chat_id == chat_id, @@ -268,10 +308,26 @@ async def is_user_sticker_banned(user_id: int, chat_id: int) -> bool: return True +async def unban_user_stickers(user_id: int, chat_id: int) -> bool: + async with async_session() as session: + stmt = select(StickerBan).where( + StickerBan.user_id == user_id, + StickerBan.chat_id == chat_id, + ) + result = await session.execute(stmt) + row = result.scalar_one_or_none() + + if row: + row.ban_until = None + row.count = 0 + row.ban_trigger = 0 + await session.commit() + return True + return False + + async def get_file_id(file_key: str) -> str | None: async with async_session() as session: - from sqlalchemy import select - stmt = select(FileIdCache).where(FileIdCache.file_key == file_key) result = await session.execute(stmt) row = result.scalar_one_or_none() @@ -280,18 +336,22 @@ async def get_file_id(file_key: str) -> str | None: async def save_file_id(file_key: str, file_id: str | None) -> None: async with async_session() as session: - from sqlalchemy import delete - - await session.execute(delete(FileIdCache).where(FileIdCache.file_key == file_key)) - if file_id is not None: - session.add(FileIdCache(file_key=file_key, file_id=file_id)) + if file_id is None: + await session.execute(delete(FileIdCache).where(FileIdCache.file_key == file_key)) + else: + stmt = ( + insert(FileIdCache) + .values(file_key=file_key, file_id=file_id) + .on_conflict_do_update(index_elements=["file_key"], set_={"file_id": file_id}) + ) + await session.execute(stmt) await session.commit() -async def save_chat_user(user_id: int, chat_id: int, username: str | None, full_name: str) -> None: +async def save_chat_user( + user_id: int, chat_id: int, username: str | None, full_name: str +) -> None: async with async_session() as session: - from sqlalchemy import select - stmt = select(ChatUser).where( ChatUser.user_id == user_id, ChatUser.chat_id == chat_id, @@ -306,20 +366,20 @@ async def save_chat_user(user_id: int, chat_id: int, username: str | None, full_ row.full_name = full_name row.last_seen = now else: - session.add(ChatUser( - user_id=user_id, - chat_id=chat_id, - username=username, - full_name=full_name, - last_seen=now, - )) + session.add( + ChatUser( + user_id=user_id, + chat_id=chat_id, + username=username, + full_name=full_name, + last_seen=now, + ) + ) await session.commit() async def get_chat_users(chat_id: int) -> list[dict]: async with async_session() as session: - from sqlalchemy import select - stmt = ( select(ChatUser) .where(ChatUser.chat_id == chat_id) @@ -334,8 +394,6 @@ async def get_chat_users(chat_id: int) -> list[dict]: async def get_or_create_dialogue(user_id: int, chat_id: int) -> dict: async with async_session() as session: - from sqlalchemy import select - stmt = select(DialogueSession).where( DialogueSession.user_id == user_id, DialogueSession.chat_id == chat_id, @@ -351,7 +409,7 @@ async def get_or_create_dialogue(user_id: int, chat_id: int) -> dict: chat_id=chat_id, phase=1, msg_count=0, - is_active=1, + is_active=True, last_activity=now, ) session.add(row) @@ -366,7 +424,7 @@ async def get_or_create_dialogue(user_id: int, chat_id: int) -> dict: "blocked_until": row.blocked_until, } - row.is_active = 1 + row.is_active = True row.last_activity = now await session.commit() @@ -380,8 +438,6 @@ async def get_or_create_dialogue(user_id: int, chat_id: int) -> dict: async def increment_dialogue_count(user_id: int, chat_id: int) -> dict: async with async_session() as session: - from sqlalchemy import select - stmt = select(DialogueSession).where( DialogueSession.user_id == user_id, DialogueSession.chat_id == chat_id, @@ -401,8 +457,6 @@ async def increment_dialogue_count(user_id: int, chat_id: int) -> dict: async def reset_dialogue_to_phase2(user_id: int, chat_id: int) -> None: async with async_session() as session: - from sqlalchemy import select - stmt = select(DialogueSession).where( DialogueSession.user_id == user_id, DialogueSession.chat_id == chat_id, @@ -419,8 +473,6 @@ async def reset_dialogue_to_phase2(user_id: int, chat_id: int) -> None: async def block_dialogue(user_id: int, chat_id: int, duration: int = 3600) -> None: async with async_session() as session: - from sqlalchemy import select - stmt = select(DialogueSession).where( DialogueSession.user_id == user_id, DialogueSession.chat_id == chat_id, @@ -430,14 +482,12 @@ async def block_dialogue(user_id: int, chat_id: int, duration: int = 3600) -> No if row: row.blocked_until = time.time() + duration - row.is_active = 0 + row.is_active = False await session.commit() async def clear_dialogue(user_id: int, chat_id: int) -> None: async with async_session() as session: - from sqlalchemy import select - stmt = select(DialogueSession).where( DialogueSession.user_id == user_id, DialogueSession.chat_id == chat_id, @@ -448,67 +498,46 @@ async def clear_dialogue(user_id: int, chat_id: int) -> None: if row: row.phase = 1 row.msg_count = 0 - row.is_active = 0 + row.is_active = False + row.blocked_until = None row.last_activity = time.time() await session.commit() -async def clear_user_context(user_id: int, chat_id: int) -> None: +async def save_conversation_summary( + user_id: int, chat_id: int, summary: str, embedding: str | None = None +) -> None: async with async_session() as session: - from sqlalchemy import delete - - stmt = delete(UserContext).where( - UserContext.user_id == user_id, - UserContext.chat_id == chat_id, + session.add( + ConversationSummary( + user_id=user_id, + chat_id=chat_id, + summary=summary, + embedding=embedding, + created_at=time.time(), + ) ) - await session.execute(stmt) - await session.commit() - - -async def save_conversation_summary(user_id: int, chat_id: int, summary: str, embedding: str | None = None) -> None: - async with async_session() as session: - session.add(ConversationSummary( - user_id=user_id, - chat_id=chat_id, - summary=summary, - embedding=embedding, - created_at=time.time(), - )) await session.commit() async def get_summaries(user_id: int, chat_id: int, limit: int = 5) -> list[dict]: async with async_session() as session: - from sqlalchemy import select - stmt = ( select(ConversationSummary) - .where(ConversationSummary.user_id == user_id, ConversationSummary.chat_id == chat_id) + .where( + ConversationSummary.user_id == user_id, + ConversationSummary.chat_id == chat_id, + ) .order_by(ConversationSummary.created_at.desc()) .limit(limit) ) result = await session.execute(stmt) return [ - {"id": s.id, "summary": s.summary, "embedding": s.embedding, "created_at": s.created_at} + { + "id": s.id, + "summary": s.summary, + "embedding": s.embedding, + "created_at": s.created_at, + } for s in result.scalars().all() ] - - -async def unban_user_stickers(user_id: int, chat_id: int) -> bool: - async with async_session() as session: - from sqlalchemy import select - - stmt = select(StickerBan).where( - StickerBan.user_id == user_id, - StickerBan.chat_id == chat_id, - ) - result = await session.execute(stmt) - row = result.scalar_one_or_none() - - if row: - row.ban_until = None - row.count = 0 - row.ban_trigger = 0 - await session.commit() - return True - return False diff --git a/bot/utils/memory.py b/bot/utils/memory.py index 4fbe007..b133833 100644 --- a/bot/utils/memory.py +++ b/bot/utils/memory.py @@ -3,8 +3,11 @@ import logging import math import re -from bot.utils.ai_client import _get_connector +import aiohttp + +from bot.utils.proxy import get_proxy_connector from bot.utils.database import get_summaries, save_conversation_summary +from bot.utils.ai_client import get_client_timeout, ask_ai_simple from config import ROUTERAI_API_KEY, ROUTERAI_BASE_URL logger = logging.getLogger(__name__) @@ -29,26 +32,36 @@ async def create_embedding(text: str) -> list[float] | None: "encoding_format": "float", } - import aiohttp - from bot.utils.ai_client import get_client_timeout - - connector = _get_connector() + connector = get_proxy_connector() try: - async with aiohttp.ClientSession(timeout=get_client_timeout(30), connector=connector) as session: + async with aiohttp.ClientSession( + timeout=get_client_timeout(30), connector=connector + ) as session: async with session.post(EMBEDDING_URL, json=payload, headers=headers) as response: if response.status != 200: error_body = await response.text() - logger.warning("Embedding error | status=%s error=%s", response.status, error_body[:200]) + logger.warning( + "Embedding error | status=%s error=%s", + response.status, + error_body[:200], + ) return None data = await response.json() embedding = data["data"][0]["embedding"] - logger.info("Embedding created | dim=%d input_len=%d", len(embedding), min(len(text), 8000)) + logger.info( + "Embedding created | dim=%d input_len=%d", + len(embedding), + min(len(text), 8000), + ) return embedding - except Exception as e: - logger.error("Embedding request error: %s", e) + except Exception as exc: + logger.error("Embedding request error: %s", exc) return None + finally: + if connector: + await connector.close() def cosine_similarity(a: list[float], b: list[float]) -> float: @@ -60,14 +73,18 @@ def cosine_similarity(a: list[float], b: list[float]) -> float: return dot / (norm_a * norm_b) -async def find_relevant_summaries(user_id: int, chat_id: int, query: str, top_k: int = 3) -> list[str]: +async def find_relevant_summaries( + user_id: int, chat_id: int, query: str, top_k: int = 3 +) -> list[str]: query_emb = await create_embedding(query) if not query_emb: return [] summaries = await get_summaries(user_id, chat_id, limit=20) if not summaries: - logger.info("Relevant summaries | user=%d chat=%d no summaries found", user_id, chat_id) + logger.info( + "Relevant summaries | user=%d chat=%d no summaries found", user_id, chat_id + ) return [] scored = [] @@ -87,24 +104,30 @@ async def find_relevant_summaries(user_id: int, chat_id: int, query: str, top_k: logger.info( "Relevant summaries | user=%d chat=%d found=%d top_score=%.3f", - user_id, chat_id, len(top), top_score, + user_id, + chat_id, + len(top), + top_score, ) return top -async def save_summary_with_embedding(user_id: int, chat_id: int, summary_text: str) -> None: +async def save_summary_with_embedding( + user_id: int, chat_id: int, summary_text: str +) -> None: emb = await create_embedding(summary_text) embedding_json = json.dumps(emb) if emb else None await save_conversation_summary(user_id, chat_id, summary_text, embedding_json) logger.info( "Summary saved | user=%d chat=%d summary_len=%d emb=%s", - user_id, chat_id, len(summary_text), "yes" if emb else "no", + user_id, + chat_id, + len(summary_text), + "yes" if emb else "no", ) async def generate_summary(messages: list[dict]) -> str | None: - from bot.utils.ai_client import ask_ai_simple - messages_text = "\n".join( f"{'Пользователь' if m['role'] == 'user' else 'Астра'}: {m['content'][:300]}" for m in messages[-50:] @@ -127,7 +150,7 @@ async def generate_summary(messages: list[dict]) -> str | None: logger.warning("Summary too short | len=%d", len(summary)) else: logger.warning("Summary generation returned None") - except Exception as e: - logger.error("Summary generation error: %s", e) + except Exception as exc: + logger.error("Summary generation error: %s", exc) return None diff --git a/bot/utils/moderation.py b/bot/utils/moderation.py deleted file mode 100644 index d5d6789..0000000 --- a/bot/utils/moderation.py +++ /dev/null @@ -1,41 +0,0 @@ -import time -from collections import defaultdict - - -class ModerationManager: - def __init__(self, limit: int, window: int, ban_duration: int): - self.limit = limit - self.window = window - self.ban_duration = ban_duration - self.counters: dict[int, list[float]] = defaultdict(list) - self.bans: dict[int, float] = {} - - def add_message(self, user_id: int) -> None: - now = time.time() - self.counters[user_id].append(now) - self.counters[user_id] = [ - t for t in self.counters[user_id] if now - t <= self.window - ] - - def is_banned(self, user_id: int) -> bool: - if user_id in self.bans: - if time.time() - self.bans[user_id] < self.ban_duration: - return True - del self.bans[user_id] - return False - - def check_and_ban(self, user_id: int) -> bool: - if len(self.counters.get(user_id, [])) >= self.limit: - self.bans[user_id] = time.time() - return True - return False - - def should_delete(self, user_id: int) -> bool: - return self.is_banned(user_id) - - -moderation = ModerationManager( - limit=25, - window=60, - ban_duration=300, -) diff --git a/bot/utils/proxy.py b/bot/utils/proxy.py new file mode 100644 index 0000000..a8849e9 --- /dev/null +++ b/bot/utils/proxy.py @@ -0,0 +1,41 @@ +import logging +from urllib.parse import urlparse + +from aiohttp_socks import ProxyConnector, ProxyType + +from config import PROXY_ENABLED, PROXY_URL + +logger = logging.getLogger(__name__) + + +def get_proxy_connector() -> ProxyConnector | None: + """Create an aiohttp SOCKS5 proxy connector from PROXY_URL if enabled.""" + if not PROXY_ENABLED or not PROXY_URL: + return None + + parsed = urlparse(PROXY_URL) + + if parsed.scheme not in ("socks5", "socks5h"): + logger.warning("Unsupported proxy scheme: %s", parsed.scheme) + return None + + host = parsed.hostname + port = parsed.port + if not host or not port: + logger.warning("Invalid proxy URL: missing host or port") + return None + + username = parsed.username + password = parsed.password + + try: + return ProxyConnector( + proxy_type=ProxyType.SOCKS5, + host=host, + port=port, + username=username, + password=password, + ) + except Exception as exc: + logger.error("Failed to create proxy connector: %s", exc) + return None diff --git a/bot/utils/s3_client.py b/bot/utils/s3_client.py index bd8b960..afe356c 100644 --- a/bot/utils/s3_client.py +++ b/bot/utils/s3_client.py @@ -1,5 +1,6 @@ import os import logging +from functools import lru_cache import boto3 from botocore.config import Config @@ -10,6 +11,8 @@ logger = logging.getLogger("s3_client") ENDPOINT_URL = "https://storage.yandexcloud.net" + +@lru_cache(maxsize=1) def _build_s3_client(): return boto3.client( "s3", @@ -23,6 +26,7 @@ def _build_s3_client(): ), ) + def upload_file(file_path: str) -> str | None: key_name = os.path.basename(file_path) try: diff --git a/bot/utils/voice.py b/bot/utils/voice.py index ff66232..00e83b6 100644 --- a/bot/utils/voice.py +++ b/bot/utils/voice.py @@ -11,83 +11,101 @@ from config import PROXY_ENABLED, PROXY_URL logger = logging.getLogger(__name__) -_model = None +_model: WhisperModel | None = None +_model_lock = asyncio.Lock() BASE_DIR = Path(__file__).resolve().parent.parent.parent MODELS_DIR = BASE_DIR / "models" / "faster-whisper" VOICE_DIR = BASE_DIR / "bot" / "data" / "voice" -def _setup_proxy(): +def _setup_proxy() -> None: if PROXY_ENABLED and PROXY_URL: os.environ["ALL_PROXY"] = PROXY_URL - logger.info("Proxy set for model download: %s", PROXY_URL[:30] + "...") + logger.info("Proxy set for model download: %s...", PROXY_URL[:30]) -def _get_model() -> WhisperModel: +def _load_model() -> WhisperModel: + """Synchronous model load/download. Must run in a thread.""" + MODELS_DIR.mkdir(parents=True, exist_ok=True) + VOICE_DIR.mkdir(parents=True, exist_ok=True) + + model_path = MODELS_DIR / "base" + model_bin = model_path / "model.bin" + + if not model_path.is_dir() or not model_bin.is_file(): + _setup_proxy() + logger.info("Model not found locally, downloading to %s...", MODELS_DIR) + return WhisperModel( + "base", + device="cpu", + cpu_threads=4, + compute_type="int8", + download_root=str(MODELS_DIR), + ) + + logger.info("Loading model from %s...", model_path) + return WhisperModel( + str(model_path), + device="cpu", + cpu_threads=4, + compute_type="int8", + ) + + +async def _get_model() -> WhisperModel: + """Thread-safe lazy initializer for the Whisper model.""" global _model if _model is None: - MODELS_DIR.mkdir(parents=True, exist_ok=True) - VOICE_DIR.mkdir(parents=True, exist_ok=True) - model_path = str(MODELS_DIR / "base") - if not os.path.isdir(model_path) or not os.path.isfile(os.path.join(model_path, "model.bin")): - _setup_proxy() - logger.info("Model not found locally, downloading to %s...", MODELS_DIR) - _model = WhisperModel( - "base", - device="cpu", - cpu_threads=4, - compute_type="int8", - download_root=str(MODELS_DIR), - ) - else: - logger.info("Loading model from %s...", model_path) - _model = WhisperModel( - model_path, - device="cpu", - cpu_threads=4, - compute_type="int8", - ) - logger.info("Whisper model loaded") + async with _model_lock: + if _model is None: + loop = asyncio.get_running_loop() + _model = await loop.run_in_executor(None, _load_model) + logger.info("Whisper model loaded") return _model async def download_voice(bot: Bot, file_id: str) -> str: VOICE_DIR.mkdir(parents=True, exist_ok=True) file = await bot.get_file(file_id) - path = str(VOICE_DIR / f"{file_id}.ogg") + safe_id = re.sub(r"[^A-Za-z0-9_-]", "_", file_id) + path = str(VOICE_DIR / f"{safe_id}.ogg") await bot.download_file(file.file_path, destination=path) - logger.info("Voice downloaded: %s -> %s", file_id, path) + logger.info("Voice downloaded: %s", path) return path async def convert_to_wav(input_path: str) -> str: - output_path = input_path.replace(".ogg", ".wav") + output_path = str(Path(input_path).with_suffix(".wav")) proc = await asyncio.create_subprocess_exec( - "ffmpeg", "-i", input_path, "-ar", "16000", "-ac", "1", - "-c:a", "pcm_s16le", output_path, "-y", + "ffmpeg", + "-i", input_path, + "-ar", "16000", + "-ac", "1", + "-c:a", "pcm_s16le", + output_path, + "-y", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) _, stderr = await proc.communicate() if proc.returncode != 0: - logger.error(f"ffmpeg error: {stderr.decode(errors='replace')}") + logger.error("ffmpeg error: %s", stderr.decode(errors="replace")) raise RuntimeError("ffmpeg conversion failed") - logger.info(f"Converted to WAV: {output_path}") + logger.info("Converted to WAV: %s", output_path) return output_path async def transcribe_audio(file_path: str) -> str: - model = _get_model() + model = await _get_model() loop = asyncio.get_running_loop() - def _transcribe(): - segments, info = model.transcribe(file_path, language="ru", beam_size=5) - text = " ".join(seg.text for seg in segments) - return text.strip() + def _transcribe() -> str: + segments, _info = model.transcribe(file_path, language="ru", beam_size=5) + return " ".join(seg.text for seg in segments).strip() text = await loop.run_in_executor(None, _transcribe) - logger.info(f"Transcription result ({len(text)} chars): {text[:100]}...") + logger.info("Transcription result (%d chars): %s...", len(text), text[:100]) return text @@ -102,11 +120,11 @@ def normalize_text(text: str) -> str: return text -async def cleanup_files(*paths: str): +async def cleanup_files(*paths: str | None) -> None: for path in paths: if path and os.path.exists(path): try: - os.unlink(path) - logger.debug(f"Cleaned up: {path}") - except Exception as e: - logger.warning(f"Cleanup failed for {path}: {e}") + await asyncio.to_thread(os.unlink, path) + logger.debug("Cleaned up: %s", path) + except Exception as exc: + logger.warning("Cleanup failed for %s: %s", path, exc) diff --git a/bot/utils/weather.py b/bot/utils/weather.py index a550a06..fb21e7c 100644 --- a/bot/utils/weather.py +++ b/bot/utils/weather.py @@ -2,10 +2,9 @@ import logging from datetime import datetime import aiohttp -from aiohttp_socks import ProxyConnector, ProxyType -from config import API_WEATHER, PROXY_ENABLED, PROXY_URL -from bot.bot import bot +from bot.utils.proxy import get_proxy_connector +from config import API_WEATHER logger = logging.getLogger(__name__) @@ -28,47 +27,26 @@ def get_wind_direction(deg: float) -> str: return "Северо-запад" -def _get_connector(): - if PROXY_ENABLED and PROXY_URL: - parsed = PROXY_URL.replace("socks5://", "").replace("socks5h://", "") - if "@" in parsed: - auth, host_port = parsed.split("@", 1) - username, password = auth.split(":", 1) - else: - username = None - password = None - host_port = parsed - - host, port = host_port.rsplit(":", 1) - port = int(port) - - return ProxyConnector( - proxy_type=ProxyType.SOCKS5, - host=host, - port=port, - username=username, - password=password, - ) - return None - - async def get_weather(message) -> None: - if len(message.text.split(maxsplit=1)) == 1: - await bot.send_message(message.chat.id, "Пожалуйста, укажите город.") + text_parts = message.text.split(maxsplit=1) + if len(text_parts) == 1: + await message.answer("Пожалуйста, укажите город.") return - city = message.text.split(maxsplit=1)[1].strip() + city = text_parts[1].strip() url = "https://api.openweathermap.org/data/2.5/weather" params = {"q": city, "appid": API_WEATHER, "units": "metric", "lang": "ru"} - timeout = aiohttp.ClientTimeout(total=30, sock_connect=15, sock_read=15) - connector = _get_connector() + connector = get_proxy_connector() try: - async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=30, sock_connect=15, sock_read=15), + connector=connector, + ) as session: async with session.get(url, params=params) as response: if response.status == 404: - await bot.send_message(message.chat.id, "Город не найден. Пожалуйста, уточните запрос.") + await message.answer("Город не найден. Пожалуйста, уточните запрос.") return response.raise_for_status() weather_data = await response.json() @@ -103,16 +81,16 @@ async def get_weather(message) -> None: f"Видимость: {visibility} м\n" f"Восход: {sunrise_time}, закат: {sunset_time}" ) - await bot.send_message(message.chat.id, weather_message, parse_mode=None) + await message.answer(weather_message, parse_mode=None) except aiohttp.ClientResponseError as exc: logger.error("HTTP ошибка погоды: %s", exc) - await bot.send_message(message.chat.id, "При получении данных произошла ошибка, попробуйте еще раз.") + await message.answer("При получении данных произошла ошибка, попробуйте еще раз.") except aiohttp.ClientError as exc: logger.error("Ошибка запроса погоды: %s", exc) - await bot.send_message(message.chat.id, "Не удалось связаться с погодным сервисом.") + await message.answer("Не удалось связаться с погодным сервисом.") except Exception as exc: logger.exception("Неожиданная ошибка погоды: %s", exc) - await bot.send_message(message.chat.id, "Произошла ошибка при обработке погоды.") + await message.answer("Произошла ошибка при обработке погоды.") finally: if connector: await connector.close() diff --git a/bot/utils/yadisk_download.py b/bot/utils/yadisk_download.py index a3ae5b9..0e6739a 100644 --- a/bot/utils/yadisk_download.py +++ b/bot/utils/yadisk_download.py @@ -5,80 +5,64 @@ from pathlib import Path from urllib.parse import urlencode, unquote import aiohttp -from aiohttp_socks import ProxyConnector, ProxyType -from config import PROXY_URL, PROXY_ENABLED +from bot.utils.proxy import get_proxy_connector DOWNLOAD_DIR = Path("bot/data/downloads") DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) -def _get_connector(): - if PROXY_ENABLED and PROXY_URL: - parsed = PROXY_URL.replace("socks5://", "").replace("socks5h://", "") - if "@" in parsed: - auth, host_port = parsed.split("@", 1) - username, password = auth.split(":", 1) - else: - username = None - password = None - host_port = parsed - - host, port = host_port.rsplit(":", 1) - port = int(port) - - return ProxyConnector( - proxy_type=ProxyType.SOCKS5, - host=host, - port=port, - username=username, - password=password, - ) - return None def sanitize_filename(filename: str) -> str: filename = os.path.basename(filename).strip() filename = re.sub(r"[\\/:*?\"<>|]+", "_", filename) return filename or "downloaded_file" + async def download_yandex_file(public_url: str, progress_callback=None) -> str: base_url = "https://cloud-api.yandex.net/v1/disk/public/resources/download?" final_url = base_url + urlencode({"public_key": public_url}) timeout = aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=300) - connector = _get_connector() + connector = get_proxy_connector() - async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: - async with session.get(final_url) as response: - response.raise_for_status() - payload = await response.json() - download_url = payload["href"] + try: + async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: + async with session.get(final_url) as response: + response.raise_for_status() + payload = await response.json() + download_url = payload["href"] - async with session.get(download_url) as response: - response.raise_for_status() + async with session.get(download_url) as response: + response.raise_for_status() - content_disposition = response.headers.get("Content-Disposition", "") - filename = download_url.split("/")[-1] + content_disposition = response.headers.get("Content-Disposition", "") + filename = download_url.split("/")[-1] - if "filename*" in content_disposition: - try: - encoded = content_disposition.split("filename*=")[1].strip() - parts = encoded.split("''", 1) - if len(parts) == 2: - filename = unquote(parts[1], encoding=parts[0] or "utf-8") - except Exception: - pass + if "filename*" in content_disposition: + try: + encoded = content_disposition.split("filename*=")[1].strip() + parts = encoded.split("''", 1) + if len(parts) == 2: + filename = unquote(parts[1], encoding=parts[0] or "utf-8") + except Exception: + pass - filename = sanitize_filename(filename) - suffix = Path(filename).suffix or ".bin" + filename = sanitize_filename(filename) + suffix = Path(filename).suffix or ".bin" - with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, dir=str(DOWNLOAD_DIR)) as tmp_file: - total_size = int(response.headers.get("Content-Length", 0)) - downloaded_size = 0 + with tempfile.NamedTemporaryFile( + delete=False, suffix=suffix, dir=str(DOWNLOAD_DIR) + ) as tmp_file: + total_size = int(response.headers.get("Content-Length", 0)) + downloaded_size = 0 - async for chunk in response.content.iter_chunked(1024 * 64): - tmp_file.write(chunk) - downloaded_size += len(chunk) - if progress_callback: - await progress_callback(downloaded_size, total_size) + async for chunk in response.content.iter_chunked(1024 * 64): + tmp_file.write(chunk) + downloaded_size += len(chunk) + if progress_callback: + await progress_callback(downloaded_size, total_size) - return tmp_file.name + return tmp_file.name + finally: + if connector: + await connector.close() diff --git a/config.py b/config.py index 204e985..ad44a4b 100644 --- a/config.py +++ b/config.py @@ -23,6 +23,11 @@ AI_PHASE2_LIMIT = 20 AI_COOLDOWN = 3600 AI_BLOCK_DEFAULT_DURATION = 86400 +# Health check periodically calls every free model to test availability. +# This consumes API tokens; disabled by default. +AI_HEALTH_CHECK_ENABLED = os.getenv("AI_HEALTH_CHECK_ENABLED", "false").lower() == "true" +AI_HEALTH_CHECK_INTERVAL = int(os.getenv("AI_HEALTH_CHECK_INTERVAL", "600")) + ACCESS_KEY = os.getenv("ACCESS_KEY", "").strip() SECRET_KEY = os.getenv("SECRET_KEY", "").strip() BUCKET_NAME = os.getenv("BUCKET_NAME", "").strip() diff --git a/docker-compose.yml b/docker-compose.yml index 2e671c1..b6aa28f 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -6,3 +6,4 @@ services: env_file: .env volumes: - ./bot/data:/app/bot/data + - ./models:/app/models diff --git a/main.py b/main.py index b7119d2..2b4954e 100644 --- a/main.py +++ b/main.py @@ -19,6 +19,7 @@ from config import PROXY_ENABLED setup_logging() logger = logging.getLogger(__name__) +# Router order matters: moderation first so it can delete messages before others process them. dp.include_router(moderation.router) dp.include_router(layout.router) dp.include_router(weather.router) @@ -37,9 +38,9 @@ async def save_user_info(message: Message): message.from_user.username, message.from_user.full_name or "", ) - except Exception: - pass - raise SkipHandler() + except Exception as exc: + logger.warning("Failed to save chat user: %s", exc) + raise SkipHandler async def main(): @@ -47,31 +48,41 @@ async def main(): logging.info("База данных инициализирована") await setup_bot_commands() - asyncio.create_task(start_model_health_check()) + health_check_task = asyncio.create_task(start_model_health_check()) try: if PROXY_ENABLED: await setup_proxy() else: me = await bot.me() - logging.info(f"Бот запущен без прокси: @{me.username}") - except TelegramNetworkError as e: + logging.info("Бот запущен без прокси: @%s", me.username) + except TelegramNetworkError: if PROXY_ENABLED: logging.error("Не могу соединиться с прокси сервером. Попробуй другой прокси.") else: logging.error("Не могу соединиться с Telegram. Проверь интернет-соединение.") + health_check_task.cancel() return - except ClientConnectionError as e: + except ClientConnectionError: if PROXY_ENABLED: logging.error("Не могу соединиться с прокси сервером. Попробуй другой прокси.") else: logging.error("Не могу соединиться с Telegram. Проверь интернет-соединение.") + health_check_task.cancel() return - except Exception as e: - logging.error(f"Ошибка при запуске бота: {e}") + except Exception as exc: + logging.error("Ошибка при запуске бота: %s", exc) + health_check_task.cancel() return - await dp.start_polling(bot) + try: + await dp.start_polling(bot) + finally: + health_check_task.cancel() + try: + await health_check_task + except asyncio.CancelledError: + pass if __name__ == "__main__": diff --git a/requirements.txt b/requirements.txt index 6cea768..6d0e6e5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,8 @@ -aiogram>=3.0.0 -python-dotenv>=1.0.0 -aiohttp-socks>=0.8.0 -sqlalchemy>=2.0.0 -aiosqlite>=0.19.0 -faster-whisper>=1.1.0 -httpx[socks] -boto3>=1.34.0 +aiogram==3.20.0 +python-dotenv==1.1.0 +aiohttp==3.12.4 +aiohttp-socks==0.10.1 +sqlalchemy==2.0.40 +aiosqlite==0.21.0 +faster-whisper==1.1.1 +boto3==1.37.25 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_layout_converter.py b/tests/test_layout_converter.py new file mode 100644 index 0000000..9a3b644 --- /dev/null +++ b/tests/test_layout_converter.py @@ -0,0 +1,23 @@ +import pytest + +from bot.utils.layout_converter import convert_layout + + +@pytest.mark.parametrize( + "input_text, expected", + [ + ("gbdj", "пиво"), + ("rjytxyj!", "конечно!"), + ("Ghbdtn", "Привет"), + ("123", "123"), + ("", ""), + ("qwerty", "йцукен"), + ], +) +def test_convert_layout(input_text: str, expected: str) -> None: + assert convert_layout(input_text) == expected + + +def test_convert_layout_preserves_case() -> None: + assert convert_layout("QWERTY") == "ЙЦУКЕН" + assert convert_layout("QwErTy") == "ЙцУкЕн" diff --git a/tests/test_voice.py b/tests/test_voice.py new file mode 100644 index 0000000..74fbf03 --- /dev/null +++ b/tests/test_voice.py @@ -0,0 +1,17 @@ +import pytest + +from bot.utils.voice import normalize_text + + +@pytest.mark.parametrize( + "input_text, expected", + [ + ("привет", "Привет."), + ("Привет!", "Привет!"), + (" hello world ", "Hello world."), + ("", ""), + ("как дела", "Как дела."), + ], +) +def test_normalize_text(input_text: str, expected: str) -> None: + assert normalize_text(input_text) == expected