Initial commit: UMB Telegram Bot
This commit is contained in:
Галингер Р.С.
@@ -0,0 +1,170 @@
|
||||
import logging
|
||||
import time
|
||||
|
||||
from aiogram import Router, F
|
||||
from aiogram.types import Message
|
||||
|
||||
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()
|
||||
|
||||
|
||||
@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
|
||||
|
||||
replied_from_bot = replied.from_user and replied.from_user.is_bot
|
||||
|
||||
if not replied_from_bot:
|
||||
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)
|
||||
|
||||
if phase == 1:
|
||||
limit = AI_DIALOGUE_LIMIT
|
||||
else:
|
||||
limit = AI_PHASE2_LIMIT
|
||||
|
||||
if current_count > limit:
|
||||
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
|
||||
|
||||
if phase == 1:
|
||||
context_limit = AI_DIALOGUE_LIMIT
|
||||
else:
|
||||
context_limit = AI_PHASE2_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 e:
|
||||
logger.error("Error fetching relevant summaries: %s", e)
|
||||
|
||||
remaining = limit - current_count
|
||||
warning = ""
|
||||
if remaining <= 5:
|
||||
warning = f"\n\n⚠️ Осталось {remaining} сообщений в этом диалоге."
|
||||
|
||||
system_prompt_extra = ""
|
||||
if extra_context:
|
||||
system_prompt_extra += extra_context
|
||||
if warning:
|
||||
system_prompt_extra += warning
|
||||
|
||||
message_text = message.text or ""
|
||||
if system_prompt_extra:
|
||||
message_text += "\n\n(Контекст)" + system_prompt_extra
|
||||
|
||||
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
|
||||
|
||||
response = await ask_ai(
|
||||
message_text,
|
||||
context_messages,
|
||||
status_callback=update_status,
|
||||
)
|
||||
|
||||
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:
|
||||
pass
|
||||
|
||||
if remaining <= 3 and remaining > 0:
|
||||
try:
|
||||
warn_msg = await message.reply(
|
||||
f"⚠️ Осталось {remaining} сообщений. Память почти заполнена.",
|
||||
parse_mode=None,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
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 e:
|
||||
logger.error("Error saving summary: %s", e)
|
||||
|
||||
await reset_dialogue_to_phase2(user_id, chat_id)
|
||||
|
||||
try:
|
||||
await status_msg.edit_text("✅ Начинаю новую сессию (осталось 20 сообщений).", parse_mode=None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
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 e:
|
||||
logger.error("Error saving final summary: %s", e)
|
||||
|
||||
await block_dialogue(user_id, chat_id, AI_COOLDOWN)
|
||||
|
||||
await message.reply("Твой лимит исчерпан. Возвращайся через час.", parse_mode=None)
|
||||
Reference in New Issue
Block a user