1699 lines
64 KiB
Python
1699 lines
64 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import time
|
||
import uuid
|
||
from datetime import datetime, timedelta, timezone
|
||
from datetime import time as dt_time
|
||
from pathlib import Path
|
||
|
||
import httpx
|
||
from telegram import (
|
||
BotCommand,
|
||
CopyTextButton,
|
||
InlineKeyboardButton,
|
||
InlineKeyboardMarkup,
|
||
InputFile,
|
||
LabeledPrice,
|
||
MenuButtonCommands,
|
||
Update,
|
||
)
|
||
from telegram.constants import ChatAction, ChatType
|
||
from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut
|
||
from telegram.ext import (
|
||
Application,
|
||
CallbackQueryHandler,
|
||
CommandHandler,
|
||
ContextTypes,
|
||
MessageHandler,
|
||
PreCheckoutQueryHandler,
|
||
TypeHandler,
|
||
filters,
|
||
)
|
||
|
||
from comfy_client import ComfyClient
|
||
from config import Settings, setup_logging
|
||
from db import Database, Request, UserAlreadyQueued, UserContext
|
||
from moderation import PromptBlocked, detect_person, is_prompt_allowed
|
||
from prompt_enhancer import PromptEnhancer
|
||
from queue_manager import QueueEvent, QueueFull, QueueManager
|
||
from workflows import WorkflowBuilder
|
||
|
||
logger = logging.getLogger(__name__)
|
||
MSK = timezone(timedelta(hours=3))
|
||
POLLING_WATCHDOG_INTERVAL = 60.0
|
||
POLLING_WATCHDOG_STUCK_AFTER = 300.0
|
||
RESTART_HINT = "/restart — сброс контекста"
|
||
CAPTION_LIMIT = 1024
|
||
|
||
|
||
def commands_need_update(
|
||
current: tuple[BotCommand, ...] | list[BotCommand],
|
||
desired: tuple[BotCommand, ...] | list[BotCommand],
|
||
) -> bool:
|
||
current_pairs = [
|
||
(command.command, command.description) for command in current
|
||
]
|
||
desired_pairs = [
|
||
(command.command, command.description) for command in desired
|
||
]
|
||
return current_pairs != desired_pairs
|
||
|
||
|
||
def format_context(context: UserContext | None) -> str:
|
||
if context is None:
|
||
return "Контекст пуст. Следующий запрос начнёт новую сцену."
|
||
|
||
parts = ["Текущий контекст:", "", f"Промпт: {context.current_prompt}"]
|
||
if (
|
||
context.base_prompt
|
||
and context.base_prompt != context.current_prompt
|
||
):
|
||
parts.extend(["", f"Базовая сцена: {context.base_prompt}"])
|
||
if context.last_request_id is not None:
|
||
mode = f" ({context.last_mode})" if context.last_mode else ""
|
||
parts.append(
|
||
f"Последний запрос: #{context.last_request_id}{mode}"
|
||
)
|
||
if context.updated_at:
|
||
try:
|
||
updated = datetime.fromisoformat(
|
||
context.updated_at
|
||
).astimezone(MSK)
|
||
except ValueError:
|
||
pass
|
||
else:
|
||
parts.append(f"Обновлён: {updated.strftime('%d.%m.%Y %H:%M')}")
|
||
parts.extend(["", "Сбросить: /restart"])
|
||
return "\n".join(parts)
|
||
|
||
|
||
def result_caption(prompt: str, *, limit: int = CAPTION_LIMIT) -> str:
|
||
value = (prompt or "").strip()
|
||
suffix = f"\n\n{RESTART_HINT}"
|
||
available = max(0, limit - len(suffix))
|
||
if len(value) > available:
|
||
value = value[: max(0, available - 1)].rstrip()
|
||
if value:
|
||
value += "…"
|
||
return f"{value}{suffix}" if value else RESTART_HINT
|
||
|
||
|
||
def polling_is_stuck(
|
||
*,
|
||
now: float,
|
||
last_update_at: float,
|
||
pending_update_count: int,
|
||
stuck_after: float,
|
||
) -> bool:
|
||
return (
|
||
pending_update_count > 0
|
||
and now - last_update_at >= stuck_after
|
||
)
|
||
|
||
START_TEXT = (
|
||
"Пришли текст — сгенерирую картинку.\n"
|
||
"Пришли фото с подписью — отредактирую его по описанию.\n"
|
||
"Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
|
||
"/balance — баланс генераций.\n"
|
||
"/invite — пригласить друга.\n"
|
||
"/status — состояние текущей задачи.\n"
|
||
"/restart — сброс контекста.\n"
|
||
"/context — текущий контекст.\n"
|
||
"/settings — качество и шаги Krea2."
|
||
)
|
||
BOT_COMMANDS = (
|
||
BotCommand("start", "Запуск и баланс"),
|
||
BotCommand("help", "Инструкция"),
|
||
BotCommand("balance", "Баланс генераций"),
|
||
BotCommand("invite", "Пригласить друга"),
|
||
BotCommand("status", "Статус текущей задачи"),
|
||
BotCommand("context", "Текущий контекст"),
|
||
BotCommand("settings", "Качество и шаги Krea2"),
|
||
BotCommand("restart", "Сброс контекста"),
|
||
)
|
||
BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
|
||
AGE_TEXT = (
|
||
"🔞 Бот доступен только пользователям 18+.\n\n"
|
||
"Подтверди, что тебе исполнилось 18 лет."
|
||
)
|
||
AGE_DENIED = "Доступ к боту разрешён только с 18 лет."
|
||
|
||
|
||
class MpakoBot:
|
||
def __init__(self, settings: Settings) -> None:
|
||
self._settings = settings
|
||
self._db = Database(
|
||
settings.database_url,
|
||
min_size=settings.db_pool_min_size,
|
||
max_size=settings.db_pool_max_size,
|
||
)
|
||
self._workflows: WorkflowBuilder | None = None
|
||
self._comfy: ComfyClient | None = None
|
||
self._queue: QueueManager | None = None
|
||
self._bot = None
|
||
self._status_messages: dict[int, tuple[int, int]] = {}
|
||
self._last_edit: dict[int, tuple[float, int]] = {}
|
||
self._daily_task: asyncio.Task[None] | None = None
|
||
self._health_task: asyncio.Task[None] | None = None
|
||
self._watchdog_task: asyncio.Task[None] | None = None
|
||
self._last_update_at = time.monotonic()
|
||
self._application: Application | None = None
|
||
|
||
def run(self) -> None:
|
||
builder = (
|
||
Application.builder()
|
||
.token(self._settings.bot_token)
|
||
.connect_timeout(10)
|
||
.read_timeout(30)
|
||
.write_timeout(30)
|
||
.pool_timeout(5)
|
||
.connection_pool_size(10)
|
||
.get_updates_connect_timeout(10)
|
||
.get_updates_read_timeout(15)
|
||
.get_updates_write_timeout(15)
|
||
.get_updates_pool_timeout(5)
|
||
.get_updates_connection_pool_size(10)
|
||
.post_init(self._post_init)
|
||
.post_shutdown(self._post_shutdown)
|
||
)
|
||
if self._settings.telegram_proxy:
|
||
builder = builder.proxy(self._settings.telegram_proxy).get_updates_proxy(
|
||
self._settings.telegram_proxy
|
||
)
|
||
application = builder.build()
|
||
self._application = application
|
||
application.add_handler(
|
||
TypeHandler(Update, self._track_update), group=-1
|
||
)
|
||
application.add_handler(CommandHandler("start", self._start))
|
||
application.add_handler(CommandHandler("help", self._help))
|
||
application.add_handler(CommandHandler("status", self._status))
|
||
application.add_handler(CommandHandler("balance", self._balance))
|
||
application.add_handler(CommandHandler("invite", self._invite))
|
||
application.add_handler(
|
||
CommandHandler("context", self._context_command)
|
||
)
|
||
application.add_handler(
|
||
CommandHandler("settings", self._settings_command)
|
||
)
|
||
application.add_handler(CommandHandler("restart", self._reset))
|
||
application.add_handler(CommandHandler("reset", self._reset))
|
||
application.add_handler(
|
||
PreCheckoutQueryHandler(self._pre_checkout)
|
||
)
|
||
application.add_handler(
|
||
CallbackQueryHandler(self._queue_choice, pattern=r"^queue:")
|
||
)
|
||
application.add_handler(
|
||
CallbackQueryHandler(self._age_choice, pattern=r"^age:")
|
||
)
|
||
application.add_handler(
|
||
CallbackQueryHandler(
|
||
self._quality_choice, pattern=r"^quality:"
|
||
)
|
||
)
|
||
application.add_handler(
|
||
CallbackQueryHandler(self._steps_choice, pattern=r"^steps:")
|
||
)
|
||
application.add_handler(
|
||
MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment)
|
||
)
|
||
application.add_handler(MessageHandler(filters.PHOTO, self._photo))
|
||
application.add_handler(
|
||
MessageHandler(filters.TEXT & ~filters.COMMAND, self._text)
|
||
)
|
||
application.add_error_handler(self._on_error)
|
||
application.run_polling(allowed_updates=Update.ALL_TYPES)
|
||
|
||
async def _post_init(self, application: Application) -> None:
|
||
self._bot = application.bot
|
||
await self._setup_commands()
|
||
self._settings.storage_dir.mkdir(parents=True, exist_ok=True)
|
||
self._workflows = WorkflowBuilder(
|
||
self._settings.workflow_t2i,
|
||
self._settings.workflow_i2i,
|
||
self._settings.workflow_mapping,
|
||
)
|
||
self._comfy = ComfyClient(
|
||
self._settings.comfy_base_url,
|
||
self._settings.comfy_ws_url,
|
||
)
|
||
await self._db.connect()
|
||
await self._db.seed_admins(self._settings.admin_user_ids)
|
||
await self._db.accrue_daily_credits()
|
||
self._daily_task = asyncio.create_task(
|
||
self._daily_credits_loop(), name="daily-credits"
|
||
)
|
||
self._health_task = asyncio.create_task(
|
||
self._telegram_health_loop(), name="telegram-health"
|
||
)
|
||
self._watchdog_task = asyncio.create_task(
|
||
self._polling_watchdog_loop(), name="polling-watchdog"
|
||
)
|
||
enhancer = (
|
||
PromptEnhancer(self._comfy, self._settings)
|
||
if self._settings.prompt_enhance_enabled
|
||
else None
|
||
)
|
||
self._queue = QueueManager(
|
||
db=self._db,
|
||
comfy=self._comfy,
|
||
workflows=self._workflows,
|
||
settings=self._settings,
|
||
on_event=self._on_queue_event,
|
||
enhancer=enhancer,
|
||
)
|
||
await self._queue.start()
|
||
logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url)
|
||
|
||
async def _post_shutdown(self, application: Application) -> None:
|
||
if self._watchdog_task is not None:
|
||
self._watchdog_task.cancel()
|
||
await asyncio.gather(self._watchdog_task, return_exceptions=True)
|
||
self._watchdog_task = None
|
||
if self._health_task is not None:
|
||
self._health_task.cancel()
|
||
await asyncio.gather(self._health_task, return_exceptions=True)
|
||
self._health_task = None
|
||
if self._daily_task is not None:
|
||
self._daily_task.cancel()
|
||
await asyncio.gather(self._daily_task, return_exceptions=True)
|
||
self._daily_task = None
|
||
if self._queue is not None:
|
||
await self._queue.stop()
|
||
if self._comfy is not None:
|
||
await self._comfy.close()
|
||
await self._db.close()
|
||
logger.info("Бот остановлен")
|
||
|
||
async def _setup_commands(self) -> None:
|
||
if self._bot is None:
|
||
return
|
||
try:
|
||
current = await self._retry(
|
||
self._bot.get_my_commands, attempts=3
|
||
)
|
||
if commands_need_update(current, BOT_COMMANDS):
|
||
await self._retry(
|
||
self._bot.set_my_commands, BOT_COMMANDS, attempts=3
|
||
)
|
||
logger.info("Команды бота обновлены")
|
||
menu = await self._retry(
|
||
self._bot.get_chat_menu_button, attempts=3
|
||
)
|
||
if not isinstance(menu, MenuButtonCommands):
|
||
await self._retry(
|
||
self._bot.set_chat_menu_button,
|
||
menu_button=MenuButtonCommands(),
|
||
attempts=3,
|
||
)
|
||
logger.info("Кнопка меню с командами включена")
|
||
except TelegramError:
|
||
logger.exception("Не удалось настроить команды бота")
|
||
|
||
async def _track_update(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
self._last_update_at = time.monotonic()
|
||
|
||
async def _polling_watchdog_loop(self) -> None:
|
||
while True:
|
||
await asyncio.sleep(POLLING_WATCHDOG_INTERVAL)
|
||
pending = await self._pending_updates_count()
|
||
if pending is None:
|
||
continue
|
||
if not polling_is_stuck(
|
||
now=time.monotonic(),
|
||
last_update_at=self._last_update_at,
|
||
pending_update_count=pending,
|
||
stuck_after=POLLING_WATCHDOG_STUCK_AFTER,
|
||
):
|
||
continue
|
||
logger.warning(
|
||
"Polling watchdog: %s обновлений не обрабатываются,"
|
||
" перезапускаю updater",
|
||
pending,
|
||
)
|
||
await self._restart_polling()
|
||
|
||
async def _pending_updates_count(self) -> int | None:
|
||
url = (
|
||
"https://api.telegram.org/bot"
|
||
f"{self._settings.bot_token}/getWebhookInfo"
|
||
)
|
||
proxy = self._settings.telegram_proxy or None
|
||
try:
|
||
async with httpx.AsyncClient(
|
||
proxy=proxy,
|
||
timeout=httpx.Timeout(10.0, connect=5.0),
|
||
) as client:
|
||
response = await client.get(url)
|
||
data = response.json()
|
||
except (httpx.HTTPError, ValueError) as exc:
|
||
logger.warning(
|
||
"Polling watchdog: проверка не удалась (%s)",
|
||
type(exc).__name__,
|
||
)
|
||
return None
|
||
if not data.get("ok"):
|
||
return None
|
||
result = data.get("result") or {}
|
||
return int(result.get("pending_update_count") or 0)
|
||
|
||
async def _restart_polling(self) -> None:
|
||
application = self._application
|
||
if application is None or application.updater is None:
|
||
return
|
||
try:
|
||
await asyncio.wait_for(
|
||
application.updater.stop(), timeout=15
|
||
)
|
||
await asyncio.wait_for(
|
||
application.updater.start_polling(
|
||
allowed_updates=Update.ALL_TYPES,
|
||
drop_pending_updates=False,
|
||
),
|
||
timeout=30,
|
||
)
|
||
except Exception:
|
||
logger.exception(
|
||
"Polling watchdog: не удалось перезапустить updater,"
|
||
" выхожу для рестарта контейнера"
|
||
)
|
||
os._exit(1)
|
||
self._last_update_at = time.monotonic()
|
||
logger.warning("Polling watchdog: updater перезапущен")
|
||
|
||
async def _daily_credits_loop(self) -> None:
|
||
while True:
|
||
now = datetime.now(MSK)
|
||
next_midnight = datetime.combine(
|
||
now.date() + timedelta(days=1), dt_time.min, MSK
|
||
)
|
||
await asyncio.sleep(max(1.0, (next_midnight - now).total_seconds()))
|
||
try:
|
||
granted = await self._db.accrue_daily_credits()
|
||
logger.info("Ежедневные генерации начислены: %s", granted)
|
||
except Exception:
|
||
logger.exception("Ошибка ежедневного начисления генераций")
|
||
|
||
async def _telegram_health_loop(self) -> None:
|
||
interval = self._settings.telegram_health_interval
|
||
max_delay = self._settings.telegram_retry_max_delay
|
||
delay = min(1.0, max_delay)
|
||
while True:
|
||
await asyncio.sleep(interval)
|
||
while True:
|
||
try:
|
||
if self._bot is not None:
|
||
await self._bot.get_me()
|
||
break
|
||
except (TimedOut, NetworkError):
|
||
logger.warning(
|
||
"Telegram health check failed, retry in %.0f s", delay
|
||
)
|
||
await asyncio.sleep(delay)
|
||
delay = min(max_delay, delay * 2)
|
||
except TelegramError:
|
||
logger.exception("Telegram health check error")
|
||
break
|
||
delay = min(1.0, max_delay)
|
||
|
||
def _start_text(self, balance: int | None = None) -> str:
|
||
parts: list[str] = []
|
||
if balance is not None:
|
||
parts.append(
|
||
f"Баланс: {balance} {self._credits_word(balance)}."
|
||
)
|
||
if self._settings.payments_enabled:
|
||
normal = self._settings.stars_price
|
||
premium = self._settings.premium_stars_price
|
||
parts.append(
|
||
"Стоимость генерации:\n"
|
||
f"• обычная очередь — {normal} звёзд;\n"
|
||
f"• премиум-очередь — {premium} звёзд."
|
||
)
|
||
parts.append(START_TEXT)
|
||
return "\n\n".join(parts)
|
||
|
||
@staticmethod
|
||
def _credits_word(count: int) -> str:
|
||
if count % 10 == 1 and count % 100 != 11:
|
||
return "генерация"
|
||
if count % 10 in (2, 3, 4) and count % 100 not in (12, 13, 14):
|
||
return "генерации"
|
||
return "генераций"
|
||
|
||
@staticmethod
|
||
def _parse_referrer(args: list[str]) -> int | None:
|
||
if not args:
|
||
return None
|
||
raw = args[0].strip()
|
||
if not raw.startswith("ref_"):
|
||
return None
|
||
try:
|
||
return int(raw[4:])
|
||
except ValueError:
|
||
return None
|
||
|
||
async def _start(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_base_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return
|
||
row = await self._db.ensure_user(user.id, user.username)
|
||
referrer_id = self._parse_referrer(context.args)
|
||
if (
|
||
referrer_id
|
||
and referrer_id != user.id
|
||
and await self._db.get_user(referrer_id) is not None
|
||
):
|
||
await self._db.set_referrer(user.id, referrer_id)
|
||
row = await self._db.ensure_user(user.id, user.username)
|
||
if row.age_declined:
|
||
await message.reply_text(AGE_DENIED)
|
||
return
|
||
if not row.age_confirmed:
|
||
await self._retry(
|
||
message.reply_text,
|
||
AGE_TEXT,
|
||
reply_markup=self._age_keyboard(),
|
||
)
|
||
return
|
||
await self._retry(
|
||
message.reply_text, self._start_text(row.credits)
|
||
)
|
||
|
||
@staticmethod
|
||
def _age_keyboard() -> InlineKeyboardMarkup:
|
||
return InlineKeyboardMarkup(
|
||
[
|
||
[
|
||
InlineKeyboardButton(
|
||
"Мне 18+", callback_data="age:yes"
|
||
),
|
||
InlineKeyboardButton(
|
||
"Мне меньше 18", callback_data="age:no"
|
||
),
|
||
]
|
||
]
|
||
)
|
||
|
||
async def _age_choice(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
query = update.callback_query
|
||
user = update.effective_user
|
||
if query is None or query.data is None or user is None:
|
||
return
|
||
row = await self._db.ensure_user(user.id, user.username)
|
||
if row.is_blocked:
|
||
await query.answer("Доступ заблокирован.", show_alert=True)
|
||
return
|
||
if query.data == "age:yes":
|
||
await self._db.set_age_confirmed(user.id, True)
|
||
await query.answer("Возраст подтверждён")
|
||
row = await self._db.get_user(user.id)
|
||
try:
|
||
await query.edit_message_text("Возраст подтверждён ✅")
|
||
except TelegramError:
|
||
pass
|
||
if row is not None:
|
||
await self._retry(
|
||
self._bot.send_message,
|
||
chat_id=user.id,
|
||
text=self._start_text(row.credits),
|
||
)
|
||
return
|
||
if query.data == "age:no":
|
||
await self._db.set_age_declined(user.id)
|
||
await query.answer()
|
||
try:
|
||
await query.edit_message_text(AGE_DENIED)
|
||
except TelegramError:
|
||
pass
|
||
|
||
async def _help(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return
|
||
row = await self._db.ensure_user(user.id, user.username)
|
||
await self._retry(
|
||
message.reply_text, self._start_text(row.credits)
|
||
)
|
||
|
||
async def _balance(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return
|
||
row = await self._db.ensure_user(user.id, user.username)
|
||
await self._retry(
|
||
message.reply_text,
|
||
f"Баланс: {row.credits} {self._credits_word(row.credits)}.",
|
||
)
|
||
|
||
async def _invite(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None or self._bot is None:
|
||
return
|
||
await self._db.ensure_user(user.id, user.username)
|
||
username = self._bot.username
|
||
if not username:
|
||
await message.reply_text("Не удалось получить ссылку. Попробуй позже.")
|
||
return
|
||
link = f"https://t.me/{username}?start=ref_{user.id}"
|
||
keyboard = InlineKeyboardMarkup(
|
||
[
|
||
[
|
||
InlineKeyboardButton(
|
||
"Скопировать ссылку",
|
||
copy_text=CopyTextButton(text=link),
|
||
)
|
||
]
|
||
]
|
||
)
|
||
await self._retry(
|
||
message.reply_text,
|
||
"Пригласи друга:\n"
|
||
f"{link}\n\n"
|
||
"Когда друг впервые оплатит генерацию, ты получишь +3 генерации.",
|
||
reply_markup=keyboard,
|
||
)
|
||
|
||
async def _reset(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return
|
||
cleared = await self._db.clear_user_context(user.id)
|
||
text = (
|
||
"Контекст сброшен. Следующий запрос начнёт новую сцену."
|
||
if cleared
|
||
else "Контекст уже пуст."
|
||
)
|
||
await self._retry(message.reply_text, text)
|
||
|
||
async def _context_command(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return
|
||
user_context = await self._db.get_user_context(user.id)
|
||
await self._retry(
|
||
message.reply_text, format_context(user_context)
|
||
)
|
||
|
||
@staticmethod
|
||
def _settings_keyboard(
|
||
quality: str, steps: int
|
||
) -> InlineKeyboardMarkup:
|
||
def quality_label(name: str, value: str) -> str:
|
||
return f"{name} ✅" if quality == value else name
|
||
|
||
def steps_label(value: int) -> str:
|
||
text = f"{value} шагов"
|
||
return f"{text} ✅" if steps == value else text
|
||
|
||
return InlineKeyboardMarkup(
|
||
[
|
||
[
|
||
InlineKeyboardButton(
|
||
quality_label("Стандартная", "standard"),
|
||
callback_data="quality:standard",
|
||
)
|
||
],
|
||
[
|
||
InlineKeyboardButton(
|
||
quality_label("Улучшенная", "improved"),
|
||
callback_data="quality:improved",
|
||
)
|
||
],
|
||
[
|
||
InlineKeyboardButton(
|
||
steps_label(8), callback_data="steps:8"
|
||
),
|
||
InlineKeyboardButton(
|
||
steps_label(10), callback_data="steps:10"
|
||
),
|
||
InlineKeyboardButton(
|
||
steps_label(12), callback_data="steps:12"
|
||
),
|
||
],
|
||
]
|
||
)
|
||
|
||
async def _settings_command(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return
|
||
row = await self._db.ensure_user(user.id, user.username)
|
||
await self._retry(
|
||
message.reply_text,
|
||
"Качество Krea2 для редактирования фото:\n"
|
||
"Стандартная — Q4\n"
|
||
"Улучшенная — Q6\n"
|
||
"Шаги: 8 / 10 / 12",
|
||
reply_markup=self._settings_keyboard(
|
||
row.generation_quality, row.generation_steps
|
||
),
|
||
)
|
||
|
||
async def _quality_choice(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
query = update.callback_query
|
||
if query is None or query.data is None:
|
||
return
|
||
try:
|
||
_, quality = query.data.split(":", 1)
|
||
except ValueError:
|
||
await query.answer("Некорректный выбор.", show_alert=True)
|
||
return
|
||
if quality not in {"standard", "improved"}:
|
||
await query.answer("Некорректный выбор.", show_alert=True)
|
||
return
|
||
user = query.from_user
|
||
row = await self._db.ensure_user(user.id, user.username)
|
||
await self._db.set_generation_quality(user.id, quality)
|
||
await query.answer()
|
||
if query.message is not None:
|
||
try:
|
||
await query.edit_message_reply_markup(
|
||
reply_markup=self._settings_keyboard(
|
||
quality, row.generation_steps
|
||
)
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось обновить настройки качества")
|
||
|
||
async def _steps_choice(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
query = update.callback_query
|
||
if query is None or query.data is None:
|
||
return
|
||
try:
|
||
_, raw_steps = query.data.split(":", 1)
|
||
steps = int(raw_steps)
|
||
except ValueError:
|
||
await query.answer("Некорректный выбор.", show_alert=True)
|
||
return
|
||
if steps not in {8, 10, 12}:
|
||
await query.answer("Некорректный выбор.", show_alert=True)
|
||
return
|
||
user = query.from_user
|
||
row = await self._db.ensure_user(user.id, user.username)
|
||
await self._db.set_generation_steps(user.id, steps)
|
||
await query.answer()
|
||
if query.message is not None:
|
||
try:
|
||
await query.edit_message_reply_markup(
|
||
reply_markup=self._settings_keyboard(
|
||
row.generation_quality, steps
|
||
)
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось обновить настройки шагов")
|
||
|
||
async def _status(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return
|
||
active = await self._db.get_active_by_user(user.id)
|
||
if active is None:
|
||
await message.reply_text("Активных задач нет.")
|
||
return
|
||
if active.status == "running":
|
||
label = " (премиум)" if active.is_premium else ""
|
||
await message.reply_text(f"Твоя картинка сейчас генерируется{label}.")
|
||
return
|
||
position = await self._queue.position(active.id) if self._queue else 0
|
||
label = "премиум-очередь" if active.is_premium else "очередь"
|
||
await message.reply_text(
|
||
f"Задача в {label}. Перед тобой: {position} {self._people(position)}."
|
||
)
|
||
|
||
async def _reward_referrer(self, user) -> None:
|
||
if not await self._db.mark_first_payment(user.id):
|
||
return
|
||
row = await self._db.get_user(user.id)
|
||
if row is None or row.referrer_id is None:
|
||
return
|
||
balance = await self._db.add_credits(row.referrer_id, 3)
|
||
if balance is None:
|
||
logger.warning("Не найден реферер %s", row.referrer_id)
|
||
return
|
||
text = (
|
||
f"Твой друг {row.display_user} сделал первую оплату. "
|
||
f"+3 генерации. Баланс: {balance}."
|
||
)
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_message, chat_id=row.referrer_id, text=text
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось уведомить реферера %s", row.referrer_id)
|
||
|
||
async def _pre_checkout(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
query = update.pre_checkout_query
|
||
if query is None:
|
||
return
|
||
pending = await self._db.get_pending_payment(query.invoice_payload)
|
||
user_row = await self._db.get_user(query.from_user.id)
|
||
allowed = (
|
||
pending is not None
|
||
and pending.user_id == query.from_user.id
|
||
and pending.status == "pending"
|
||
and query.currency == "XTR"
|
||
and query.total_amount == pending.amount
|
||
and user_row is not None
|
||
and user_row.age_confirmed
|
||
and not user_row.age_declined
|
||
and not user_row.is_blocked
|
||
)
|
||
try:
|
||
await query.answer(
|
||
ok=allowed,
|
||
error_message=None
|
||
if allowed
|
||
else "Счёт устарел. Отправь запрос заново.",
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось ответить на pre_checkout_query")
|
||
|
||
async def _queue_choice(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
query = update.callback_query
|
||
if query is None or query.data is None:
|
||
return
|
||
try:
|
||
_, payload, tier = query.data.split(":", 2)
|
||
except ValueError:
|
||
await query.answer("Некорректный выбор.", show_alert=True)
|
||
return
|
||
if tier not in {"normal", "premium", "credit"}:
|
||
await query.answer("Некорректный выбор.", show_alert=True)
|
||
return
|
||
|
||
pending = await self._db.get_pending_payment(payload)
|
||
if (
|
||
pending is None
|
||
or pending.user_id != query.from_user.id
|
||
or pending.status != "pending"
|
||
or pending.amount != 0
|
||
):
|
||
await query.answer("Счёт уже создан или устарел.", show_alert=True)
|
||
return
|
||
|
||
active = await self._db.get_active_by_user(query.from_user.id)
|
||
if active is not None:
|
||
await query.answer(
|
||
"У тебя уже есть активная задача.", show_alert=True
|
||
)
|
||
return
|
||
|
||
user_row = await self._db.get_user(query.from_user.id)
|
||
if user_row is not None and user_row.is_blocked:
|
||
await query.answer("Доступ заблокирован.", show_alert=True)
|
||
return
|
||
if user_row is None or not user_row.age_confirmed:
|
||
await query.answer(
|
||
"Сначала подтверди возраст.", show_alert=True
|
||
)
|
||
if query.message is not None:
|
||
await query.message.reply_text(
|
||
AGE_TEXT, reply_markup=self._age_keyboard()
|
||
)
|
||
return
|
||
|
||
if tier == "credit":
|
||
remaining = await self._db.consume_credit(query.from_user.id)
|
||
if remaining is None:
|
||
await query.answer(
|
||
"Генерации закончились. Выбери оплату.", show_alert=True
|
||
)
|
||
return
|
||
await self._db.mark_pending_cancelled(payload)
|
||
await query.answer()
|
||
try:
|
||
await query.edit_message_text(
|
||
f"Списана 1 генерация. Осталось: {remaining}."
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось обновить сообщение выбора очереди")
|
||
accepted = await self._submit(
|
||
update,
|
||
mode=pending.mode,
|
||
prompt=pending.prompt,
|
||
input_path=pending.input_file_path,
|
||
skip_payment=True,
|
||
is_premium=False,
|
||
use_context=bool(pending.context_prompt),
|
||
context_base_prompt=pending.context_base_prompt,
|
||
context_prompt=pending.context_prompt,
|
||
)
|
||
if not accepted:
|
||
await self._db.add_credits(query.from_user.id, 1)
|
||
try:
|
||
await query.edit_message_text(
|
||
"Не удалось поставить задачу. Генерация возвращена."
|
||
)
|
||
except TelegramError:
|
||
pass
|
||
return
|
||
|
||
is_premium = tier == "premium"
|
||
amount = (
|
||
self._settings.premium_stars_price
|
||
if is_premium
|
||
else self._settings.stars_price
|
||
)
|
||
if not await self._db.set_pending_choice(payload, is_premium, amount):
|
||
await query.answer("Счёт уже создан.", show_alert=True)
|
||
return
|
||
await query.answer()
|
||
|
||
label = "Премиум-очередь" if is_premium else "Обычная очередь"
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_invoice,
|
||
chat_id=query.from_user.id,
|
||
title=label,
|
||
description=(
|
||
"Приоритетная генерация." if is_premium
|
||
else "Обычная генерация."
|
||
),
|
||
payload=payload,
|
||
provider_token="",
|
||
currency="XTR",
|
||
prices=[LabeledPrice(label=label, amount=amount)],
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось выставить счёт пользователю %s", query.from_user.id)
|
||
await self._db.mark_pending_cancelled(payload)
|
||
await query.edit_message_text("Не удалось выставить счёт. Попробуй позже.")
|
||
return
|
||
|
||
await query.edit_message_text(
|
||
f"Счёт на {amount} звёзд отправлен. Очередь: {label.lower()}."
|
||
)
|
||
|
||
async def _successful_payment(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
payment = message.successful_payment if message is not None else None
|
||
if message is None or user is None or payment is None:
|
||
return
|
||
|
||
payload = payment.invoice_payload
|
||
pending = await self._db.get_pending_payment(payload)
|
||
if pending is None or pending.user_id != user.id:
|
||
logger.error(
|
||
"Платёж %s не найден или принадлежит другому пользователю",
|
||
payload,
|
||
)
|
||
return
|
||
if pending.status == "paid":
|
||
await message.reply_text("Этот платёж уже обработан.")
|
||
return
|
||
if pending.status != "pending":
|
||
return
|
||
|
||
user_row = await self._db.get_user(user.id)
|
||
if (
|
||
user_row is None
|
||
or not user_row.age_confirmed
|
||
or user_row.age_declined
|
||
or user_row.is_blocked
|
||
):
|
||
refunded = await self._refund_stars(
|
||
user.id, payment.telegram_payment_charge_id
|
||
)
|
||
if refunded:
|
||
await self._db.mark_pending_refunded(payload)
|
||
note = " Звёзды возвращены." if refunded else ""
|
||
await message.reply_text("Доступ заблокирован." + note)
|
||
return
|
||
|
||
charge_id = payment.telegram_payment_charge_id
|
||
if payment.currency != "XTR" or payment.total_amount != pending.amount:
|
||
refunded = await self._refund_stars(user.id, charge_id)
|
||
if refunded:
|
||
await self._db.mark_pending_refunded(payload)
|
||
note = " Звёзды возвращены." if refunded else ""
|
||
await message.reply_text("Платёж не принят." + note)
|
||
return
|
||
|
||
if not await self._db.mark_pending_paid(payload, charge_id):
|
||
await message.reply_text("Этот платёж уже обработан.")
|
||
return
|
||
|
||
accepted = await self._submit(
|
||
update,
|
||
mode=pending.mode,
|
||
prompt=pending.prompt,
|
||
input_path=pending.input_file_path,
|
||
skip_payment=True,
|
||
payment_charge_id=charge_id,
|
||
payment_payload=payload,
|
||
is_premium=pending.is_premium,
|
||
use_context=bool(pending.context_prompt),
|
||
context_base_prompt=pending.context_base_prompt,
|
||
context_prompt=pending.context_prompt,
|
||
)
|
||
if accepted:
|
||
await self._reward_referrer(user)
|
||
|
||
async def _text(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None or not message.text:
|
||
return
|
||
|
||
prompt = message.text.strip()
|
||
if not prompt:
|
||
await message.reply_text("Промпт пустой. Напиши, что нарисовать.")
|
||
return
|
||
if len(prompt) > self._settings.max_prompt_length:
|
||
await message.reply_text(
|
||
f"Промпт слишком длинный: максимум {self._settings.max_prompt_length} символов."
|
||
)
|
||
return
|
||
|
||
reply = message.reply_to_message
|
||
attachment = None
|
||
suffix = ".jpg"
|
||
if reply is not None:
|
||
if reply.photo:
|
||
attachment = reply.photo[-1]
|
||
elif reply.document is not None and (
|
||
reply.document.mime_type or ""
|
||
).startswith("image/"):
|
||
attachment = reply.document
|
||
suffix = Path(reply.document.file_name or "").suffix.lower() or ".png"
|
||
|
||
if attachment is None:
|
||
await self._submit(
|
||
update, mode="t2i", prompt=prompt, input_path=None
|
||
)
|
||
return
|
||
|
||
input_path = await self._download_attachment(
|
||
message, attachment, user.id, suffix=suffix
|
||
)
|
||
if input_path is None:
|
||
return
|
||
accepted = await self._submit(
|
||
update,
|
||
mode="i2i",
|
||
prompt=prompt,
|
||
input_path=str(input_path),
|
||
)
|
||
if not accepted:
|
||
input_path.unlink(missing_ok=True)
|
||
|
||
async def _photo(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None or not message.photo:
|
||
return
|
||
|
||
caption = (message.caption or "").strip()
|
||
if not caption:
|
||
await message.reply_text(
|
||
"Пришли фото с подписью — что именно изменить на изображении."
|
||
)
|
||
return
|
||
if len(caption) > self._settings.max_prompt_length:
|
||
await message.reply_text(
|
||
f"Описание слишком длинное: максимум {self._settings.max_prompt_length} символов."
|
||
)
|
||
return
|
||
|
||
photo = message.photo[-1]
|
||
input_path = await self._download_attachment(
|
||
message, photo, user.id, suffix=".jpg"
|
||
)
|
||
if input_path is None:
|
||
return
|
||
|
||
accepted = await self._submit(
|
||
update,
|
||
mode="i2i",
|
||
prompt=caption,
|
||
input_path=str(input_path),
|
||
use_context=False,
|
||
)
|
||
if not accepted:
|
||
input_path.unlink(missing_ok=True)
|
||
|
||
async def _download_attachment(
|
||
self,
|
||
message,
|
||
attachment,
|
||
user_id: int,
|
||
*,
|
||
suffix: str = ".jpg",
|
||
) -> Path | None:
|
||
if attachment.file_size and attachment.file_size > self._settings.max_photo_bytes:
|
||
limit_mb = self._settings.max_photo_bytes / (1024 * 1024)
|
||
try:
|
||
await self._retry(
|
||
message.reply_text,
|
||
f"Файл слишком большой. Максимум {limit_mb:.0f} МБ.",
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось сообщить о размере файла")
|
||
return None
|
||
|
||
suffix = suffix.lower()
|
||
if not suffix.startswith(".") or len(suffix) > 8:
|
||
suffix = ".png"
|
||
|
||
user_dir = self._settings.storage_dir / str(user_id)
|
||
input_path = user_dir / f"{uuid.uuid4().hex}_input{suffix}"
|
||
try:
|
||
user_dir.mkdir(parents=True, exist_ok=True)
|
||
await self._retry(
|
||
message.chat.send_action, ChatAction.TYPING, attempts=3
|
||
)
|
||
|
||
async def download() -> None:
|
||
input_path.unlink(missing_ok=True)
|
||
file = await attachment.get_file()
|
||
await file.download_to_drive(custom_path=input_path)
|
||
|
||
await self._retry(download, attempts=5)
|
||
except (TelegramError, OSError):
|
||
logger.exception("Не удалось скачать вложение пользователя %s", user_id)
|
||
input_path.unlink(missing_ok=True)
|
||
try:
|
||
await self._retry(
|
||
message.reply_text,
|
||
"Не удалось скачать файл, попробуй ещё раз.",
|
||
)
|
||
except TelegramError:
|
||
pass
|
||
return None
|
||
return input_path
|
||
|
||
async def _submit(
|
||
self,
|
||
update: Update,
|
||
*,
|
||
mode: str,
|
||
prompt: str,
|
||
input_path: str | None,
|
||
skip_payment: bool = False,
|
||
payment_charge_id: str | None = None,
|
||
payment_payload: str | None = None,
|
||
is_premium: bool = False,
|
||
use_context: bool = True,
|
||
context_base_prompt: str | None = None,
|
||
context_prompt: str | None = None,
|
||
) -> bool:
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None or self._queue is None:
|
||
return False
|
||
|
||
user_row = await self._db.ensure_user(user.id, user.username)
|
||
if user_row.age_declined:
|
||
await message.reply_text(AGE_DENIED)
|
||
return False
|
||
if not user_row.age_confirmed:
|
||
await message.reply_text(
|
||
AGE_TEXT, reply_markup=self._age_keyboard()
|
||
)
|
||
return False
|
||
is_admin = user_row.is_admin
|
||
if user_row.is_blocked:
|
||
await message.reply_text("Доступ заблокирован.")
|
||
return False
|
||
if not use_context:
|
||
context_base_prompt = None
|
||
context_prompt = None
|
||
elif context_prompt is None:
|
||
user_context = await self._db.get_user_context(user.id)
|
||
if user_context is not None:
|
||
context_base_prompt = user_context.base_prompt
|
||
context_prompt = user_context.current_prompt
|
||
credits = user_row.credits
|
||
has_person = False
|
||
if (
|
||
mode == "i2i"
|
||
and input_path
|
||
and not is_admin
|
||
and self._settings.moderation_enabled
|
||
):
|
||
has_person = await asyncio.to_thread(detect_person, input_path)
|
||
|
||
if (
|
||
not skip_payment
|
||
and not is_admin
|
||
and self._settings.payments_enabled
|
||
):
|
||
allowed, reason = is_prompt_allowed(
|
||
prompt,
|
||
mode=mode,
|
||
has_person=has_person,
|
||
extra_public_figures=self._settings.blocked_public_figures,
|
||
)
|
||
if not allowed:
|
||
await self._db.log_blocked(
|
||
user_id=user.id,
|
||
prompt=prompt,
|
||
reason=reason or "unknown",
|
||
mode=mode,
|
||
has_person=has_person,
|
||
)
|
||
await message.reply_text(BLOCKED_TEXT)
|
||
return False
|
||
|
||
active = await self._db.get_active_by_user(user.id)
|
||
if active is not None:
|
||
if active.status == "running":
|
||
text = "Твоя картинка уже генерируется."
|
||
else:
|
||
ahead = await self._queue.position(active.id)
|
||
text = f"Ты уже в очереди. Перед тобой: {ahead} {self._people(ahead)}."
|
||
await message.reply_text(text)
|
||
return False
|
||
|
||
counts = await self._queue.queue_counts()
|
||
payload = uuid.uuid4().hex
|
||
await self._db.create_pending_payment(
|
||
payload=payload,
|
||
user_id=user.id,
|
||
mode=mode,
|
||
prompt=prompt,
|
||
username=user.username,
|
||
input_file_path=input_path,
|
||
amount=0,
|
||
context_base_prompt=context_base_prompt,
|
||
context_prompt=context_prompt,
|
||
)
|
||
keyboard = InlineKeyboardMarkup(
|
||
[
|
||
[
|
||
InlineKeyboardButton(
|
||
(
|
||
f"Обычная — 1 генерация (осталось {credits})"
|
||
if credits > 0
|
||
else f"Обычная — {self._settings.stars_price} звёзд"
|
||
),
|
||
callback_data=(
|
||
f"queue:{payload}:credit"
|
||
if credits > 0
|
||
else f"queue:{payload}:normal"
|
||
),
|
||
)
|
||
],
|
||
[
|
||
InlineKeyboardButton(
|
||
f"Премиум — {self._settings.premium_stars_price} звёзд",
|
||
callback_data=f"queue:{payload}:premium",
|
||
)
|
||
],
|
||
]
|
||
)
|
||
text = (
|
||
f"Баланс: {credits} {self._credits_word(credits)}.\n"
|
||
f"Обычная очередь: {counts['normal']} чел.\n"
|
||
f"Премиум-очередь: {counts['premium']} чел.\n\n"
|
||
"Выбери очередь:"
|
||
)
|
||
try:
|
||
await self._retry(
|
||
message.reply_text,
|
||
text,
|
||
reply_markup=keyboard,
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось показать выбор очереди %s", user.id)
|
||
await self._db.mark_pending_cancelled(payload)
|
||
return False
|
||
return True
|
||
|
||
if is_admin:
|
||
is_premium = False
|
||
|
||
try:
|
||
request, position = await self._queue.enqueue(
|
||
user_id=user.id,
|
||
username=user.username,
|
||
mode=mode,
|
||
prompt=prompt,
|
||
input_file_path=input_path,
|
||
has_person=has_person,
|
||
is_admin=is_admin,
|
||
payment_charge_id=payment_charge_id,
|
||
is_premium=is_premium,
|
||
context_base_prompt=context_base_prompt,
|
||
context_prompt=context_prompt,
|
||
)
|
||
except PromptBlocked:
|
||
note = await self._refund_if_paid(
|
||
user.id, payment_charge_id, payment_payload
|
||
)
|
||
await message.reply_text(BLOCKED_TEXT + note)
|
||
return False
|
||
except UserAlreadyQueued:
|
||
note = await self._refund_if_paid(
|
||
user.id, payment_charge_id, payment_payload
|
||
)
|
||
active = await self._db.get_active_by_user(user.id)
|
||
text = "Ты уже в очереди."
|
||
if active is not None:
|
||
if active.status == "running":
|
||
text += " Твоя картинка уже генерируется."
|
||
else:
|
||
ahead = await self._queue.position(active.id)
|
||
text += f" Перед тобой: {ahead} {self._people(ahead)}."
|
||
await message.reply_text(text + note)
|
||
return False
|
||
except QueueFull as exc:
|
||
note = await self._refund_if_paid(
|
||
user.id, payment_charge_id, payment_payload
|
||
)
|
||
await message.reply_text(exc.user_message + note)
|
||
return False
|
||
except Exception:
|
||
logger.exception("Не удалось поставить задачу пользователю %s", user.id)
|
||
note = await self._refund_if_paid(
|
||
user.id, payment_charge_id, payment_payload
|
||
)
|
||
await message.reply_text(
|
||
"Не удалось поставить задачу. Попробуй позже." + note
|
||
)
|
||
return False
|
||
|
||
if payment_payload:
|
||
await self._db.attach_pending_request(payment_payload, request.id)
|
||
|
||
try:
|
||
status_message = await self._retry(
|
||
message.reply_text, self._ack_text(position, is_premium)
|
||
)
|
||
except TelegramError:
|
||
logger.exception(
|
||
"Не удалось отправить подтверждение по запросу %s", request.id
|
||
)
|
||
return True
|
||
self._status_messages[request.id] = (
|
||
status_message.chat_id,
|
||
status_message.message_id,
|
||
)
|
||
|
||
# Воркер мог стартовать раньше, чем мы сохранили message_id.
|
||
current = await self._db.get_request(request.id)
|
||
if current is None:
|
||
return True
|
||
if current.status == "running":
|
||
await self._edit_status(request.id, "Генерация…", percent=0)
|
||
elif current.status == "done":
|
||
await self._clear_status(request.id, delete=True)
|
||
elif current.status == "failed":
|
||
await self._finish_failed(
|
||
current,
|
||
current.error_text or "Не удалось сгенерировать изображение.",
|
||
)
|
||
return True
|
||
|
||
async def _refund_if_paid(
|
||
self,
|
||
user_id: int,
|
||
charge_id: str | None,
|
||
payload: str | None,
|
||
) -> str:
|
||
if not charge_id:
|
||
return ""
|
||
refunded = await self._refund_stars(user_id, charge_id)
|
||
if refunded and payload:
|
||
await self._db.mark_pending_refunded(payload)
|
||
if refunded:
|
||
return " Звёзды возвращены."
|
||
return " Не удалось вернуть звёзды, напиши админу."
|
||
|
||
async def _refund_stars(self, user_id: int, charge_id: str) -> bool:
|
||
if self._bot is None:
|
||
return False
|
||
try:
|
||
return await self._retry(
|
||
self._bot.refund_star_payment,
|
||
user_id=user_id,
|
||
telegram_payment_charge_id=charge_id,
|
||
attempts=2,
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось вернуть звёзды пользователю %s", user_id)
|
||
return False
|
||
|
||
def _ack_text(self, position: int, is_premium: bool = False) -> str:
|
||
label = "Премиум-очередь. " if is_premium else ""
|
||
if position <= 0:
|
||
return f"Принято. {label}Ты первый, начинаю…"
|
||
return f"Принято. {label}Перед тобой: {position} {self._people(position)}."
|
||
|
||
@staticmethod
|
||
def _people(count: int) -> str:
|
||
if count % 10 == 1 and count % 100 != 11:
|
||
return "человек"
|
||
if count % 10 in (2, 3, 4) and count % 100 not in (12, 13, 14):
|
||
return "человека"
|
||
return "человек"
|
||
|
||
async def _on_queue_event(self, event: QueueEvent) -> None:
|
||
if event.kind == "started":
|
||
await self._edit_status(event.request.id, "Генерация…", percent=0)
|
||
elif event.kind == "enhancing":
|
||
await self._edit_status(event.request.id, "Улучшаю промпт…")
|
||
elif event.kind == "progress":
|
||
percent = event.percent if event.percent is not None else 0
|
||
await self._edit_status(
|
||
event.request.id,
|
||
f"Генерация… {percent}%",
|
||
percent=percent,
|
||
)
|
||
elif event.kind == "done":
|
||
await self._finish_done(event.request)
|
||
elif event.kind == "failed":
|
||
await self._finish_failed(
|
||
event.request,
|
||
event.error or "Не удалось сгенерировать изображение.",
|
||
)
|
||
|
||
async def _retry(
|
||
self,
|
||
func,
|
||
*args,
|
||
attempts: int | None = None,
|
||
**kwargs,
|
||
):
|
||
max_attempts = (
|
||
self._settings.telegram_retry_attempts
|
||
if attempts is None
|
||
else attempts
|
||
)
|
||
infinite = max_attempts <= 0
|
||
max_delay = self._settings.telegram_retry_max_delay
|
||
delay = min(1.0, max_delay)
|
||
attempt = 0
|
||
while True:
|
||
attempt += 1
|
||
try:
|
||
return await func(*args, **kwargs)
|
||
except (TimedOut, NetworkError):
|
||
if not infinite and attempt >= max_attempts:
|
||
raise
|
||
logger.warning(
|
||
"Telegram timeout (attempt %s%s), retry in %.1f s",
|
||
attempt,
|
||
"" if infinite else f"/{max_attempts}",
|
||
delay,
|
||
)
|
||
await asyncio.sleep(delay)
|
||
delay = min(max_delay, delay * 2)
|
||
|
||
async def _edit_status(
|
||
self,
|
||
request_id: int,
|
||
text: str,
|
||
*,
|
||
percent: int | None = None,
|
||
) -> None:
|
||
if self._bot is None:
|
||
return
|
||
target = self._status_messages.get(request_id)
|
||
if target is None:
|
||
return
|
||
|
||
now = time.monotonic()
|
||
last = self._last_edit.get(request_id)
|
||
if (
|
||
last is not None
|
||
and percent is not None
|
||
and percent - last[1] < 3
|
||
and now - last[0] < 2.5
|
||
):
|
||
return
|
||
|
||
chat_id, message_id = target
|
||
try:
|
||
await self._retry(
|
||
self._bot.edit_message_text,
|
||
chat_id=chat_id,
|
||
message_id=message_id,
|
||
text=text,
|
||
attempts=3,
|
||
)
|
||
except BadRequest as exc:
|
||
if "not modified" not in str(exc).lower():
|
||
logger.debug("Не удалось обновить статус %s: %s", request_id, exc)
|
||
return
|
||
except TelegramError as exc:
|
||
logger.debug("Не удалось обновить статус %s: %s", request_id, exc)
|
||
return
|
||
self._last_edit[request_id] = (
|
||
now,
|
||
percent if percent is not None else -1,
|
||
)
|
||
|
||
async def _finish_done(self, request: Request) -> None:
|
||
target = self._status_messages.get(request.id)
|
||
chat_id = target[0] if target is not None else request.user_id
|
||
output_path = request.output_file_path
|
||
|
||
if not output_path or not Path(output_path).exists():
|
||
await self._finish_failed(request, "Файл результата не найден.")
|
||
return
|
||
|
||
if self._bot is None:
|
||
return
|
||
|
||
prompt = request.enhanced_prompt or request.prompt_text
|
||
caption = result_caption(prompt)
|
||
filename = Path(output_path).name
|
||
try:
|
||
content = await asyncio.to_thread(Path(output_path).read_bytes)
|
||
except OSError:
|
||
logger.exception("Не удалось прочитать результат запроса %s", request.id)
|
||
await self._finish_failed(request, "Файл результата не читается.")
|
||
return
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_photo,
|
||
chat_id=chat_id,
|
||
photo=InputFile(content, filename=filename),
|
||
caption=caption,
|
||
read_timeout=120,
|
||
write_timeout=120,
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось отправить фото по запросу %s", request.id)
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_document,
|
||
chat_id=chat_id,
|
||
document=InputFile(content, filename=filename),
|
||
caption=caption,
|
||
read_timeout=120,
|
||
write_timeout=120,
|
||
)
|
||
except TelegramError:
|
||
await self._finish_failed(
|
||
request, "Не удалось отправить результат в Telegram."
|
||
)
|
||
return
|
||
|
||
await self._clear_status(request.id, delete=True)
|
||
await self._save_context(request)
|
||
|
||
async def _save_context(self, request: Request) -> None:
|
||
current = request.enhanced_prompt or request.prompt_text
|
||
base = request.context_base_prompt or current
|
||
try:
|
||
await self._db.save_user_context(
|
||
user_id=request.user_id,
|
||
base_prompt=base,
|
||
current_prompt=current,
|
||
last_request_id=request.id,
|
||
last_mode=request.mode,
|
||
)
|
||
except Exception:
|
||
logger.exception(
|
||
"Не удалось сохранить контекст запроса %s", request.id
|
||
)
|
||
|
||
async def _finish_failed(self, request: Request, error: str) -> None:
|
||
target = self._status_messages.get(request.id)
|
||
chat_id = target[0] if target is not None else request.user_id
|
||
refund_note = ""
|
||
if request.payment_charge_id and not request.payment_refunded:
|
||
refunded = await self._refund_stars(
|
||
request.user_id, request.payment_charge_id
|
||
)
|
||
if refunded:
|
||
await self._db.set_refunded(request.id)
|
||
refund_note = " Звёзды возвращены."
|
||
else:
|
||
refund_note = " Не удалось вернуть звёзды, напиши админу."
|
||
text = f"Не получилось: {error}{refund_note}"
|
||
if target is not None and self._bot is not None:
|
||
try:
|
||
await self._retry(
|
||
self._bot.edit_message_text,
|
||
chat_id=target[0],
|
||
message_id=target[1],
|
||
text=text,
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось показать ошибку по запросу %s", request.id)
|
||
if self._bot is not None:
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_message, chat_id=chat_id, text=text
|
||
)
|
||
except TelegramError:
|
||
logger.exception(
|
||
"Не удалось отправить ошибку по запросу %s", request.id
|
||
)
|
||
await self._clear_status(request.id, delete=False)
|
||
return
|
||
|
||
if self._bot is not None:
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_message, chat_id=chat_id, text=text
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось отправить ошибку по запросу %s", request.id)
|
||
|
||
async def _clear_status(self, request_id: int, *, delete: bool) -> None:
|
||
target = self._status_messages.pop(request_id, None)
|
||
self._last_edit.pop(request_id, None)
|
||
if delete and target is not None and self._bot is not None:
|
||
try:
|
||
await self._retry(
|
||
self._bot.delete_message,
|
||
chat_id=target[0],
|
||
message_id=target[1],
|
||
attempts=2,
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось удалить статус %s", request_id)
|
||
|
||
async def _check_base_access(self, update: Update) -> bool:
|
||
message = update.effective_message
|
||
chat = update.effective_chat
|
||
user = update.effective_user
|
||
if message is None or chat is None or user is None:
|
||
return False
|
||
if chat.type != ChatType.PRIVATE:
|
||
await message.reply_text("Напиши боту в личку.")
|
||
return False
|
||
if (
|
||
self._settings.allowed_user_ids
|
||
and user.id not in self._settings.allowed_user_ids
|
||
):
|
||
await message.reply_text("Доступ к боту ограничен.")
|
||
return False
|
||
row = await self._db.get_user(user.id)
|
||
if row is not None and row.is_blocked:
|
||
await message.reply_text("Доступ заблокирован.")
|
||
return False
|
||
return True
|
||
|
||
async def _check_access(self, update: Update) -> bool:
|
||
if not await self._check_base_access(update):
|
||
return False
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return False
|
||
row = await self._db.get_user(user.id)
|
||
if row is not None and row.age_declined:
|
||
await message.reply_text(AGE_DENIED)
|
||
return False
|
||
if row is None or not row.age_confirmed:
|
||
await message.reply_text(
|
||
AGE_TEXT, reply_markup=self._age_keyboard()
|
||
)
|
||
return False
|
||
return True
|
||
|
||
async def _on_error(
|
||
self, update: object, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
logger.error("Ошибка обработчика", exc_info=context.error)
|
||
if isinstance(update, Update) and update.effective_message is not None:
|
||
try:
|
||
await update.effective_message.reply_text(
|
||
"Внутренняя ошибка. Попробуй ещё раз."
|
||
)
|
||
except TelegramError:
|
||
pass
|
||
|
||
|
||
def main() -> None:
|
||
settings = Settings.from_env()
|
||
setup_logging(settings.log_level)
|
||
MpakoBot(settings).run()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|