umb/bot/routers/dialogue.py

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,
    )