Initial commit: UMB Telegram Bot
This commit is contained in:
Галингер Р.С.
@@ -0,0 +1,260 @@
|
||||
import logging
|
||||
|
||||
from aiogram import Router, F
|
||||
from aiogram.types import Message
|
||||
from aiogram.filters import Command
|
||||
|
||||
from bot.utils.ai_client import ask_ai
|
||||
from bot.utils.database import (
|
||||
add_context_message,
|
||||
get_user_context,
|
||||
is_ai_blocked,
|
||||
block_user_from_ai,
|
||||
unblock_user_from_ai,
|
||||
get_chat_users,
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = Router()
|
||||
|
||||
|
||||
def extract_username(text: str) -> str | None:
|
||||
parts = text.split()
|
||||
for part in parts:
|
||||
if part.startswith("@"):
|
||||
return part[1:]
|
||||
if part.startswith("id") and part[2:].isdigit():
|
||||
return part[2:]
|
||||
return None
|
||||
|
||||
|
||||
async def get_user_id_by_username(message: Message, username: str) -> int | None:
|
||||
try:
|
||||
username_clean = username.lstrip("@")
|
||||
admins = await message.bot.get_chat_administrators(message.chat.id)
|
||||
for admin in admins:
|
||||
if admin.user.username and admin.user.username.lower() == username_clean.lower():
|
||||
return admin.user.id
|
||||
if str(admin.user.id) == username_clean:
|
||||
return admin.user.id
|
||||
|
||||
if username_clean.isdigit():
|
||||
return int(username_clean)
|
||||
except Exception as e:
|
||||
logger.error("Error resolving username %s: %s", username, e)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
async def _is_creator(message: Message) -> bool:
|
||||
if message.from_user.id == message.chat.id:
|
||||
return True
|
||||
try:
|
||||
admins = await message.bot.get_chat_administrators(message.chat.id)
|
||||
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)
|
||||
return False
|
||||
|
||||
|
||||
async def _not_creator(message: Message) -> bool:
|
||||
return not await _is_creator(message)
|
||||
|
||||
|
||||
@router.message(Command("ai"))
|
||||
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 "")
|
||||
|
||||
if await is_ai_blocked(user_id, chat_id):
|
||||
await message.answer("Тебе временно недоступен AI. Обратись к владельцу чата.", parse_mode=None)
|
||||
return
|
||||
|
||||
text_parts = message.text.split(maxsplit=1)
|
||||
has_direct_question = len(text_parts) > 1 and text_parts[1].strip()
|
||||
|
||||
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)
|
||||
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
|
||||
prompt = f"Проанализируй это сообщение:\n\n{target_text}"
|
||||
context_text = target_text
|
||||
elif has_direct_question:
|
||||
prompt = text_parts[1].strip()
|
||||
context_text = prompt
|
||||
else:
|
||||
await _start_dialogue(message)
|
||||
return
|
||||
|
||||
status_msg = await message.reply("Думаю...", parse_mode=None)
|
||||
|
||||
context = await get_user_context(user_id, chat_id, AI_CONTEXT_LIMIT)
|
||||
|
||||
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)
|
||||
|
||||
async def update_status(text: str):
|
||||
try:
|
||||
await status_msg.edit_text(text, parse_mode=None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
response = await ask_ai(prompt, context, status_callback=update_status)
|
||||
|
||||
if not response:
|
||||
response = "Не удалось получить ответ от AI."
|
||||
|
||||
await add_context_message(user_id, chat_id, context_text, "user")
|
||||
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)
|
||||
except Exception:
|
||||
await status_msg.edit_text(response, parse_mode=None, disable_web_page_preview=True)
|
||||
|
||||
|
||||
async def _start_dialogue(message: Message):
|
||||
user_id = message.from_user.id
|
||||
chat_id = message.chat.id
|
||||
|
||||
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
|
||||
await message.answer(
|
||||
f"Твой лимит диалога исчерпан. Попробуй через {mins} мин.",
|
||||
parse_mode=None,
|
||||
)
|
||||
return
|
||||
|
||||
if dialogue["phase"] == 2 and dialogue["msg_count"] >= AI_PHASE2_LIMIT:
|
||||
await message.answer("Твой лимит исчерпан.", parse_mode=None)
|
||||
return
|
||||
|
||||
await message.answer(
|
||||
"💬 Диалог начат! Отвечай на мои сообщения, чтобы продолжать.\n"
|
||||
"Отправь /aiclear чтобы завершить.",
|
||||
parse_mode=None,
|
||||
)
|
||||
|
||||
|
||||
@router.message(Command("aiuser"))
|
||||
async def cmd_aiuser(message: Message):
|
||||
if await _not_creator(message):
|
||||
return
|
||||
|
||||
try:
|
||||
users = await get_chat_users(message.chat.id)
|
||||
|
||||
lines = []
|
||||
for i, u in enumerate(users, 1):
|
||||
username_str = f"@{u['username']}" if u.get("username") else "NO DATA"
|
||||
lines.append(f"{i}. {u['full_name']} | {username_str} | {u['user_id']}")
|
||||
|
||||
text = "\n".join(lines) if lines else "Нет данных о пользователях."
|
||||
|
||||
await message.bot.send_message(
|
||||
message.from_user.id,
|
||||
f"📋 Список пользователей чата ({len(users)}):\n\n{text}",
|
||||
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)
|
||||
|
||||
|
||||
@router.message(Command("aiclear"))
|
||||
async def cmd_aiclear(message: Message):
|
||||
user_id = message.from_user.id
|
||||
chat_id = message.chat.id
|
||||
|
||||
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)
|
||||
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))
|
||||
else:
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@router.message(Command("aino"))
|
||||
async def cmd_aino(message: Message):
|
||||
if await _not_creator(message):
|
||||
return
|
||||
|
||||
username = extract_username(message.text)
|
||||
if not username:
|
||||
await message.answer("Укажи пользователя: /aino @username или /aino id<UID>", 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)
|
||||
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)
|
||||
|
||||
|
||||
@router.message(Command("aiyes"))
|
||||
async def cmd_aiyes(message: Message):
|
||||
if await _not_creator(message):
|
||||
return
|
||||
|
||||
username = extract_username(message.text)
|
||||
if not username:
|
||||
await message.answer("Укажи пользователя: /aiyes @username или /aiyes id<UID>", 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)
|
||||
return
|
||||
|
||||
unblocked = await unblock_user_from_ai(target_user_id, message.chat.id)
|
||||
if unblocked:
|
||||
await message.answer(f"Пользователь {username} разблокирован для AI.", parse_mode=None)
|
||||
else:
|
||||
await message.answer(f"Пользователь {username} не был заблокирован от AI.", parse_mode=None)
|
||||
Reference in New Issue
Block a user