Премиум-очередь
Build and push Docker image / docker (push) Successful in 10s
Details
Build and push Docker image / docker (push) Successful in 10s
Details
This commit is contained in:
parent
dc16bb7e0c
commit
71a7cd5054
|
|
@ -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
|
||||
|
|
|
|||
11
README.md
11
README.md
|
|
@ -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
162
bot.py
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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] Возврат звёзд при ошибке постановки в очередь и генерации.
|
||||
|
|
|
|||
|
|
@ -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
83
db.py
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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 %}
|
||||
|
|
|
|||
Loading…
Reference in New Issue