diff --git a/configs/bot.json b/configs/bot.json index d7bb064..865e72e 100644 --- a/configs/bot.json +++ b/configs/bot.json @@ -5,8 +5,5 @@ "temperature_range": [0.0, 0.5, 0.8, 1.0, 1.2], "top_p_range": [0.8, 0.9, 0.95, 0.98, 1.0], "timezone": "Europe/Moscow", - "sub_price_rub": 500, - "sub_price_stars": 250, - "sub_duration": 604800, "output_chunk_size": 3500 } diff --git a/configs/localization.json b/configs/localization.json index ad60d51..4b348ca 100644 --- a/configs/localization.json +++ b/configs/localization.json @@ -17,8 +17,10 @@ "ACTIVE_SUB": "Подписка активирована! Осталось {remaining_hours}ч", "REMAINING_MESSAGES": "Осталось запросов к {model}: {remaining_count}", "SET_EMAIL": "Пожалуйста, сначала задайте свою почту через '/setemail ...'. Туда придёт чек.", - "SUB_TITLE": "Покупка подписки в боте 'Сайга' на неделю для пользователя {user_id}", - "SUB_SHORT_TITLE": "Покупка подписки на неделю", + "SUB_WEEK_TITLE": "Покупка подписки в боте 'Сайга' на неделю для пользователя {user_id}", + "SUB_WEEK_SHORT_TITLE": "Покупка подписки на неделю", + "SUB_MONTH_TITLE": "Покупка подписки в боте 'Сайга' на месяц для пользователя {user_id}", + "SUB_MONTH_SHORT_TITLE": "Покупка подписки на месяц", "SUB_NOT_CHAT": "Подписку можно купить только в переписке с самим ботом!", "SUB_SUCCESS": "Спасибо за оплату, подписка выдана! Узнать статус: /subinfo", "MODEL_NOT_SUPPORTED": "Выбранная модель больше не поддерживается, переключите на другую с помощью /setmodel", @@ -46,10 +48,12 @@ "DALLE_LIMIT": "Лимит по генерации картинок исчерпан, восстановится через 24 часа", "DALLE_PROMPT": "Генерирую картинку по промпту: {displayed_prompt}", "DALLE_ERROR": "Ошибка при вызове DALL-E: {error}", - "BUY_WITH_STARS": "Купить (за Telegram Stars)", - "BUY_WITH_RUB": "Купить (за рубли)", + "BUY_WEEK_WITH_STARS": "Купить 7 дней за Telegram Stars", + "BUY_WEEK_WITH_RUB": "Купить 7 дней за рубли", + "BUY_MONTH_WITH_STARS": "Купить 31 день за Telegram Stars", + "BUY_MONTH_WITH_RUB": "Купить 31 день за рубли", "WRONG_COMMAND": "Такой команды у бота нет. Если вы не пытались ввести команду, уберите '/' из начала сообщения.", - "PAY_SUPPORT": "Возврат средств за подписку возможен в течение 12 часов с момента её оформления. Контакт: @YallenGusev", + "PAY_SUPPORT": "Возврат средств за подписку возможен в течение 12 часов с момента её оформления, и если с момента оформления было сделано менее 20 запросов к любым моделями. Контакт: @YallenGusev", "FILE_IS_TOO_BIG": "Слишком большой файл! Телеграм-боты не могут скачивать файлы больше 20 МБ.", "PRIVACY": { "template_name": "ru_privacy" diff --git a/scripts/fetch_textual_conversations.py b/scripts/fetch_textual_conversations.py index 12e515b..736e134 100644 --- a/scripts/fetch_textual_conversations.py +++ b/scripts/fetch_textual_conversations.py @@ -1,6 +1,7 @@ import json import fire +from tqdm import tqdm from src.database import Database @@ -21,10 +22,10 @@ def merge_messages(messages): def main(db_path: str, output_path: str, min_timestamp: int = None, fetch_chats: bool = False): db = Database(db_path) - conversations = set(db.get_all_conv_ids()) + conversations = set(db.get_all_conv_ids(min_timestamp=min_timestamp)) records = [] first_messages = set() - for conv_id in conversations: + for conv_id in tqdm(conversations): messages = db.fetch_conversation(conv_id) timestamps = {m.get("timestamp", None) for m in messages} diff --git a/src/bot.py b/src/bot.py index 277b92d..fba6a0d 100644 --- a/src/bot.py +++ b/src/bot.py @@ -7,6 +7,7 @@ import io import re from email.utils import parseaddr +from enum import Enum from typing import cast, List, Dict, Any, Optional, Union, Callable, Coroutine, Tuple, BinaryIO from dataclasses import dataclass, field @@ -50,6 +51,19 @@ ChatMessage = Dict[str, Any] ChatMessages = List[ChatMessage] +@dataclass +class SubConfig: + price: int = 500 + currency: str = "RUB" + duration: int = 7 * 86400 + + +class SubKey(str, Enum): + RUB_WEEK = "rub_week" + RUB_MONTH = "rub_month" + XTR_WEEK = "xtr_week" + XTR_MONTH = "xtr_month" + @dataclass class BotConfig: @@ -60,10 +74,13 @@ class BotConfig: top_p_range: List[float] freq_penalty_range: List[float] = field(default_factory=lambda: [0.0, 0.05, 0.1, 0.2, 0.5, 1.0]) timezone: str = "Europe/Moscow" - sub_price_rub: int = 500 - sub_price_stars: int = 250 - sub_duration: int = 7 * 86400 output_chunk_size: int = 3500 + sub_configs: Dict[SubKey, SubConfig] = field(default_factory=lambda: { + SubKey.RUB_WEEK: SubConfig(500, "RUB", 7 * 86400), + SubKey.RUB_MONTH: SubConfig(2100, "RUB", 31 * 86400), + SubKey.XTR_WEEK: SubConfig(250, "XTR", 7 * 86400), + SubKey.XTR_MONTH: SubConfig(1000, "XTR", 31 * 86400), + }) def _crop_content(content: str) -> str: @@ -159,7 +176,8 @@ def __init__( self.freq_penalty_kb.add(InlineKeyboardButton(text=str(value), callback_data=f"setfreqpenalty:{value}")) self.buy_kb = InlineKeyboardBuilder() - self.buy_kb.add(InlineKeyboardButton(text=self.localization.BUY_WITH_STARS, callback_data="buy:stars")) + self.buy_kb.add(InlineKeyboardButton(text=self.localization.BUY_WEEK_WITH_STARS, callback_data="buy:stars:xtr_week")) + self.buy_kb.add(InlineKeyboardButton(text=self.localization.BUY_MONTH_WITH_STARS, callback_data="buy:stars:xtr_month")) self.bot = Bot(token=self.config.token, default=DefaultBotProperties(parse_mode=None)) self.bot_info: Optional[User] = None @@ -216,10 +234,12 @@ def __init__( self.scheduler: Optional[AsyncIOScheduler] = None self.yookassa: Optional[YookassaHandler] = None if yookassa_config_path and os.path.exists(yookassa_config_path): - self.buy_kb.add(InlineKeyboardButton(text=self.localization.BUY_WITH_RUB, callback_data="buy:yookassa")) + self.buy_kb.add(InlineKeyboardButton(text=self.localization.BUY_WEEK_WITH_RUB, callback_data="buy:yookassa:rub_week")) + self.buy_kb.add(InlineKeyboardButton(text=self.localization.BUY_MONTH_WITH_RUB, callback_data="buy:yookassa:rub_month")) with open(yookassa_config_path) as r: config = json.load(r) self.yookassa = YookassaHandler(**config) + self.buy_kb.adjust(2) async def start_polling(self) -> None: self.scheduler = AsyncIOScheduler(timezone=self.config.timezone) @@ -532,13 +552,21 @@ async def sub_buy(self, message: Message) -> None: limits = {name: provider.limits for name, provider in self.providers.items()} sub_limits = self.localization.LIMITS.render(limits=limits, mode="subscribed").strip() - description = self.localization.SUB_DESCRIPTION.render(sub_limits=sub_limits, price=self.config.sub_price_rub) + description = self.localization.SUB_DESCRIPTION.render( + sub_limits=sub_limits, + price_week=self.config.sub_configs[SubKey.RUB_WEEK].price, + price_month=self.config.sub_configs[SubKey.RUB_MONTH].price, + ) await message.reply(description, parse_mode=ParseMode.MARKDOWN, reply_markup=self.buy_kb.as_markup()) async def stars_sub_buy_proceed(self, callback: CallbackQuery) -> None: assert callback.from_user assert callback.message assert isinstance(callback.message, Message) + assert callback.data + assert "buy:stars:" in callback.data + + sub_key_str = callback.data.split(":")[2] user_id = callback.from_user.id chat_id = callback.message.chat.id is_chat = chat_id != user_id @@ -551,16 +579,26 @@ async def stars_sub_buy_proceed(self, callback: CallbackQuery) -> None: await callback.message.reply(self.localization.ACTIVE_SUB.format(remaining_hours=remaining_seconds // 3600)) return - title = self.localization.SUB_SHORT_TITLE - description = self.localization.SUB_TITLE.format(user_id=user_id) + key = SubKey(sub_key_str) + key_to_short_title = { + SubKey.XTR_WEEK: self.localization.SUB_WEEK_SHORT_TITLE, + SubKey.XTR_MONTH: self.localization.SUB_MONTH_SHORT_TITLE, + } + key_to_title = { + SubKey.XTR_WEEK: self.localization.SUB_WEEK_TITLE, + SubKey.XTR_MONTH: self.localization.SUB_MONTH_TITLE + } + title = key_to_short_title[key] + description = key_to_title[key].format(user_id=user_id) + sub = self.config.sub_configs[key] await self.bot.send_invoice( chat_id, title=title, description=description, - prices=[LabeledPrice(label=title, amount=self.config.sub_price_stars)], + prices=[LabeledPrice(label=title, amount=sub.price)], provider_token="", currency="XTR", - payload=str(user_id), + payload=f"{user_id}#{sub.duration}", reply_to_message_id=callback.message.message_id, ) @@ -581,15 +619,22 @@ async def successful_payment_handler(self, message: Message) -> None: payload = successful_payment.invoice_payload charge_id = successful_payment.telegram_payment_charge_id self.db.add_charge(user_id, charge_id) - assert user_id == int(payload) - self.db.subscribe_user(user_id, self.config.sub_duration) + payload_user_id, payload_duration = payload.split("#") + assert user_id == int(payload_user_id) + self.db.subscribe_user(user_id, int(payload_duration)) await self.bot.send_message(chat_id, self.localization.SUB_SUCCESS) async def yookassa_sub_buy_proceed(self, callback: CallbackQuery) -> None: assert self.yookassa assert callback.from_user assert callback.message + assert callback.data + assert "buy:yookassa:" in callback.data assert isinstance(callback.message, Message) + assert self.bot_info + assert self.bot_info.username + + sub_key_str = callback.data.split(":")[2] user_id = callback.from_user.id email = self.db.get_email(user_id) if not email: @@ -607,14 +652,18 @@ async def yookassa_sub_buy_proceed(self, callback: CallbackQuery) -> None: await callback.message.reply(self.localization.ACTIVE_SUB.format(remaining_hours=remaining_seconds // 3600)) return - timestamp = self.db.get_current_ts() - title = self.localization.SUB_TITLE.format(user_id=user_id) - assert self.bot_info - assert self.bot_info.username + key = SubKey(sub_key_str) + sub = self.config.sub_configs[key] + key_to_title = { + SubKey.RUB_WEEK: self.localization.SUB_WEEK_TITLE, + SubKey.RUB_MONTH: self.localization.SUB_MONTH_TITLE + } + title = key_to_title[key].format(user_id=user_id) payment_data = self.yookassa.create_payment( - self.config.sub_price_rub, title, email=email, bot_username=self.bot_info.username + sub.price, title, email=email, bot_username=self.bot_info.username ) payment_id = payment_data["id"] + timestamp = self.db.get_current_ts() try: url = payment_data["confirmation"]["confirmation_url"] status = payment_data["status"] @@ -636,7 +685,8 @@ async def yookassa_check_payments(self) -> None: payment_id=payment.payment_id, status=status, internal_status=payment.internal_status ) if status == YookassaStatus.SUCCEEDED: - self.db.subscribe_user(payment.user_id, self.config.sub_duration) + sub_key = SubKey(self.yookassa.get_sub_key(payment.payment_id)) + self.db.subscribe_user(payment.user_id, self.config.sub_configs[sub_key].duration) await self.bot.send_message(chat_id=payment.chat_id, text=self.localization.SUB_SUCCESS) self.db.set_payment_status(payment.payment_id, status=status.value, internal_status="completed") elif status == YookassaStatus.CANCELED: diff --git a/src/localization.py b/src/localization.py index 2760b82..bc203d0 100644 --- a/src/localization.py +++ b/src/localization.py @@ -35,8 +35,10 @@ class Localization: ACTIVE_SUB: str REMAINING_MESSAGES: str SET_EMAIL: str - SUB_TITLE: str - SUB_SHORT_TITLE: str + SUB_WEEK_TITLE: str + SUB_MONTH_TITLE: str + SUB_WEEK_SHORT_TITLE: str + SUB_MONTH_SHORT_TITLE: str SUB_NOT_CHAT: str SUB_SUCCESS: str MODEL_NOT_SUPPORTED: str @@ -67,8 +69,10 @@ class Localization: DALLE_LIMIT: str DALLE_PROMPT: str DALLE_ERROR: str - BUY_WITH_STARS: str - BUY_WITH_RUB: str + BUY_WEEK_WITH_STARS: str + BUY_WEEK_WITH_RUB: str + BUY_MONTH_WITH_STARS: str + BUY_MONTH_WITH_RUB: str WRONG_COMMAND: str PAY_SUPPORT: str PRIVACY: Template diff --git a/src/payments.py b/src/payments.py index bc598f9..b3fbd39 100644 --- a/src/payments.py +++ b/src/payments.py @@ -47,3 +47,9 @@ def cancel_payment(self, payment_id: str) -> None: def check_payment(self, payment_id: str) -> YookassaStatus: payment = json.loads((Payment.find_one(payment_id)).json()) return YookassaStatus(payment["status"]) + + def get_sub_key(self, payment_id: str) -> str: + payment = json.loads((Payment.find_one(payment_id)).json()) + if "недел" in payment["description"]: + return "rub_week" + return "rub_month" diff --git a/templates/ru_sub.jinja b/templates/ru_sub.jinja index b8022ba..ff73d66 100644 --- a/templates/ru_sub.jinja +++ b/templates/ru_sub.jinja @@ -1,8 +1,8 @@ -*Покупка подписки в боте 'Сайга' на неделю* +*Покупка подписки в боте 'Сайга'* Лимиты общения с моделями станут такими: {{sub_limits}} -Подписка будет действовать ровно 7 дней. +Подписка будет действовать ровно 7 дней для подписки на неделю и ровно 31 день для подписки на месяц. -Цена подписки: *{{price}} рублей* +Цена подписки: *{{price_week}} рублей* за неделю и *{{price_month}} рублей* за месяц