From 2161bbe0152b42b9387808113beb1df4c952b0ed Mon Sep 17 00:00:00 2001 From: MpakobecEx Date: Sun, 18 Oct 2026 20:09:46 +0500 Subject: [PATCH] =?UTF-8?q?=D0=9E=D0=BF=D0=BB=D0=B0=D1=82=D0=B0=20=D0=B3?= =?UTF-8?q?=D0=B5=D0=BD=D0=B5=D1=80=D0=B0=D1=86=D0=B8=D0=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 4 + README.md | 17 ++++ bot.py | 217 +++++++++++++++++++++++++++++++++++++++++++++-- checklist.md | 17 +++- config.py | 4 + db.py | 153 ++++++++++++++++++++++++++++++++- queue_manager.py | 2 + tests/test_db.py | 45 ++++++++++ 8 files changed, 445 insertions(+), 14 deletions(-) diff --git a/.env.example b/.env.example index 6277596..1ed16aa 100644 --- a/.env.example +++ b/.env.example @@ -40,6 +40,10 @@ MODERATION_ENABLED=true # Дополнительные имена публичных людей через запятую BLOCKED_PUBLIC_FIGURES= +# Оплата Telegram Stars +PAYMENTS_ENABLED=true +STARS_PRICE=5 + # Веб-история: только локальная сеть, без авторизации WEB_BIND_IP=192.168.31.158 WEB_PAGE_SIZE=50 diff --git a/README.md b/README.md index 7c75b30..f94070d 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ Telegram-бот для генерации и редактирования изо - **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений. - **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI. - **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов. +- **Оплата**: 5 Telegram Stars за генерацию, админам бесплатно, возврат при ошибке. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - SQLite: запросы, статусы, seed, пути к файлам. @@ -121,6 +122,22 @@ BLOCKED_PUBLIC_FIGURES= При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`. +## Оплата + +Генерация стоит 5 Telegram Stars: + +```dotenv +PAYMENTS_ENABLED=true +STARS_PRICE=5 +``` + +- Админы из `ADMIN_USER_IDS` генерируют бесплатно. +- Не-админу выставляется счёт через Telegram Stars (`XTR`). +- Фильтрация промпта выполняется до выставления счёта. +- Если поставить задачу в очередь не удалось, звёзды возвращаются. +- При ошибке генерации звёзды также возвращаются. +- Оплата и `telegram_payment_charge_id` сохраняются в `pending_payments` и `requests`. + ## Веб-история После `docker compose up -d` доступна страница: diff --git a/bot.py b/bot.py index a5d2023..abaa230 100644 --- a/bot.py +++ b/bot.py @@ -6,7 +6,7 @@ import time import uuid from pathlib import Path -from telegram import InputFile, Update +from telegram import InputFile, LabeledPrice, Update from telegram.constants import ChatAction, ChatType from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut from telegram.ext import ( @@ -14,13 +14,14 @@ from telegram.ext import ( CommandHandler, ContextTypes, MessageHandler, + PreCheckoutQueryHandler, filters, ) from comfy_client import ComfyClient from config import Settings, setup_logging from db import Database, Request, UserAlreadyQueued -from moderation import PromptBlocked, detect_person +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 @@ -67,6 +68,12 @@ class MpakoBot: application.add_handler(CommandHandler("start", self._start)) application.add_handler(CommandHandler("help", self._help)) application.add_handler(CommandHandler("status", self._status)) + application.add_handler( + PreCheckoutQueryHandler(self._pre_checkout) + ) + 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) @@ -111,6 +118,15 @@ class MpakoBot: await self._db.close() logger.info("Бот остановлен") + def _start_text(self) -> str: + text = START_TEXT + if self._settings.payments_enabled: + text += ( + f"\n\nСтоимость генерации: {self._settings.stars_price} звёзд. " + "Админам бесплатно." + ) + return text + async def _start( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: @@ -118,7 +134,7 @@ class MpakoBot: return message = update.effective_message if message is not None: - await message.reply_text(START_TEXT) + await message.reply_text(self._start_text()) async def _help( self, update: Update, context: ContextTypes.DEFAULT_TYPE @@ -127,7 +143,7 @@ class MpakoBot: return message = update.effective_message if message is not None: - await message.reply_text(START_TEXT) + await message.reply_text(self._start_text()) async def _status( self, update: Update, context: ContextTypes.DEFAULT_TYPE @@ -150,6 +166,76 @@ class MpakoBot: f"Задача в очереди. Перед тобой: {position} {self._people(position)}." ) + 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) + 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 + ) + try: + await query.answer( + ok=allowed, + error_message=None + if allowed + else "Счёт устарел. Отправь запрос заново.", + ) + except TelegramError: + logger.exception("Не удалось ответить на pre_checkout_query") + + 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 + + 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 + + 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, + ) + async def _text( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: @@ -290,6 +376,9 @@ class MpakoBot: mode: str, prompt: str, input_path: str | None, + skip_payment: bool = False, + payment_charge_id: str | None = None, + payment_payload: str | None = None, ) -> bool: message = update.effective_message user = update.effective_user @@ -306,6 +395,59 @@ class MpakoBot: ): 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 + + payload = uuid.uuid4().hex + amount = self._settings.stars_price + 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=amount, + ) + 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)], + ) + except TelegramError: + logger.exception("Не удалось выставить счёт пользователю %s", user.id) + await self._db.mark_pending_cancelled(payload) + await message.reply_text( + "Не удалось выставить счёт. Попробуй позже." + ) + return False + return True + try: request, position = await self._queue.enqueue( user_id=user.id, @@ -315,11 +457,18 @@ class MpakoBot: input_file_path=input_path, has_person=has_person, is_admin=is_admin, + payment_charge_id=payment_charge_id, ) except PromptBlocked: - await message.reply_text(BLOCKED_TEXT) + 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: @@ -328,16 +477,27 @@ class MpakoBot: else: ahead = await self._queue.position(active.id) text += f" Перед тобой: {ahead} {self._people(ahead)}." - await message.reply_text(text) + await message.reply_text(text + note) return False except QueueFull as exc: - await message.reply_text(exc.user_message) + 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) - await message.reply_text("Не удалось поставить задачу. Попробуй позже.") + 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) @@ -367,6 +527,35 @@ class MpakoBot: ) 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) -> str: if position <= 0: return "Принято. Ты первый, начинаю…" @@ -513,7 +702,17 @@ class MpakoBot: 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 - text = f"Не получилось: {error}" + 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( diff --git a/checklist.md b/checklist.md index 3be93ac..940227f 100644 --- a/checklist.md +++ b/checklist.md @@ -6,10 +6,10 @@ ## Текущий статус -- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14) и модерация (15). +- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16). - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. -- 45 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, вспомогательные функции ComfyUI. -- Не проверено: очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. +- 46 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, вспомогательные функции ComfyUI. +- Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. ## 0. Допущения и решения по умолчанию @@ -329,6 +329,17 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); **DoD:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра. +## 16. Оплата Stars + +- [x] `PAYMENTS_ENABLED` / `STARS_PRICE`, админы бесплатно. +- [x] Invoice `XTR` + `pre_checkout_query` + `successful_payment`. +- [x] `pending_payments`: payload, сумма, статус, charge_id, request_id. +- [x] Возврат звёзд при ошибке постановки в очередь и генерации. +- [x] `payment_charge_id` и `payment_refunded` в `requests`. +- [x] Тесты жизненного цикла платежа. + +**DoD:** не-админ получает счёт на 5 звёзд, после оплаты задача ставится в очередь; при сбое звёзды возвращаются. + ## Порядок работ 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow. diff --git a/config.py b/config.py index 56eca36..10bf059 100644 --- a/config.py +++ b/config.py @@ -94,6 +94,8 @@ class Settings: max_photo_bytes: int = 15 * 1024 * 1024 moderation_enabled: bool = True blocked_public_figures: frozenset[str] = field(default_factory=frozenset) + payments_enabled: bool = True + stars_price: int = 5 web_page_size: int = 50 web_bind_ip: str = "127.0.0.1" @@ -152,6 +154,8 @@ class Settings: max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)), moderation_enabled=_bool("MODERATION_ENABLED", True), blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"), + payments_enabled=_bool("PAYMENTS_ENABLED", True), + stars_price=max(1, _int("STARS_PRICE", 5)), 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", diff --git a/db.py b/db.py index 21bcb6e..7cbbf82 100644 --- a/db.py +++ b/db.py @@ -37,6 +37,8 @@ class Request: output_file_path: str | None status: str comfy_prompt_id: str | None + payment_charge_id: str | None + payment_refunded: bool mode: str seed: int | None error_text: str | None @@ -59,6 +61,8 @@ class Request: output_file_path=row["output_file_path"], status=row["status"], comfy_prompt_id=row["comfy_prompt_id"], + payment_charge_id=row["payment_charge_id"], + payment_refunded=bool(row["payment_refunded"]), mode=row["mode"], seed=row["seed"], error_text=row["error_text"], @@ -101,6 +105,37 @@ class BlockedRequest: ) +@dataclass(frozen=True) +class PendingPayment: + payload: str + user_id: int + timestamp: str + mode: str + prompt: str + username: str | None + input_file_path: str | None + amount: int + status: str + telegram_payment_charge_id: str | None + request_id: int | None + + @classmethod + def from_row(cls, row: aiosqlite.Row) -> PendingPayment: + return cls( + payload=row["payload"], + user_id=row["user_id"], + timestamp=row["timestamp"], + mode=row["mode"], + prompt=row["prompt"], + username=row["username"], + input_file_path=row["input_file_path"], + amount=int(row["amount"]), + status=row["status"], + telegram_payment_charge_id=row["telegram_payment_charge_id"], + request_id=row["request_id"], + ) + + SCHEMA = """ CREATE TABLE IF NOT EXISTS requests ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -114,6 +149,8 @@ CREATE TABLE IF NOT EXISTS requests ( status TEXT NOT NULL DEFAULT 'queued' CHECK (status IN ('queued', 'running', 'done', 'failed')), comfy_prompt_id TEXT, + payment_charge_id TEXT, + payment_refunded INTEGER NOT NULL DEFAULT 0, mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')), seed INTEGER, error_text TEXT, @@ -141,6 +178,24 @@ CREATE TABLE IF NOT EXISTS blocked_requests ( CREATE INDEX IF NOT EXISTS idx_blocked_user_time ON blocked_requests(user_id, id); + +CREATE TABLE IF NOT EXISTS pending_payments ( + payload TEXT PRIMARY KEY, + user_id INTEGER NOT NULL, + timestamp TEXT NOT NULL, + mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')), + prompt TEXT NOT NULL, + username TEXT, + input_file_path TEXT, + amount INTEGER NOT NULL, + status TEXT NOT NULL DEFAULT 'pending' + CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')), + telegram_payment_charge_id TEXT, + request_id INTEGER +); + +CREATE INDEX IF NOT EXISTS idx_pending_user_time + ON pending_payments(user_id, timestamp); """ @@ -173,6 +228,12 @@ class Database: await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT") if "enhanced_prompt" not in columns: await conn.execute("ALTER TABLE requests ADD COLUMN enhanced_prompt TEXT") + if "payment_charge_id" not in columns: + await conn.execute("ALTER TABLE requests ADD COLUMN payment_charge_id TEXT") + if "payment_refunded" not in columns: + await conn.execute( + "ALTER TABLE requests ADD COLUMN payment_refunded INTEGER NOT NULL DEFAULT 0" + ) async def close(self) -> None: if self._conn is not None: @@ -193,6 +254,7 @@ class Database: prompt_text: str, input_file_path: str | None, seed: int, + payment_charge_id: str | None = None, ) -> Request: conn = self._connection() async with self._write_lock: @@ -201,8 +263,8 @@ class Database: """ INSERT INTO requests (user_id, username, timestamp, prompt_text, - input_file_path, status, mode, seed) - VALUES (?, ?, ?, ?, ?, 'queued', ?, ?) + input_file_path, status, mode, seed, payment_charge_id) + VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?) """, ( user_id, @@ -212,6 +274,7 @@ class Database: input_file_path, mode, seed, + payment_charge_id, ), ) request_id = cursor.lastrowid @@ -389,6 +452,92 @@ class Database: (prompt, request_id), ) + async def set_refunded(self, request_id: int) -> None: + await self._update( + "UPDATE requests SET payment_refunded = 1 WHERE id = ?", + (request_id,), + ) + + async def create_pending_payment( + self, + *, + payload: str, + user_id: int, + mode: str, + prompt: str, + username: str | None, + input_file_path: str | None, + amount: int, + ) -> PendingPayment: + conn = self._connection() + async with self._write_lock: + await conn.execute( + """ + INSERT INTO pending_payments + (payload, user_id, timestamp, mode, prompt, username, + input_file_path, amount, status) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending') + """, + ( + payload, + user_id, + utc_now(), + mode, + prompt, + username, + input_file_path, + amount, + ), + ) + await conn.commit() + pending = await self.get_pending_payment(payload) + if pending is None: + raise RuntimeError("Не удалось прочитать счёт") + return pending + + async def get_pending_payment(self, payload: str) -> PendingPayment | None: + conn = self._connection() + async with conn.execute( + "SELECT * FROM pending_payments WHERE payload = ?", (payload,) + ) as cursor: + row = await cursor.fetchone() + return PendingPayment.from_row(row) if row is not None else None + + async def mark_pending_paid(self, payload: str, charge_id: str) -> bool: + conn = self._connection() + async with self._write_lock: + cursor = await conn.execute( + """ + UPDATE pending_payments + SET status = 'paid', telegram_payment_charge_id = ? + WHERE payload = ? AND status = 'pending' + """, + (charge_id, payload), + ) + await conn.commit() + return cursor.rowcount > 0 + + async def mark_pending_refunded(self, payload: str) -> None: + await self._update( + "UPDATE pending_payments SET status = 'refunded' WHERE payload = ?", + (payload,), + ) + + async def mark_pending_cancelled(self, payload: str) -> None: + await self._update( + """ + UPDATE pending_payments SET status = 'cancelled' + WHERE payload = ? AND status = 'pending' + """, + (payload,), + ) + + async def attach_pending_request(self, payload: str, request_id: int) -> None: + await self._update( + "UPDATE pending_payments SET request_id = ? WHERE payload = ?", + (request_id, payload), + ) + async def set_done(self, request_id: int, output_file_path: str) -> None: await self._update( """ diff --git a/queue_manager.py b/queue_manager.py index b515b90..d7a60c0 100644 --- a/queue_manager.py +++ b/queue_manager.py @@ -107,6 +107,7 @@ class QueueManager: username: str | None = None, has_person: bool = False, is_admin: bool = False, + payment_charge_id: str | None = None, ) -> tuple[Request, int]: async with self._lock: if self._stopping: @@ -141,6 +142,7 @@ class QueueManager: prompt_text=prompt, input_file_path=input_file_path, seed=seed, + payment_charge_id=payment_charge_id, ) self._queue.append(request.id) self._active_users.add(user_id) diff --git a/tests/test_db.py b/tests/test_db.py index 1b8416d..23e5a87 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -165,6 +165,51 @@ async def test_history_filters_and_stats(db: Database) -> None: assert alice.display_user == "@alice" +async def test_pending_payment_lifecycle(db: Database) -> None: + pending = await db.create_pending_payment( + payload="payload-1", + user_id=1, + mode="t2i", + prompt="apple", + username="alice", + input_file_path=None, + amount=5, + ) + assert pending.status == "pending" + assert pending.amount == 5 + + assert await db.mark_pending_paid("payload-1", "charge-1") is True + loaded = await db.get_pending_payment("payload-1") + assert loaded is not None + assert loaded.status == "paid" + assert loaded.telegram_payment_charge_id == "charge-1" + + await db.attach_pending_request("payload-1", 42) + loaded = await db.get_pending_payment("payload-1") + assert loaded is not None + assert loaded.request_id == 42 + + await db.mark_pending_refunded("payload-1") + loaded = await db.get_pending_payment("payload-1") + assert loaded is not None + assert loaded.status == "refunded" + + request = await db.create_request( + user_id=2, + mode="t2i", + prompt_text="paid apple", + input_file_path=None, + seed=1, + payment_charge_id="charge-2", + ) + assert request.payment_charge_id == "charge-2" + assert request.payment_refunded is False + await db.set_refunded(request.id) + loaded_request = await db.get_request(request.id) + assert loaded_request is not None + assert loaded_request.payment_refunded is True + + async def test_migrates_legacy_schema_without_username(tmp_path) -> None: path = tmp_path / "legacy.db" connection = sqlite3.connect(path)