diff --git a/README.md b/README.md index d610427..df7838f 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ Telegram-бот для генерации и редактирования изо - **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов. - **Оплата**: 5 Telegram Stars за обычную очередь, 15 за премиум; админам бесплатно, возврат при ошибке. - **Премиум-очередь**: чередуется с обычной (одна премиум, одна обычная); если одна пуста — работает другая. +- **Баланс и рефералы**: 3 стартовые генерации, +1 ежедневно в 00:00 МСК, +3 за первую оплату друга. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - SQLite: запросы, статусы, seed, пути к файлам. @@ -123,6 +124,18 @@ BLOCKED_PUBLIC_FIGURES= При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`. +## Баланс и рефералы + +- Новому пользователю начисляется 3 генерации. +- Каждый день в 00:00 МСК — ещё 1 генерация. +- Обычная генерация может быть оплачена 1 генерацией с баланса или 5 звёздами. +- Премиум-очередь всегда оплачивается 15 звёздами; баланс при этом не расходуется. +- `/balance` — текущий баланс. +- `/invite` — реферальная ссылка. +- За первую оплату приглашённого друга начисляется +3 генерации, реферер получает уведомление. + +Баланс и рефералы хранятся в таблице `users`. Пропущенные дни начисляются при следующем запуске/полуночи. + ## Оплата Генерация стоит 5 Telegram Stars в обычной очереди и 15 звёзд в премиум-очереди: diff --git a/bot.py b/bot.py index 57e3c27..5dd3b48 100644 --- a/bot.py +++ b/bot.py @@ -4,9 +4,12 @@ import asyncio import logging import time import uuid +from datetime import datetime, timedelta, timezone +from datetime import time as dt_time from pathlib import Path from telegram import ( + CopyTextButton, InlineKeyboardButton, InlineKeyboardMarkup, InputFile, @@ -34,12 +37,15 @@ from queue_manager import QueueEvent, QueueFull, QueueManager from workflows import WorkflowBuilder logger = logging.getLogger(__name__) +MSK = timezone(timedelta(hours=3)) START_TEXT = ( "Пришли текст — сгенерирую картинку.\n" "Пришли фото с подписью — отредактирую его по описанию.\n" "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n" - "/status — узнать состояние текущей задачи." + "/balance — баланс генераций.\n" + "/invite — пригласить друга.\n" + "/status — состояние текущей задачи." ) BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны." @@ -54,6 +60,7 @@ class MpakoBot: 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 def run(self) -> None: builder = ( @@ -75,6 +82,8 @@ 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(CommandHandler("balance", self._balance)) + application.add_handler(CommandHandler("invite", self._invite)) application.add_handler( PreCheckoutQueryHandler(self._pre_checkout) ) @@ -104,6 +113,10 @@ class MpakoBot: self._settings.comfy_ws_url, ) await self._db.connect() + await self._db.accrue_daily_credits() + self._daily_task = asyncio.create_task( + self._daily_credits_loop(), name="daily-credits" + ) enhancer = ( PromptEnhancer(self._comfy, self._settings) if self._settings.prompt_enhance_enabled @@ -121,6 +134,10 @@ class MpakoBot: logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url) async def _post_shutdown(self, application: Application) -> 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: @@ -128,17 +145,55 @@ class MpakoBot: await self._db.close() logger.info("Бот остановлен") - def _start_text(self) -> str: + 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("Ошибка ежедневного начисления генераций") + + 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 - return ( + parts.append( "Стоимость генерации:\n" f"• обычная очередь — {normal} звёзд;\n" - f"• премиум-очередь — {premium} звёзд.\n\n" - + START_TEXT + f"• премиум-очередь — {premium} звёзд." ) - return START_TEXT + 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 @@ -146,8 +201,19 @@ class MpakoBot: if not await self._check_access(update): return message = update.effective_message - if message is not None: - await message.reply_text(self._start_text()) + 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) + await message.reply_text(self._start_text(row.credits)) async def _help( self, update: Update, context: ContextTypes.DEFAULT_TYPE @@ -155,8 +221,57 @@ class MpakoBot: if not await self._check_access(update): return message = update.effective_message - if message is not None: - await message.reply_text(self._start_text()) + user = update.effective_user + if message is None or user is None: + return + row = await self._db.ensure_user(user.id, user.username) + await 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 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 message.reply_text( + "Пригласи друга:\n" + f"{link}\n\n" + "Когда друг впервые оплатит генерацию, ты получишь +3 генерации.", + reply_markup=keyboard, + ) async def _status( self, update: Update, context: ContextTypes.DEFAULT_TYPE @@ -181,6 +296,27 @@ class MpakoBot: 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: @@ -216,7 +352,7 @@ class MpakoBot: except ValueError: await query.answer("Некорректный выбор.", show_alert=True) return - if tier not in {"normal", "premium"}: + if tier not in {"normal", "premium", "credit"}: await query.answer("Некорректный выбор.", show_alert=True) return @@ -237,6 +373,39 @@ class MpakoBot: ) 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, + ) + 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 @@ -309,7 +478,7 @@ class MpakoBot: await message.reply_text("Этот платёж уже обработан.") return - await self._submit( + accepted = await self._submit( update, mode=pending.mode, prompt=pending.prompt, @@ -319,6 +488,8 @@ class MpakoBot: payment_payload=payload, is_premium=pending.is_premium, ) + if accepted: + await self._reward_referrer(user) async def _text( self, update: Update, context: ContextTypes.DEFAULT_TYPE @@ -471,6 +642,8 @@ class MpakoBot: return False is_admin = user.id in self._settings.admin_user_ids + user_row = await self._db.ensure_user(user.id, user.username) + credits = user_row.credits has_person = False if ( mode == "i2i" @@ -527,8 +700,16 @@ class MpakoBot: [ [ InlineKeyboardButton( - f"Обычная — {self._settings.stars_price} звёзд", - callback_data=f"queue:{payload}:normal", + ( + 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" + ), ) ], [ @@ -540,6 +721,7 @@ class MpakoBot: ] ) text = ( + f"Баланс: {credits} {self._credits_word(credits)}.\n" f"Обычная очередь: {counts['normal']} чел.\n" f"Премиум-очередь: {counts['premium']} чел.\n\n" "Выбери очередь:" @@ -557,7 +739,7 @@ class MpakoBot: return True if is_admin: - is_premium = True + is_premium = False try: request, position = await self._queue.enqueue( diff --git a/checklist.md b/checklist.md index 1d8caaa..5063a2f 100644 --- a/checklist.md +++ b/checklist.md @@ -6,10 +6,10 @@ ## Текущий статус -- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16). +- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15), оплата Stars (16) и баланс/рефералы (17). - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. -- 49 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, вспомогательные функции ComfyUI. -- Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. +- 50 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, вспомогательные функции ComfyUI. +- Не проверено: реальная оплата Stars и реферальная выплата в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. ## 0. Допущения и решения по умолчанию @@ -341,6 +341,16 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); **DoD:** не-админ получает счёт на 5 звёзд, после оплаты задача ставится в очередь; при сбое звёзды возвращаются. +## 17. Баланс и рефералы + +- [x] `users`: 3 стартовые генерации, ежедневное +1 в 00:00 МСК. +- [x] Списание 1 генерации за обычную очередь, премиум — только звёзды. +- [x] `/balance`, `/invite`, реферальная ссылка `?start=ref_`. +- [x] +3 генерации рефереру за первую оплату друга и уведомление. +- [x] Тесты баланса, ежедневного начисления и рефералов. + +**DoD:** новый пользователь получает 3 генерации, ежедневно +1, друг после первой оплаты приносит рефереру +3. + ## Порядок работ 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow. diff --git a/db.py b/db.py index e8034b8..89e68ce 100644 --- a/db.py +++ b/db.py @@ -4,7 +4,7 @@ import asyncio import logging import sqlite3 from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import UTC, date, datetime, timedelta, timezone from pathlib import Path from typing import Any @@ -13,12 +13,17 @@ import aiosqlite logger = logging.getLogger(__name__) ACTIVE_STATUSES = ("queued", "running") +MSK = timezone(timedelta(hours=3)) def utc_now() -> str: return datetime.now(UTC).isoformat() +def msk_today() -> date: + return datetime.now(MSK).date() + + class UserAlreadyQueued(RuntimeError): def __init__(self, user_id: int) -> None: super().__init__(f"У пользователя {user_id} уже есть активный запрос") @@ -107,6 +112,33 @@ class BlockedRequest: ) +@dataclass(frozen=True) +class User: + user_id: int + username: str | None + credits: int + last_daily_credit: str + referrer_id: int | None + first_paid_at: str | None + created_at: str + + @property + def display_user(self) -> str: + return f"@{self.username}" if self.username else str(self.user_id) + + @classmethod + def from_row(cls, row: aiosqlite.Row) -> User: + return cls( + user_id=int(row["user_id"]), + username=row["username"], + credits=int(row["credits"]), + last_daily_credit=row["last_daily_credit"], + referrer_id=row["referrer_id"], + first_paid_at=row["first_paid_at"], + created_at=row["created_at"], + ) + + @dataclass(frozen=True) class PendingPayment: payload: str @@ -202,6 +234,19 @@ CREATE TABLE IF NOT EXISTS pending_payments ( CREATE INDEX IF NOT EXISTS idx_pending_user_time ON pending_payments(user_id, timestamp); + +CREATE TABLE IF NOT EXISTS users ( + user_id INTEGER PRIMARY KEY, + username TEXT, + credits INTEGER NOT NULL DEFAULT 3, + last_daily_credit TEXT NOT NULL, + referrer_id INTEGER, + first_paid_at TEXT, + created_at TEXT NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_users_referrer + ON users(referrer_id); """ @@ -250,6 +295,16 @@ class Database: await conn.execute( "ALTER TABLE pending_payments ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0" ) + await conn.execute( + """ + INSERT OR IGNORE INTO users + (user_id, username, credits, last_daily_credit, + referrer_id, first_paid_at, created_at) + SELECT user_id, NULL, 3, ?, NULL, NULL, ? + FROM (SELECT DISTINCT user_id FROM requests) + """, + (msk_today().isoformat(), utc_now()), + ) async def close(self) -> None: if self._conn is not None: @@ -327,6 +382,115 @@ class Database: row = await cursor.fetchone() return Request.from_row(row) if row is not None else None + async def get_user(self, user_id: int) -> User | None: + conn = self._connection() + async with conn.execute( + "SELECT * FROM users WHERE user_id = ?", (user_id,) + ) as cursor: + row = await cursor.fetchone() + return User.from_row(row) if row is not None else None + + async def ensure_user( + self, user_id: int, username: str | None = None + ) -> User: + conn = self._connection() + today = msk_today().isoformat() + async with self._write_lock: + await conn.execute( + """ + INSERT OR IGNORE INTO users + (user_id, username, credits, last_daily_credit, created_at) + VALUES (?, ?, 3, ?, ?) + """, + (user_id, username, today, utc_now()), + ) + if username: + await conn.execute( + "UPDATE users SET username = ? WHERE user_id = ?", + (username, user_id), + ) + await conn.commit() + user = await self.get_user(user_id) + if user is None: + raise RuntimeError("Не удалось прочитать пользователя") + return user + + async def add_credits(self, user_id: int, amount: int) -> int | None: + conn = self._connection() + async with self._write_lock: + cursor = await conn.execute( + "UPDATE users SET credits = credits + ? WHERE user_id = ?", + (amount, user_id), + ) + await conn.commit() + if cursor.rowcount == 0: + return None + user = await self.get_user(user_id) + return user.credits if user is not None else None + + async def consume_credit(self, user_id: int) -> int | None: + conn = self._connection() + async with self._write_lock: + cursor = await conn.execute( + "UPDATE users SET credits = credits - 1" + " WHERE user_id = ? AND credits > 0", + (user_id,), + ) + await conn.commit() + if cursor.rowcount == 0: + return None + user = await self.get_user(user_id) + return user.credits if user is not None else None + + async def set_referrer(self, user_id: int, referrer_id: int) -> bool: + conn = self._connection() + async with self._write_lock: + cursor = await conn.execute( + """ + UPDATE users SET referrer_id = ? + WHERE user_id = ? AND referrer_id IS NULL AND user_id != ? + """, + (referrer_id, user_id, referrer_id), + ) + await conn.commit() + return cursor.rowcount > 0 + + async def mark_first_payment(self, user_id: int) -> bool: + conn = self._connection() + async with self._write_lock: + cursor = await conn.execute( + "UPDATE users SET first_paid_at = ?" + " WHERE user_id = ? AND first_paid_at IS NULL", + (utc_now(), user_id), + ) + await conn.commit() + return cursor.rowcount > 0 + + async def accrue_daily_credits(self, today: date | None = None) -> int: + conn = self._connection() + current = today or msk_today() + current_text = current.isoformat() + async with self._write_lock: + async with conn.execute( + "SELECT user_id, last_daily_credit FROM users" + " WHERE last_daily_credit < ?", + (current_text,), + ) as cursor: + rows = await cursor.fetchall() + for row in rows: + try: + last = date.fromisoformat(row["last_daily_credit"]) + days = max(1, (current - last).days) + except ValueError: + days = 1 + await conn.execute( + "UPDATE users SET credits = credits + ?," + " last_daily_credit = ? WHERE user_id = ?", + (days, current_text, int(row["user_id"])), + ) + await conn.commit() + return len(rows) + async def list_queued(self) -> list[Request]: conn = self._connection() async with conn.execute( diff --git a/tests/test_db.py b/tests/test_db.py index bea616a..3182f4e 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -1,6 +1,7 @@ from __future__ import annotations import sqlite3 +from datetime import date, timedelta import pytest import pytest_asyncio @@ -242,6 +243,35 @@ async def test_pending_payment_lifecycle(db: Database) -> None: assert loaded_request.payment_refunded is True +async def test_user_credits_referral_and_daily_accrual(db: Database) -> None: + user = await db.ensure_user(1, "alice") + assert user.credits == 3 + assert user.last_daily_credit + + assert await db.add_credits(1, 2) == 5 + assert await db.consume_credit(1) == 4 + assert await db.consume_credit(1) == 3 + + friend = await db.ensure_user(2, "bob") + assert friend.credits == 3 + assert await db.set_referrer(2, 1) is True + assert await db.set_referrer(2, 1) is False + assert await db.mark_first_payment(2) is True + assert await db.mark_first_payment(2) is False + + friend = await db.get_user(2) + assert friend is not None + assert friend.referrer_id == 1 + assert friend.first_paid_at is not None + + future = date.today() + timedelta(days=2) + assert await db.accrue_daily_credits(future) >= 2 + user = await db.get_user(1) + assert user is not None + assert user.credits == 5 + assert await db.accrue_daily_credits(future) == 0 + + async def test_migrates_legacy_schema_without_username(tmp_path) -> None: path = tmp_path / "legacy.db" connection = sqlite3.connect(path)