350 lines
14 KiB
Python
350 lines
14 KiB
Python
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)
|
||
|
||
# time.sleep(2)
|
||
#
|
||
# await pre_checkout_query.bot.send_message(chat_id=pre_checkout_query.from_user.id, text="Спасибо, пользутесь нашим сервисом!")
|
||
#
|
||
# text, reply_markup = await get_page_content(page=0, name="main", user_id=pre_checkout_query.from_user.id)
|
||
#
|
||
# await pre_checkout_query.bot.send_message(
|
||
# chat_id=pre_checkout_query.from_user.id,
|
||
# text=text,
|
||
# reply_markup=reply_markup,
|
||
# )
|
||
|
||
@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 отправьте текст с вашим вопросом и получите на него ответ)")
|
||
|
||
@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",
|
||
'1o mini': "o1-mini",
|
||
'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())
|
||
|
||
|
||
|
||
|
||
|