import logging 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, get_user_context, get_or_create_dialogue, increment_dialogue_count, reset_dialogue_to_phase2, block_dialogue, is_ai_blocked, save_chat_user, clear_user_context, ) 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__) 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 "", ) if await is_ai_blocked(user_id, chat_id): return if not await _is_astra_message(replied): return dialogue = await get_or_create_dialogue(user_id, chat_id) if not dialogue["is_active"]: return count_result = await increment_dialogue_count(user_id, chat_id) 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, ) limit = AI_DIALOGUE_LIMIT if phase == 1 else AI_PHASE2_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, ) 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, ) await _end_dialogue(message, user_id, chat_id) return 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 ) if relevant: extra_context = "\n\nИз прошлых диалогов:\n" + "\n---\n".join(relevant[:2]) except Exception as exc: logger.error("Error fetching relevant summaries: %s", exc) remaining = limit - current_count warning = "" if 0 < remaining <= 5: warning = f"\n\n⚠️ Осталось {remaining} сообщений в этом диалоге." system_extras = [] if extra_context: system_extras.append(extra_context) if warning: system_extras.append(warning) 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 as exc: logger.debug("Status update failed: %s", exc) response = await ask_ai( message.text or "", context_messages, status_callback=update_status, extra_system_content=extra_system_content, ) if not response: response = "Не могу ответить сейчас." await add_context_message(user_id, chat_id, message.text or "", "user") await add_context_message(user_id, chat_id, response, "assistant") try: await message.reply(response, parse_mode="HTML", disable_web_page_preview=True) except Exception: await message.reply(response, parse_mode=None, disable_web_page_preview=True) try: await status_msg.delete() except Exception as exc: logger.debug("Status delete failed: %s", exc) if 0 < remaining <= 3: try: await message.reply( f"⚠️ Осталось {remaining} сообщений. Память почти заполнена.", parse_mode=None, ) 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 ) try: context = await get_user_context(user_id, chat_id, limit=AI_DIALOGUE_LIMIT) if context and len(context) >= 4: summary = await generate_summary(context) if summary: await save_summary_with_embedding(user_id, chat_id, summary) await clear_user_context(user_id, chat_id) 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 as exc: logger.debug("Status edit failed: %s", exc) async def _end_dialogue(message: Message, user_id: int, chat_id: int): context = await get_user_context(user_id, chat_id, limit=AI_PHASE2_LIMIT) if context and len(context) >= 4: try: summary = await generate_summary(context) if summary: await save_summary_with_embedding(user_id, chat_id, summary) 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, )