import csv from aiogram.fsm.context import FSMContext from bot.handlers.user.states import ChooseModelState, SendMessageState import time from db.db import * from aiogram import Router, types, F from aiogram.filters import CommandStart, Command from aiogram.types import LabeledPrice, PreCheckoutQuery, InlineKeyboardButton, FSInputFile from aiogram.utils.keyboard import InlineKeyboardBuilder from db.db import user_exists, add_user_to_db, set_curren_model, get_temp_prompts from bot.handlers.user.model_processing import PayCallBack, BuyCallBack, ChooseCallBack from bot.handlers.user.page_processing import get_page_content from bot.kbs.inline import PageCallBack from config import shop_api_token, admins from bot.utils.openai_tasks import get_model_suggestion_from_openai, get_answer_to_question user = Router() @user.message(CommandStart()) async def start(message: types.Message, state: FSMContext): await state.clear() # TODO isAdded = await user_exists(message.from_user.id) if not isAdded: # Извлекаем информацию о пользователе first_name = message.from_user.first_name or "" last_name = message.from_user.last_name or "" user_name = message.from_user.username or "" # Добавляем пользователя в базу данных await add_user_to_db( user_id=message.from_user.id, first_name=first_name, last_name=last_name, user_name=user_name ) text, reply_markup = await get_page_content(page=0,name="main", user_id=message.from_user.id) await message.answer(text, reply_markup=reply_markup) @user.callback_query(PageCallBack.filter()) async def user_pages(callback: types.CallbackQuery, callback_data: PageCallBack, state: FSMContext): await state.clear() text, reply_markup = await get_page_content( page=callback_data.page, name=callback_data.page_name, user_id=callback.from_user.id, ) await callback.message.edit_text(text, reply_markup=reply_markup) await callback.answer() @user.callback_query(ChooseCallBack.filter()) async def choose(callback: types.CallbackQuery, callback_data: ChooseCallBack): await set_curren_model(user_id=callback.from_user.id, name_of_model=callback_data.name) text, reply_markup = await get_page_content( page=2, user_id=callback.from_user.id, name='tarif', ) await callback.message.edit_text(text, reply_markup=reply_markup) await callback.answer() @user.message(Command(commands=['tarif'])) async def tarif(message: types.Message): # Получаем текст и клавиатуру для страницы тарифа text, reply_markup = await get_page_content( page=2, user_id=message.from_user.id, # Используем message.from_user.id вместо callback name='tarif', ) # Отправляем новое сообщение с текстом и кнопками await message.answer(text, reply_markup=reply_markup) @user.callback_query(BuyCallBack.filter()) async def buy(callback: types.CallbackQuery, callback_data: BuyCallBack): kbds = InlineKeyboardBuilder() buttons = [ InlineKeyboardButton(text="Купить", callback_data=PayCallBack( name=callback_data.name, option_of_model=callback_data.option_of_model, tokens=callback_data.tokens, amount=callback_data.amount ).pack()), InlineKeyboardButton(text="⬅️Назад", callback_data=PageCallBack(page=2, page_name="tarif").pack()), ] for button in buttons: kbds.add(button) await callback.message.edit_text( text = f"Модель: {callback_data.name}\nТариф: {callback_data.option_of_model}\nКол-ов запросов: {callback_data.tokens}\nСтоимость: {callback_data.amount}RUB", reply_markup=kbds.adjust(*(1,)).as_markup() ) await callback.answer() @user.callback_query(PayCallBack.filter()) async def pay(callback: types.CallbackQuery, callback_data: PayCallBack): temp = await callback.message.bot.send_invoice( chat_id=callback.from_user.id, description=f'{callback_data.tokens} токенов за {callback_data.amount} RUB для модели {callback_data.name}', title=f'{callback_data.option_of_model} подписка на {callback_data.name}', payload=f'sub_{callback_data.name}_{callback.from_user.id}_{callback_data.tokens}', # Добавляем id пользователя и количество токенов в payload provider_token=shop_api_token, start_parameter='test', currency="RUB", prices=[LabeledPrice( label=f'Оплата токенов для {callback_data.option_of_model}', amount=int(callback_data.amount) * 100, )], ) # if temp.successful_payment is None: # await callback.bot.send_message(chat_id=callback.from_user.id,text="Something went wrong") @user.pre_checkout_query() async def process_pre_checkout_query(pre_checkout_query: PreCheckoutQuery): await pre_checkout_query.answer(ok=True) @user.message(F.successful_payment) async def process_payment(message: types.Message): if message.successful_payment is None: await message.bot.send_message(chat_id=message.from_user.id,text="Что то пошло не так") return payload_data = message.successful_payment.invoice_payload.split('_') if len(payload_data) != 4: return # Не валидный payload model_name = payload_data[1] # Имя модели user_id = int(payload_data[2]) # ID пользователя tokens = int(payload_data[3]) # Количество токенов await add_tokens_to_user(user_id, model_name, tokens) time.sleep(2) await message.bot.send_message(chat_id=message.from_user.id, text="Спасибо, что пользуетесь нашим сервисом!") await message.bot.send_message(chat_id=message.from_user.id, text="Для работы с chatGPT отправьте текст с вашим вопросом и получите на него ответ)") await set_curren_model(user_id, model_name) @user.message(Command(commands=["failed"])) async def process_payment(message: types.Message): await message.bot.send_message(chat_id=message.from_user.id, text="Что-то пошло не так( Попробуйте еще раз!") text, reply_markup = await get_page_content(page=0, name="main", user_id=message.from_user.id) await message.bot.send_message( chat_id=message.from_user.id, text=text, reply_markup=reply_markup, ) @user.callback_query(lambda c: c.data == 'choose') async def chat_test_gpt(callback: types.CallbackQuery, state: FSMContext): kbds = InlineKeyboardBuilder() kbds.add(InlineKeyboardButton( text="⬅️Назад", callback_data=PageCallBack( page=3, page_name='more_about_models', ).pack() )) await callback.message.edit_text("Вы находитесь в разделе подбора модели. Опишите ваши задачи, и я помогу выбрать модель.", reply_markup=kbds.adjust(*(1,)).as_markup()) await callback.answer() await state.set_state(ChooseModelState.choosing_model) @user.message(ChooseModelState.choosing_model) async def process_model_question(message: types.Message): user_id = message.from_user.id remaining_queries = await get_temp_prompts(user_id=user_id) kbds = InlineKeyboardBuilder() kbds.add(InlineKeyboardButton( text="⬅️Назад", callback_data=PageCallBack( page=3, page_name='more_about_models' ).pack() )) if remaining_queries > 0: # Обрабатываем вопрос с помощью OpenAI GPT openai_response = await get_model_suggestion_from_openai(message.text) # Отправляем ответ пользователю # await message.answer(text=f"Ответ OpenAI: {openai_response}" + f"\nОсталось {remaining_queries-1} запросов.", reply_markup=kbds.adjust(*(1,)).as_markup()) # Уменьшаем количество запросов if "Ошибка при обращении к OpenAI" not in openai_response: await message.answer( text=f"Ответ OpenAI: {openai_response}" + f"\nОсталось {remaining_queries - 1} запросов.", reply_markup=kbds.adjust(*(1,)).as_markup()) # await decrease_test_queries(user_id) ... else: await message.answer( text=f"Попробуйте еще раз или чуть подождите!") else: await message.answer(text="У вас закончились доступные запросы для подбора модели.", reply_markup=kbds.adjust(*(1,)).as_markup()) @user.message(Command(commands=['export_users'])) async def send_csv_to_admin(message: types.Message): # Генерируем CSV файл if message.from_user.id in admins: users = await get_all_users() # Генерируем CSV файл with open('data/users.csv', 'w', newline='', encoding='utf-8') as csvfile: csvwriter = csv.writer(csvfile, delimiter=';', quoting=csv.QUOTE_MINIMAL) csvwriter.writerow(['ID', 'First Name', 'Last Name', 'Username']) for user in users: csvwriter.writerow([user.user_id, user.first_name, user.last_name, user.user_name]) # Отправляем CSV файл администратору csv_file = FSInputFile('data/users.csv') # Обернуть путь к файлу как InputFile await message.answer_document(csv_file) @user.message(Command(commands=['send_message_to_all'])) async def send_messages_to_users(message: types.Message, state: FSMContext): if message.from_user.id in admins: # Устанавливаем состояние, что бот ждет текст сообщения для рассылки await state.set_state(SendMessageState.waiting_for_message) await message.answer("Введите сообщение для рассылки:") @user.message(SendMessageState.waiting_for_message) async def get_message_for_users(msg: types.Message, state: FSMContext): # Получаем текст сообщения для рассылки text_to_send = msg.text users = await get_all_users() # Рассылаем сообщение всем пользователям for user in users: if user.user_id == msg.from_user.id: continue try: await msg.bot.send_message(chat_id=user.user_id, text=f'Сообщение от админа:\n{text_to_send}') except Exception as e: # Логируем ошибки для недоступных пользователей print(f"Не удалось отправить сообщение пользователю {user.user_id}: {e}") await msg.answer("Сообщение успешно отправлено всем пользователям.") # Очистить состояние после завершения рассылки await state.clear() @user.message(~Command(commands=["export_users", 'send_message_to_all'])) # Игнорируем команды async def chatgpt_chatting(message: types.Message): user_id = message.from_user.id db_model = await get_current_model(user_id) kbds = InlineKeyboardBuilder() kbds.add(InlineKeyboardButton( text="Выбрать тариф", callback_data=PageCallBack( page=2, page_name='tarif' ).pack())) kbds.add(InlineKeyboardButton( text="⬅️На главную", callback_data=PageCallBack( page=0, page_name='main' ).pack() )) if db_model is None: await message.bot.send_message(chat_id=message.from_user.id, text="У вас еще не выбрана модель", reply_markup=kbds.adjust(*(1,)).as_markup()) models = { '1o' : "o1-preview-2024-09-12", '1o mini': "o1-mini-2024-09-12", '4o': "gpt-4o", '4o mini': "gpt-4o-mini" } current_model = models[db_model] remaining_queries = await get_current_model_tokens(user_id=user_id, name_of_model=db_model) if remaining_queries > 0: openai_response = await get_answer_to_question(message.text, current_model) # Уменьшаем количество запросов if "Ошибка при обращении к OpenAI" not in openai_response: await decrease_current_model_tokens(user_id, name_of_model=db_model) await message.answer(text=f"Ответ OpenAI: {openai_response}") await message.answer(text=f"Модель: {current_model}\nОсталось {remaining_queries - 1} запросов.", reply_markup=kbds.adjust(*(1,)).as_markup()) else: await message.answer( text=f"{openai_response}") await message.answer( text=f"Попробуйте еще раз или вы можете выбрать другую доступную для вас модель!") else: await message.answer(text=f"У вас закончились доступные запросы для модели {current_model}. Купите новые токены или выберете другую модель", reply_markup=kbds.adjust(*(1,)).as_markup())