chatgpt_bot/bot/handlers/user/user_handlers.py

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())