Премиум-очередь
Build and push Docker image / docker (push) Successful in 10s Details

This commit is contained in:
MpakobecEx 2026-10-18 20:42:23 +05:00
parent dc16bb7e0c
commit 71a7cd5054
11 changed files with 297 additions and 46 deletions

View File

@ -43,6 +43,7 @@ BLOCKED_PUBLIC_FIGURES=
# Оплата Telegram Stars
PAYMENTS_ENABLED=true
STARS_PRICE=5
PREMIUM_STARS_PRICE=15
# Веб-история: только локальная сеть, без авторизации
WEB_BIND_IP=192.168.31.158

View File

@ -10,7 +10,8 @@ Telegram-бот для генерации и редактирования изо
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
- **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов.
- **Оплата**: 5 Telegram Stars за генерацию, админам бесплатно, возврат при ошибке.
- **Оплата**: 5 Telegram Stars за обычную очередь, 15 за премиум; админам бесплатно, возврат при ошибке.
- **Премиум-очередь**: выполняется раньше обычной; при пустой премиум-очереди работает обычная.
- Глобальная очередь с одной активной задачей на пользователя.
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
- SQLite: запросы, статусы, seed, пути к файлам.
@ -124,14 +125,18 @@ BLOCKED_PUBLIC_FIGURES=
## Оплата
Генерация стоит 5 Telegram Stars:
Генерация стоит 5 Telegram Stars в обычной очереди и 15 звёзд в премиум-очереди:
```dotenv
PAYMENTS_ENABLED=true
STARS_PRICE=5
PREMIUM_STARS_PRICE=15
```
- Админы из `ADMIN_USER_IDS` генерируют бесплатно.
- Перед оплатой бот показывает количество людей в обычной и премиум-очередях.
- Пользователь выбирает очередь кнопкой, после чего получает счёт.
- Премиум-задачи выполняются раньше обычных.
- Админы из `ADMIN_USER_IDS` генерируют бесплатно и попадают в премиум-очередь.
- Не-админу выставляется счёт через Telegram Stars (`XTR`).
- Фильтрация промпта выполняется до выставления счёта.
- Если поставить задачу в очередь не удалось, звёзды возвращаются.

162
bot.py
View File

@ -6,11 +6,18 @@ import time
import uuid
from pathlib import Path
from telegram import InputFile, LabeledPrice, Update
from telegram import (
InlineKeyboardButton,
InlineKeyboardMarkup,
InputFile,
LabeledPrice,
Update,
)
from telegram.constants import ChatAction, ChatType
from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut
from telegram.ext import (
Application,
CallbackQueryHandler,
CommandHandler,
ContextTypes,
MessageHandler,
@ -71,6 +78,9 @@ class MpakoBot:
application.add_handler(
PreCheckoutQueryHandler(self._pre_checkout)
)
application.add_handler(
CallbackQueryHandler(self._queue_choice, pattern=r"^queue:")
)
application.add_handler(
MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment)
)
@ -120,8 +130,14 @@ class MpakoBot:
def _start_text(self) -> str:
if self._settings.payments_enabled:
price = self._settings.stars_price
return f"Стоимость одной генерации: {price} звёзд.\n\n" + START_TEXT
normal = self._settings.stars_price
premium = self._settings.premium_stars_price
return (
"Стоимость генерации:\n"
f"• обычная очередь — {normal} звёзд;\n"
f"• премиум-очередь — {premium} звёзд.\n\n"
+ START_TEXT
)
return START_TEXT
async def _start(
@ -156,11 +172,13 @@ class MpakoBot:
await message.reply_text("Активных задач нет.")
return
if active.status == "running":
await message.reply_text("Твоя картинка сейчас генерируется.")
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"Задача в очереди. Перед тобой: {position} {self._people(position)}."
f"Задача в {label}. Перед тобой: {position} {self._people(position)}."
)
async def _pre_checkout(
@ -187,6 +205,74 @@ class MpakoBot:
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"}:
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
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:
@ -231,6 +317,7 @@ class MpakoBot:
skip_payment=True,
payment_charge_id=charge_id,
payment_payload=payload,
is_premium=pending.is_premium,
)
async def _text(
@ -376,6 +463,7 @@ class MpakoBot:
skip_payment: bool = False,
payment_charge_id: str | None = None,
payment_payload: str | None = None,
is_premium: bool = False,
) -> bool:
message = update.effective_message
user = update.effective_user
@ -414,8 +502,18 @@ class MpakoBot:
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
amount = self._settings.stars_price
await self._db.create_pending_payment(
payload=payload,
user_id=user.id,
@ -423,28 +521,44 @@ class MpakoBot:
prompt=prompt,
username=user.username,
input_file_path=input_path,
amount=amount,
amount=0,
)
keyboard = InlineKeyboardMarkup(
[
[
InlineKeyboardButton(
f"Обычная — {self._settings.stars_price} звёзд",
callback_data=f"queue:{payload}:normal",
)
],
[
InlineKeyboardButton(
f"Премиум — {self._settings.premium_stars_price} звёзд",
callback_data=f"queue:{payload}:premium",
)
],
]
)
text = (
f"Обычная очередь: {counts['normal']} чел.\n"
f"Премиум-очередь: {counts['premium']} чел.\n\n"
"Выбери очередь:"
)
try:
await self._retry(
self._bot.send_invoice,
chat_id=user.id,
title="Генерация изображения",
description="Одна генерация или редактирование изображения.",
payload=payload,
provider_token="",
currency="XTR",
prices=[LabeledPrice(label="Генерация", amount=amount)],
message.reply_text,
text,
reply_markup=keyboard,
)
except TelegramError:
logger.exception("Не удалось выставить счёт пользователю %s", user.id)
logger.exception("Не удалось показать выбор очереди %s", user.id)
await self._db.mark_pending_cancelled(payload)
await message.reply_text(
"Не удалось выставить счёт. Попробуй позже."
)
return False
return True
if is_admin:
is_premium = True
try:
request, position = await self._queue.enqueue(
user_id=user.id,
@ -455,6 +569,7 @@ class MpakoBot:
has_person=has_person,
is_admin=is_admin,
payment_charge_id=payment_charge_id,
is_premium=is_premium,
)
except PromptBlocked:
note = await self._refund_if_paid(
@ -497,7 +612,7 @@ class MpakoBot:
try:
status_message = await self._retry(
message.reply_text, self._ack_text(position)
message.reply_text, self._ack_text(position, is_premium)
)
except TelegramError:
logger.exception(
@ -553,10 +668,11 @@ class MpakoBot:
logger.exception("Не удалось вернуть звёзды пользователю %s", user_id)
return False
def _ack_text(self, position: int) -> str:
def _ack_text(self, position: int, is_premium: bool = False) -> str:
label = "Премиум-очередь. " if is_premium else ""
if position <= 0:
return "Принято. Ты первый, начинаю…"
return f"Принято. Перед тобой: {position} {self._people(position)}."
return f"Принято. {label}Ты первый, начинаю…"
return f"Принято. {label}Перед тобой: {position} {self._people(position)}."
@staticmethod
def _people(count: int) -> str:

View File

@ -8,7 +8,7 @@
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16).
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
- 46 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, вспомогательные функции ComfyUI.
- 48 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, вспомогательные функции ComfyUI.
- Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
## 0. Допущения и решения по умолчанию
@ -332,6 +332,7 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
## 16. Оплата Stars
- [x] `PAYMENTS_ENABLED` / `STARS_PRICE`, админы бесплатно.
- [x] Премиум-очередь: `PREMIUM_STARS_PRICE`, приоритетное выполнение, счётчики очередей до оплаты.
- [x] Invoice `XTR` + `pre_checkout_query` + `successful_payment`.
- [x] `pending_payments`: payload, сумма, статус, charge_id, request_id.
- [x] Возврат звёзд при ошибке постановки в очередь и генерации.

View File

@ -96,6 +96,7 @@ class Settings:
blocked_public_figures: frozenset[str] = field(default_factory=frozenset)
payments_enabled: bool = True
stars_price: int = 5
premium_stars_price: int = 15
web_page_size: int = 50
web_bind_ip: str = "127.0.0.1"
@ -156,6 +157,7 @@ class Settings:
blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"),
payments_enabled=_bool("PAYMENTS_ENABLED", True),
stars_price=max(1, _int("STARS_PRICE", 5)),
premium_stars_price=max(1, _int("PREMIUM_STARS_PRICE", 15)),
web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))),
web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip()
or "127.0.0.1",

83
db.py
View File

@ -39,6 +39,7 @@ class Request:
comfy_prompt_id: str | None
payment_charge_id: str | None
payment_refunded: bool
is_premium: bool
mode: str
seed: int | None
error_text: str | None
@ -63,6 +64,7 @@ class Request:
comfy_prompt_id=row["comfy_prompt_id"],
payment_charge_id=row["payment_charge_id"],
payment_refunded=bool(row["payment_refunded"]),
is_premium=bool(row["is_premium"]),
mode=row["mode"],
seed=row["seed"],
error_text=row["error_text"],
@ -116,6 +118,7 @@ class PendingPayment:
input_file_path: str | None
amount: int
status: str
is_premium: bool
telegram_payment_charge_id: str | None
request_id: int | None
@ -131,6 +134,7 @@ class PendingPayment:
input_file_path=row["input_file_path"],
amount=int(row["amount"]),
status=row["status"],
is_premium=bool(row["is_premium"]),
telegram_payment_charge_id=row["telegram_payment_charge_id"],
request_id=row["request_id"],
)
@ -151,6 +155,7 @@ CREATE TABLE IF NOT EXISTS requests (
comfy_prompt_id TEXT,
payment_charge_id TEXT,
payment_refunded INTEGER NOT NULL DEFAULT 0,
is_premium INTEGER NOT NULL DEFAULT 0,
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
seed INTEGER,
error_text TEXT,
@ -190,6 +195,7 @@ CREATE TABLE IF NOT EXISTS pending_payments (
amount INTEGER NOT NULL,
status TEXT NOT NULL DEFAULT 'pending'
CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')),
is_premium INTEGER NOT NULL DEFAULT 0,
telegram_payment_charge_id TEXT,
request_id INTEGER
);
@ -234,6 +240,16 @@ class Database:
await conn.execute(
"ALTER TABLE requests ADD COLUMN payment_refunded INTEGER NOT NULL DEFAULT 0"
)
if "is_premium" not in columns:
await conn.execute(
"ALTER TABLE requests ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0"
)
async with conn.execute("PRAGMA table_info(pending_payments)") as cursor:
pending_columns = {row["name"] async for row in cursor}
if "is_premium" not in pending_columns:
await conn.execute(
"ALTER TABLE pending_payments ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0"
)
async def close(self) -> None:
if self._conn is not None:
@ -255,6 +271,7 @@ class Database:
input_file_path: str | None,
seed: int,
payment_charge_id: str | None = None,
is_premium: bool = False,
) -> Request:
conn = self._connection()
async with self._write_lock:
@ -263,8 +280,9 @@ class Database:
"""
INSERT INTO requests
(user_id, username, timestamp, prompt_text,
input_file_path, status, mode, seed, payment_charge_id)
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?)
input_file_path, status, mode, seed, payment_charge_id,
is_premium)
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?, ?)
""",
(
user_id,
@ -275,6 +293,7 @@ class Database:
mode,
seed,
payment_charge_id,
int(is_premium),
),
)
request_id = cursor.lastrowid
@ -418,16 +437,42 @@ class Database:
async def count_ahead(self, request_id: int) -> int:
conn = self._connection()
async with conn.execute(
"""
request = await self.get_request(request_id)
if request is None:
return 0
if request.is_premium:
query = """
SELECT COUNT(*) AS total FROM requests
WHERE id < ? AND status IN ('queued', 'running')
""",
(request_id,),
) as cursor:
WHERE is_premium = 1 AND status IN ('queued', 'running')
AND id < ?
"""
params = (request_id,)
else:
query = """
SELECT COUNT(*) AS total FROM requests
WHERE status IN ('queued', 'running')
AND (is_premium = 1 OR (is_premium = 0 AND id < ?))
"""
params = (request_id,)
async with conn.execute(query, params) as cursor:
row = await cursor.fetchone()
return int(row["total"]) if row is not None else 0
async def queue_counts(self) -> dict[str, int]:
conn = self._connection()
counts = {"premium": 0, "normal": 0}
async with conn.execute(
"""
SELECT is_premium, COUNT(*) AS total FROM requests
WHERE status IN ('queued', 'running')
GROUP BY is_premium
"""
) as cursor:
async for row in cursor:
key = "premium" if row["is_premium"] else "normal"
counts[key] = int(row["total"])
return counts
async def set_running(self, request_id: int, comfy_prompt_id: str | None = None) -> None:
await self._update(
"""
@ -468,6 +513,7 @@ class Database:
username: str | None,
input_file_path: str | None,
amount: int,
is_premium: bool = False,
) -> PendingPayment:
conn = self._connection()
async with self._write_lock:
@ -475,8 +521,8 @@ class Database:
"""
INSERT INTO pending_payments
(payload, user_id, timestamp, mode, prompt, username,
input_file_path, amount, status)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending')
input_file_path, amount, status, is_premium)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?)
""",
(
payload,
@ -487,6 +533,7 @@ class Database:
username,
input_file_path,
amount,
int(is_premium),
),
)
await conn.commit()
@ -503,6 +550,22 @@ class Database:
row = await cursor.fetchone()
return PendingPayment.from_row(row) if row is not None else None
async def set_pending_choice(
self, payload: str, is_premium: bool, amount: int
) -> bool:
conn = self._connection()
async with self._write_lock:
cursor = await conn.execute(
"""
UPDATE pending_payments
SET is_premium = ?, amount = ?
WHERE payload = ? AND status = 'pending' AND amount = 0
""",
(int(is_premium), amount, payload),
)
await conn.commit()
return cursor.rowcount > 0
async def mark_pending_paid(self, payload: str, charge_id: str) -> bool:
conn = self._connection()
async with self._write_lock:

View File

@ -61,7 +61,8 @@ class QueueManager:
self._on_event = on_event
self._enhancer = enhancer
self._queue: deque[int] = deque()
self._premium: deque[int] = deque()
self._normal: deque[int] = deque()
self._active_users: set[int] = set()
self._wake = asyncio.Event()
self._lock = asyncio.Lock()
@ -74,7 +75,8 @@ class QueueManager:
recovered = await self._db.recover()
async with self._lock:
for request in recovered:
self._queue.append(request.id)
target = self._premium if request.is_premium else self._normal
target.append(request.id)
self._active_users.add(request.user_id)
if recovered:
logger.info("Восстановлено из очереди: %s", len(recovered))
@ -108,6 +110,7 @@ class QueueManager:
has_person: bool = False,
is_admin: bool = False,
payment_charge_id: str | None = None,
is_premium: bool = False,
) -> tuple[Request, int]:
async with self._lock:
if self._stopping:
@ -131,7 +134,7 @@ class QueueManager:
active = await self._db.get_active_by_user(user_id)
if active is not None:
raise UserAlreadyQueued(user_id)
if len(self._queue) >= self._settings.max_queue_size:
if len(self._premium) + len(self._normal) >= self._settings.max_queue_size:
raise QueueFull()
seed = secrets.randbelow(2**32)
@ -143,13 +146,18 @@ class QueueManager:
input_file_path=input_file_path,
seed=seed,
payment_charge_id=payment_charge_id,
is_premium=is_premium,
)
self._queue.append(request.id)
target = self._premium if is_premium else self._normal
target.append(request.id)
self._active_users.add(user_id)
self._wake.set()
position = await self._db.count_ahead(request.id)
return request, position
async def queue_counts(self) -> dict[str, int]:
return await self._db.queue_counts()
async def position(self, request_id: int) -> int:
return await self._db.count_ahead(request_id)
@ -166,8 +174,10 @@ class QueueManager:
async def _next_request(self) -> int | None:
while not self._stopping:
if self._queue:
return self._queue.popleft()
if self._premium:
return self._premium.popleft()
if self._normal:
return self._normal.popleft()
self._wake.clear()
await self._wake.wait()
return None

View File

@ -165,6 +165,32 @@ async def test_history_filters_and_stats(db: Database) -> None:
assert alice.display_user == "@alice"
async def test_premium_queue_priority_counts(db: Database) -> None:
normal = await db.create_request(
user_id=1,
mode="t2i",
prompt_text="normal",
input_file_path=None,
seed=1,
)
premium = await db.create_request(
user_id=2,
mode="t2i",
prompt_text="premium",
input_file_path=None,
seed=2,
is_premium=True,
)
counts = await db.queue_counts()
assert counts == {"premium": 1, "normal": 1}
assert await db.count_ahead(premium.id) == 0
assert await db.count_ahead(normal.id) == 1
await db.set_done(normal.id, "/tmp/normal.png")
counts = await db.queue_counts()
assert counts == {"premium": 1, "normal": 0}
async def test_pending_payment_lifecycle(db: Database) -> None:
pending = await db.create_pending_payment(
payload="payload-1",
@ -173,10 +199,16 @@ async def test_pending_payment_lifecycle(db: Database) -> None:
prompt="apple",
username="alice",
input_file_path=None,
amount=5,
amount=0,
)
assert pending.status == "pending"
assert pending.amount == 5
assert pending.amount == 0
assert await db.set_pending_choice("payload-1", True, 15) is True
loaded = await db.get_pending_payment("payload-1")
assert loaded is not None
assert loaded.is_premium is True
assert loaded.amount == 15
assert await db.mark_pending_paid("payload-1", "charge-1") is True
loaded = await db.get_pending_payment("payload-1")

View File

@ -251,3 +251,22 @@ async def test_admin_prompt_is_not_filtered(db: Database, tmp_path: Path) -> Non
assert await db.list_blocked() == []
await queue.stop()
async def test_premium_queue_goes_first(db: Database, tmp_path: Path) -> None:
events: list[QueueEvent] = []
queue = make_queue(db, make_settings(tmp_path), events)
normal, _ = await queue.enqueue(
user_id=1, mode="t2i", prompt="normal"
)
premium, _ = await queue.enqueue(
user_id=2, mode="t2i", prompt="premium", is_premium=True
)
first = await queue._next_request() # type: ignore[attr-defined]
second = await queue._next_request() # type: ignore[attr-defined]
assert first == premium.id
assert second == normal.id
await queue.stop()

View File

@ -11,6 +11,8 @@
<dd><span class="badge {{ item.status }}">{{ item.status }}</span></dd>
<dt>Режим</dt>
<dd class="mono">{{ item.mode }}</dd>
<dt>Очередь</dt>
<dd>{% if item.is_premium %}премиум{% else %}обычная{% endif %}</dd>
<dt>Создан</dt>
<dd>{{ item.timestamp|fmt_dt }}</dd>
<dt>Завершён</dt>

View File

@ -77,7 +77,7 @@
</td>
<td class="nowrap">{{ item.timestamp|fmt_dt }}</td>
<td><a href="/?user_id={{ item.user_id }}" title="ID: {{ item.user_id }}">{{ item.display_user }}</a></td>
<td class="mono">{{ item.mode }}</td>
<td class="mono">{{ item.mode }}{% if item.is_premium %} <span class="mark">P</span>{% endif %}</td>
<td><span class="badge {{ item.status }}">{{ item.status }}</span></td>
<td class="prompt" title="{{ item.prompt_text }}">
{% if item.enhanced_prompt %}